fix: wrong dp-rank condition when enable pp

Author: liuqh16

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -442,7 +442,6 @@ def _criterion(outputs, inputs):
                 self.plugin.pp_size > 1
                 and self.booster.plugin.stage_manager.is_last_stage()
                 and self.tp_rank == 0
-                and self.dp_rank == 0
             ):
                 reward = all_reduce_mean(reward.mean(), self.plugin)
                 format_acc = all_reduce_mean(format_acc.mean(), self.plugin)
@@ -469,7 +468,7 @@ def _criterion(outputs, inputs):
                 self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0
             ):
                 if (not self.plugin.pp_size > 1 and self.rank == 0) or (
-                    self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0
+                    self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0 and self.dp_rank == 0
                 ):
                     raw_batch_reward_mean = torch.cat(self.raw_train_batch_reward, dim=0).mean().cpu().item()
                     raw_batch_format_acc_mean = torch.cat(self.raw_train_batch_format_acc, dim=0).mean().cpu().item()