fix if-check

Author: H-Huang

torchtitan/distributed/pipeline_parallel.py

@@ -100,7 +100,7 @@ def build_pipeline_schedule(
 
     if (
         job_config.parallelism.pipeline_parallel_expert_parallel_overlap
-        and job_config.parallelism.pipeline_parallel_schedule == "dualpipev"
+        and isinstance(schedule, ScheduleDualPipeV)
     ):
         schedule.register_custom_function(OVERLAP_F_B, overlap_callback)
 
@@ -466,6 +466,7 @@ def overlap_callback(action: _Action, ctx: _PipelineContext):
     Custom callback for OVERLAP_F_B computation that allows expert parallel communication
     and pipeline parallel computation to overlap.
     """
+    print("calling into overlap callback")
     schedule = ctx.schedule_ref
     assert isinstance(schedule, _PipelineScheduleRuntime)
     stage_index_to_stage: dict[int, _PipelineStageBase] = {