github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-09 21:09:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

use consumer global step

Browse Source
This commit is contained in:
YeAnbang
2025-05-15 14:15:40 +08:00
parent 094f119b3a
commit 4ec73298b2
3 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
applications/ColossalChat/coati/distributed/grpo_consumer.py
View File

@@ -266,7 +266,6 @@ class GRPOConsumer(BaseConsumer):
total_samples = all_reduce_sum(torch.sum(torch.ones_like(loss_mask, device=loss_mask.device)), self.plugin)
self.effective_sample_count += effective_samples.item()
self.total_sample_count += total_samples.item()
pbar.set_postfix(
{
"Global Step": self.global_step,
@@ -522,7 +521,6 @@ class GRPOConsumer(BaseConsumer):
# All gather excessive prompts index across DP ranks.
excessive_prompts_idx = [idx + self.dp_rank * self.minibatch_size for idx in excessive_prompts_idx]
excessive_prompts_idx = all_gather_tensors(excessive_prompts_idx, self.plugin)
return loss_scalar, excessive_prompts_idx
else:
return None, excessive_prompts_idx