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

address conversation

Browse Source
This commit is contained in:
YeAnbang
2025-05-29 10:25:59 +08:00
parent 58f8c9bb43
commit ee939d9aa5
3 changed files with 36 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -84,7 +84,6 @@ class GRPOConsumer(BaseConsumer):
self.project_name = project_name
self.effective_sample_count = 0
self.effective_prompt_count = 0
self.total_sample_count = 0
self.project_name = project_name
self.run_name = run_name
self.wandb_group_name = wandb_group_name
@@ -429,11 +428,9 @@ class GRPOConsumer(BaseConsumer):
self.optimizer.step()
self.optimizer.zero_grad()
self.global_step += 1
# no need to run all_reduce_sum on total_sample_count, because all dp ranks recieves a complete inference batch from all producers.
sample_utilization = self.effective_sample_count / self.total_sample_count
sample_utilization = self.effective_sample_count / len(self.raw_train_batch_reward) / self.num_generations
self.effective_prompt_count = 0
self.effective_sample_count = 0
self.total_sample_count = 0
loss_scalar = self.accum_loss.item()
if not self.plugin.pp_size > 1 or (
self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0
@@ -520,35 +517,6 @@ class GRPOConsumer(BaseConsumer):
rollout["ans_acc"] = ans_acc.view((-1, 1))
return rollout
def prompt_level_filtering(self, rollout_group: Dict[str, Any]) -> Dict[str, Any]:
"""
rollout_group: Dict[str, Any]
a group of samples generated by the model from the same prompt
contain the following keys:
"input_ids": torch.Tensor, [num_of_generation, prompt_length + response_length]
"attention_mask": torch.Tensor, [num_of_generation, prompt_length + response_length]
"action_mask": torch.Tensor, [num_of_generation, response_length]
"action_log_probs": torch.Tensor, [num_of_generation, response_length]
"response_idx": int, torch.Tensor, [num_of_generation, 2]
"gt_answer": torch.Tensor, [num_of_generation, 128]
"temperature": torch.Tensor, [] (scalar)
"reward": torch.Tensor, [num_of_generation]
"format_acc": torch.Tensor, [num_of_generation]
"ans_acc": torch.Tensor, [num_of_generation]
"""
self.total_sample_count += rollout_group["input_ids"].size(0)
if self.filter_range is not None:
# filter prompt whoes accuracy is too high or too low (out of range)
group_ans_acc = torch.mean(rollout_group["ans_acc"])
if group_ans_acc < self.filter_range[0] or group_ans_acc > self.filter_range[1]:
# filter out the prompt
return None
else:
return rollout_group
else:
# no filter
return rollout_group
def state_dict(self):
self.policy_model._force_wait_all_gather()
model = self.policy_model.unwrap()