github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-07-15 15:29:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

add entropy

Browse Source
This commit is contained in:
YeAnbang
2025-07-16 16:44:23 +08:00
parent dba0c0c4ed
commit ddda79c36f
1 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -346,6 +346,7 @@ class GRPOConsumer(BaseConsumer):
data_policy_forward["reference_action_log_probs"] = reference_action_log_probs
kl = []
policy_model_logits = torch.empty_like(input_ids_forward_micro_batch, device=self.device)
def _criterion(outputs, inputs):
action_logits = outputs.logits
@@ -425,6 +426,20 @@ class GRPOConsumer(BaseConsumer):
kl = all_reduce_mean(torch.mean(torch.stack(kl)).to(loss.device), self.plugin).data
mean_kl.append(kl)
mean_loss.append(all_reduce_mean(loss, self.plugin).data)
mini_batch_entropies.append(
all_reduce_mean(
(
(
(
entropy_from_logits(policy_model_logits[:, -num_action:])
* action_mask_forward_micro_batch
).sum(-1)
)
/ action_mask_forward_micro_batch.sum(-1)
).detach(),
self.plugin,
)
)
else:
policy_model_logits = self.policy_model(
input_ids=input_ids_forward_micro_batch,