fix ut

lhb8125 · lhb8125 · commit 12a2a22ab909 · 2025-12-08T07:38:30.000-08:00
Signed-off-by: Hongbin Liu &lt;hongbinl@nvidia.com&gt;
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -505,7 +505,8 @@ def run(
                 b_grad=b_grad,
                 is_last_layer_in_bwd=(i == b_num_layers - 1),
             )
-            b_layer.release_state()
+            if i < b_num_layers - 1:
+                b_layer.release_state()
             torch.cuda.nvtx.range_pop()
 
         # backward pass for the remaining layers

Original file line number	Diff line number	Diff line change
`@@ -505,7 +505,8 @@ def run(`
`505`	`505`	`b_grad=b_grad,`
`506`	`506`	`is_last_layer_in_bwd=(i == b_num_layers - 1),`
`507`	`507`	`)`
`508`		`- b_layer.release_state()`
	`508`	`+ if i < b_num_layers - 1:`
	`509`	`+ b_layer.release_state()`
`509`	`510`	`torch.cuda.nvtx.range_pop()`
`510`	`511`
`511`	`512`	`# backward pass for the remaining layers`