Merge branch 'malay/fwd_step_fix' into 'main'

shanmugamr1992 · shanmugamr1992 · commit ec16db0186f4 · 2024-05-02T15:29:31.000-07:00
forward step missing arg

See merge request ADLR/megatron-lm!1371
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
@@ -907,7 +907,9 @@ def backward_step_helper(microbatch_id):
             )
 
         else:  # no p2p overlap
-            output_tensor = forward_step_helper(forward_k, checkpoint_activations_microbatch)
+            output_tensor = forward_step_helper(
+                forward_k, current_microbatch, checkpoint_activations_microbatch
+            )
 
             # Backward pass.
             backward_k = k

Original file line number	Diff line number	Diff line change
`@@ -907,7 +907,9 @@ def backward_step_helper(microbatch_id):`
`907`	`907`	`)`
`908`	`908`
`909`	`909`	`else: # no p2p overlap`
`910`		`- output_tensor = forward_step_helper(forward_k, checkpoint_activations_microbatch)`
	`910`	`+ output_tensor = forward_step_helper(`
	`911`	`+ forward_k, current_microbatch, checkpoint_activations_microbatch`
	`912`	`+ )`
`911`	`913`
`912`	`914`	`# Backward pass.`
`913`	`915`	`backward_k = k`