fix: Make pp related import optional

Author: rrutmann

src/modalities/models/parallelism/pipeline_parallelism.py

@@ -14,11 +14,15 @@
 from torch.distributed.pipelining.schedules import (
     PipelineScheduleMulti,
     PipelineScheduleSingle,
-    ScheduleDualPipeV,
     ScheduleZBVZeroBubble,
     get_schedule_class,
 )
 
+try:
+    from torch.distributed.pipelining.schedules import ScheduleDualPipeV
+except ImportError:
+    ScheduleDualPipeV = None
+
 from modalities.loss_functions import Loss
 from modalities.models.model import NNModel
 from modalities.models.parallelism.stages_generator import StagesGenerator
@@ -152,7 +156,10 @@ def _get_stage_ids_of_pp_rank(
         num_stages: int,
         schedule_class: Type[PipelineScheduleSingle | PipelineScheduleMulti],
     ) -> list[int]:
-        style = "v" if schedule_class in (ScheduleZBVZeroBubble, ScheduleDualPipeV) else "loop"
+        v_schedules = [ScheduleZBVZeroBubble]
+        if ScheduleDualPipeV is not None:
+            v_schedules.append(ScheduleDualPipeV)
+        style = "v" if schedule_class in tuple(v_schedules) else "loop"
         pp_size = pp_mesh.size()
         pp_rank = pp_mesh.get_local_rank()
         stages_per_rank = num_stages // pp_size