fix pp memory issue (#6344)

TongLi3701 · Tong Li · YeAnbang · commit 3a4681fdd9c6 · 2025-08-05T14:04:09.000+08:00
Co-authored-by: Tong Li &lt;tong.li35271158@gmail.com&gt;
diff --git a/applications/ColossalChat/coati/distributed/grpo_consumer.py b/applications/ColossalChat/coati/distributed/grpo_consumer.py
@@ -358,7 +358,7 @@ def _criterion(outputs, inputs):
                         criterion=_criterion,
                         optimizer=self.optimizer,
                         return_loss=True,
-                        return_outputs=True,
+                        return_outputs=False,
                     )
                     loss = policy_model_outputs["loss"]
 

Original file line number	Diff line number	Diff line change
`@@ -358,7 +358,7 @@ def _criterion(outputs, inputs):`
`358`	`358`	`criterion=_criterion,`
`359`	`359`	`optimizer=self.optimizer,`
`360`	`360`	`return_loss=True,`
`361`		`- return_outputs=True,`
	`361`	`+ return_outputs=False,`
`362`	`362`	`)`
`363`	`363`	`loss = policy_model_outputs["loss"]`
`364`	`364`