Update

[ghstack-poisoned]
pytorch · Feb 12, 2025 · 2f4d1ce · 2f4d1ce
1 parent c131309
commit 2f4d1ce
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/torchtitan/models/llama/pipeline_llama.py b/torchtitan/models/llama/pipeline_llama.py
@@ -13,7 +13,11 @@
 import torch.nn as nn
 from torch.distributed import DeviceMesh
 from torch.distributed.pipelining import PipelineStage
-from torch.distributed.pipelining.schedules import _PipelineSchedule, get_schedule_class, ScheduleZBVZeroBubble
+from torch.distributed.pipelining.schedules import (
+    _PipelineSchedule,
+    get_schedule_class,
+    ScheduleZBVZeroBubble,
+)
 
 from torchtitan.config_manager import JobConfig
 from torchtitan.logging import logger

diff --git a/torchtitan/train_spec.py b/torchtitan/train_spec.py
@@ -36,7 +36,8 @@ class ModelProtocol(Protocol):
     """
 
     @staticmethod
-    def from_model_args(args: BaseModelArgs) -> nn.Module: ...
+    def from_model_args(args: BaseModelArgs) -> nn.Module:
+        ...
 
 
 OptimizersBuilder: TypeAlias = Callable[