Update on "[pipelining] Improve schedule csv loading"

H-Huang · H-Huang · commit e2b470cc3f9f · 2024-12-03T16:58:24.000-08:00
Add small changes based on feedback from Less when testing out pytorch/torchtitan#707 - expose `validate_schedule` as a function - handle spaces around actions in csv file - add error arrow to `_format_pipeline_schedule()` to better show where the step errored cc awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
diff --git a/test/distributed/pipelining/test_schedule.py b/test/distributed/pipelining/test_schedule.py
@@ -25,6 +25,7 @@
     _PipelineSchedule,
     _PipelineScheduleRuntime,
     _simulate_comms_compute,
+    _validate_schedule,
     B,
     F,
     get_schedule_class,
@@ -34,7 +35,6 @@
     RESHARD,
     SEND_B,
     UNSHARD,
-    validate_schedule,
     W,
 )
 from torch.distributed.pipelining.stage import _PipelineStageBase, PipelineStage
@@ -873,7 +873,7 @@ def test_valid_schedule(self):
         pp_group_size = 2
         num_stages = 2
         num_microbatches = 1
-        validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
+        _validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
 
     def test_invalid_schedule_missing_rank(self):
         actions = {
@@ -883,7 +883,7 @@ def test_invalid_schedule_missing_rank(self):
         num_stages = 2
         num_microbatches = 1
         with self.assertRaises(AssertionError):
-            validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
+            _validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
 
     def test_invalid_schedule_missing_action(self):
         actions = {
@@ -894,14 +894,7 @@ def test_invalid_schedule_missing_action(self):
         num_stages = 2
         num_microbatches = 1
         with self.assertRaises(AssertionError):
-            validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
-
-    def test_format(self):
-        actions = {
-            0: [_Action(0, F, 0), _Action(0, B, 0)],
-            1: [_Action(1, F, 0), _Action(1, B, 0)],
-        }
-        print(_format_pipeline_order(actions, error_step_number=0))
+            _validate_schedule(actions, pp_group_size, num_stages, num_microbatches)
 
 
 instantiate_parametrized_tests(TestScheduleLowering)
diff --git a/torch/distributed/pipelining/schedules.py b/torch/distributed/pipelining/schedules.py
@@ -1000,8 +1000,8 @@ def _ready_to_schedule(
     return comm_actions
 
 
-def validate_schedule(
-    actions: Dict[int, List[_Action | None]],
+def _validate_schedule(
+    actions: Dict[int, List[Optional[_Action]]],
     pp_group_size: int,
     num_stages: int,
     num_microbatches: int,
@@ -1157,7 +1157,7 @@ def _load_csv(self, filename, format="compute_only"):
             reader = csv.reader(csvfile)
             for rank, row in enumerate(reader):
                 self.pipeline_order[rank] = [_Action.from_str(s) for s in row]
-        validate_schedule(
+        _validate_schedule(
             self.pipeline_order,
             self.pp_group_size,
             self._num_stages,
