135 lines
4.5 KiB
Python
135 lines
4.5 KiB
Python
"""GRPO/CISPO core functions to implement and align with MiniMind source."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
|
|
def group_relative_advantages(
|
|
rewards: torch.Tensor,
|
|
num_generations: int,
|
|
eps: float = 1e-4,
|
|
) -> torch.Tensor:
|
|
"""Compute group-relative advantages from flat rewards.
|
|
|
|
Align with: trainer/train_grpo.py:121-124
|
|
|
|
Args:
|
|
rewards: Flat tensor with shape [batch * num_generations].
|
|
num_generations: Number of responses sampled per prompt.
|
|
eps: Stabilizer added to group std.
|
|
|
|
Returns:
|
|
Tensor with shape [batch * num_generations].
|
|
|
|
Formula:
|
|
A_i = (R_i - mean(group)) / (std(group) + eps)
|
|
"""
|
|
groups=rewards.view(-1,num_generations)
|
|
group_mean=groups.mean(dim=1,keepdim=True)
|
|
group_var=((groups-group_mean)**2).sum(dim=1)/num_generations
|
|
group_var=torch.sqrt(group_var)
|
|
group_mean=group_mean.repeat_interleave(num_generations)
|
|
group_var=group_var.repeat_interleave(num_generations)
|
|
return (rewards-group_mean)/(group_var+eps)
|
|
raise NotImplementedError("Implement in lesson 24: GRPO and CISPO.")
|
|
|
|
|
|
def token_kl_penalty(policy_logps: torch.Tensor, ref_logps: torch.Tensor) -> torch.Tensor:
|
|
"""Compute MiniMind's token-level reference KL penalty tensor.
|
|
|
|
Align with: trainer/train_grpo.py:132-133
|
|
|
|
Args:
|
|
policy_logps: Current policy logprobs, shape [batch * num_generations, response_len].
|
|
ref_logps: Frozen reference logprobs, same shape.
|
|
|
|
Returns:
|
|
Per-token KL penalty tensor with the same shape.
|
|
|
|
Formula:
|
|
delta = ref_logps - policy_logps
|
|
penalty = exp(delta) - delta - 1
|
|
"""
|
|
delta = ref_logps - policy_logps
|
|
penalty = torch.exp(delta) - delta - 1
|
|
return penalty
|
|
raise NotImplementedError("Implement in lesson 24: GRPO and CISPO.")
|
|
|
|
|
|
def grpo_policy_loss(
|
|
policy_logps: torch.Tensor,
|
|
old_logps: torch.Tensor,
|
|
ref_logps: torch.Tensor,
|
|
advantages: torch.Tensor,
|
|
completion_mask: torch.Tensor,
|
|
epsilon: float,
|
|
beta: float,
|
|
) -> torch.Tensor:
|
|
"""Compute MiniMind's GRPO token loss reduced to a scalar.
|
|
|
|
Align with: trainer/train_grpo.py:132-143
|
|
|
|
Args:
|
|
policy_logps: Current policy logprobs, shape [batch * num_generations, response_len].
|
|
old_logps: Rollout policy logprobs, same shape.
|
|
ref_logps: Frozen reference model logprobs, same shape.
|
|
advantages: Sequence-level advantages, shape [batch * num_generations].
|
|
completion_mask: Valid response-token mask, shape [batch * num_generations, response_len].
|
|
epsilon: PPO-style clip epsilon.
|
|
beta: KL penalty coefficient.
|
|
|
|
Returns:
|
|
Scalar policy loss.
|
|
"""
|
|
ratio=policy_logps-old_logps
|
|
ratio=torch.exp(ratio)
|
|
loss1=torch.clamp(ratio,1-epsilon,1+epsilon)*advantages.unsqueeze(1)
|
|
loss2=ratio*advantages.unsqueeze(1)
|
|
kl_pl=token_kl_penalty(policy_logps,ref_logps)
|
|
per_token_loss=-(torch.min(loss1,loss2)-beta*kl_pl)
|
|
loss = (
|
|
(per_token_loss * completion_mask).sum(dim=1)
|
|
/ completion_mask.sum(dim=1).clamp(min=1)
|
|
).mean()
|
|
return loss
|
|
raise NotImplementedError("Implement in lesson 24: GRPO and CISPO.")
|
|
|
|
|
|
def cispo_policy_loss(
|
|
policy_logps: torch.Tensor,
|
|
old_logps: torch.Tensor,
|
|
ref_logps: torch.Tensor,
|
|
advantages: torch.Tensor,
|
|
completion_mask: torch.Tensor,
|
|
epsilon_high: float,
|
|
beta: float,
|
|
) -> torch.Tensor:
|
|
"""Compute MiniMind's CISPO token loss reduced to a scalar.
|
|
|
|
Align with: trainer/train_grpo.py:132-143
|
|
|
|
Args:
|
|
policy_logps: Current policy logprobs, shape [batch * num_generations, response_len].
|
|
old_logps: Rollout policy logprobs, same shape.
|
|
ref_logps: Frozen reference model logprobs, same shape.
|
|
advantages: Sequence-level advantages, shape [batch * num_generations].
|
|
completion_mask: Valid response-token mask, shape [batch * num_generations, response_len].
|
|
epsilon_high: Upper bound used to clamp ratio as a detached weight.
|
|
beta: KL penalty coefficient.
|
|
|
|
Returns:
|
|
Scalar policy loss.
|
|
"""
|
|
ratio=policy_logps-old_logps
|
|
ratio=torch.exp(ratio)
|
|
climp_ratio=torch.clamp(ratio,max=epsilon_high).detach()
|
|
kl_pl=token_kl_penalty(policy_logps,ref_logps)
|
|
per_token_loss=-climp_ratio*advantages.unsqueeze(1)*policy_logps+beta*kl_pl
|
|
loss = (
|
|
(per_token_loss * completion_mask).sum(dim=1)
|
|
/ completion_mask.sum(dim=1).clamp(min=1)
|
|
).mean()
|
|
return loss
|
|
raise NotImplementedError("Implement in lesson 24: GRPO and CISPO.")
|