Update PP to release memory earlier (#1922)

H-Huang · web-flow · commit e5ef99aab5bf · 2025-10-22T09:28:36.000-07:00
Uses the API added in pytorch/pytorch#165822, since we do not return any output from PP step(). This allows us to release the memory earlier,
diff --git a/torchtitan/train.py b/torchtitan/train.py
@@ -460,12 +460,14 @@ def forward_backward_step(
                         **extra_kwargs,
                         target=targets,
                         losses=losses,
+                        return_outputs=False,
                     )
                 else:
                     self.pp_schedule.step(
                         **extra_kwargs,
                         target=targets,
                         losses=losses,
+                        return_outputs=False,
                     )
 
             # accumulate losses across pipeline microbatches

Original file line number	Diff line number	Diff line change
`@@ -460,12 +460,14 @@ def forward_backward_step(`
`460`	`460`	`**extra_kwargs,`
`461`	`461`	`target=targets,`
`462`	`462`	`losses=losses,`
	`463`	`+ return_outputs=False,`
`463`	`464`	`)`
`464`	`465`	`else:`
`465`	`466`	`self.pp_schedule.step(`
`466`	`467`	`**extra_kwargs,`
`467`	`468`	`target=targets,`
`468`	`469`	`losses=losses,`
	`470`	`+ return_outputs=False,`
`469`	`471`	`)`
`470`	`472`
`471`	`473`	`# accumulate losses across pipeline microbatches`