fix bug, tested

This commit is contained in:
YeAnbang 2025-06-09 09:37:28 +08:00
parent dc3033e68a
commit 3bed6ae9ee
3 changed files with 7 additions and 5 deletions

4
.gitignore vendored
View File

@ -167,3 +167,7 @@ applications/ColossalChat/wandb
applications/ColossalChat/model applications/ColossalChat/model
applications/ColossalChat/eval applications/ColossalChat/eval
applications/ColossalChat/rollouts applications/ColossalChat/rollouts
applications/ColossalChat/*.txt
applications/ColossalChat/*.db
applications/ColossalChat/stdin
applications/ColossalChat/*.zip

View File

@ -132,9 +132,7 @@ class BaseConsumer:
format_acc = raw_batch["format_acc"][:, :, 0] format_acc = raw_batch["format_acc"][:, :, 0]
ans_acc = raw_batch["ans_acc"][:, :, 0] ans_acc = raw_batch["ans_acc"][:, :, 0]
response_len = ( response_len = (
raw_batch["response_idx"][:, :, 1] raw_batch["response_idx"][:, :, 1] - raw_batch["response_idx"][:, :, 0] + 1
- raw_batch["response_idx"][:, :, 0]
+ 1
).type(torch.float32) ).type(torch.float32)
effective_group_mask = None effective_group_mask = None
if self.filter_range is not None and self.grpo_config.get("dynamic_batching", True): if self.filter_range is not None and self.grpo_config.get("dynamic_batching", True):
@ -160,7 +158,7 @@ class BaseConsumer:
) )
if effective_group_mask is not None: if effective_group_mask is not None:
print( print(
f"[T{dist.get_rank()}] Filter recv data: {len(raw_batch_with_reward)} -> {torch.sum(effective_group_mask).cpu().item()} effective groups" f"[T{dist.get_rank()}] Filter recv data: {len(raw_batch)} -> {torch.sum(effective_group_mask).cpu().item()} effective groups"
) )
# mapping the effective group to the raw group for indexing # mapping the effective group to the raw group for indexing
effective_group_to_raw_group_mapping = {} effective_group_to_raw_group_mapping = {}

View File

@ -291,7 +291,7 @@ class BaseProducer:
reward_model_output = self.reward_model( reward_model_output = self.reward_model(
outputs["input_ids"].view((-1, outputs["input_ids"].size(-1))), outputs["input_ids"].view((-1, outputs["input_ids"].size(-1))),
gt_answer=gt_answer, gt_answer=gt_answer,
response_idx=outputs["response_idx"], response_idx=outputs["response_idx"].view((-1, 2)),
) )
outputs["reward"] = ( outputs["reward"] = (
torch.tensor([value[0] for value in reward_model_output]) torch.tensor([value[0] for value in reward_model_output])