File size: 8,169 Bytes
60b21d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """
GRPO trainer for OpenTSLM (planted on top of the SFT-initialized OpenTSLMSP model).
Group Relative Policy Optimization (Shao et al., 2024): for each prompt we sample a
group of N rollouts, score each with the dual reward (answer + faithfulness, see
reward.py), normalize rewards within the group to get advantages, and take a
policy-gradient step. No value network is needed.
This reuses OpenTSLMSP's confirmed interfaces only:
model.pad_and_apply_batch(batch) -> (inputs_embeds [B,L,H], attention_mask [B,L])
model.generate(batch, max_new_tokens, **gen_kwargs) -> List[str]
model.llm (HF CausalLM with LoRA), model.tokenizer, model.device
so it works without modifying their model code. It is GPU-only at run time (needs
the LLM forward/backward), but the logic is plain torch and can be reviewed on CPU.
Wiring into curriculum_learning.py: in CurriculumTrainer._train_stage, replace the
inner SFT step
optimizer.zero_grad(); loss = model.compute_loss(batch)
loss.backward(); optimizer.step()
with
optimizer.zero_grad()
loss, stats = grpo.grpo_loss(batch)
loss.backward(); optimizer.step()
for a dedicated `stage6_grpo` stage (see GRPO_INTEGRATION.md).
"""
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional
import torch
import torch.nn.functional as F
from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList
class SafeLogitsProcessor(LogitsProcessor):
"""Clamp NaN/inf logits during RL rollouts (sampled generation occasionally produces
them and crashes generate). Adapted from the NyayaRL GRPO recipe."""
def __call__(self, input_ids, scores):
scores = torch.nan_to_num(scores, nan=0.0, posinf=50.0, neginf=-50.0)
return torch.clamp(scores, -50.0, 50.0)
@dataclass
class GRPOConfig:
num_rollouts: int = 8 # group size N (DeepSeek-R1 default)
max_new_tokens: int = 400 # rollout length (match generation; 300 for M4)
temperature: float = 1.0
top_p: float = 1.0
kl_coef: float = 0.1 # KL penalty to frozen SFT ref (beta); 0 to disable
adv_eps: float = 1e-4 # std floor for advantage normalization
max_grad_norm: float = 0.1 # tight grad clip for RL stability (apply in train loop)
# composite reward is computed by the injected reward_fn (see reward.py); these are
# the paper's headline weights, surfaced here for logging/reference.
w_answer: float = 0.7
w_faith: float = 0.3
class GRPOTrainer:
def __init__(
self,
model, # OpenTSLMSP (policy)
reward_fn: Callable[[str, dict], Dict[str, float]],
# reward_fn(completion_text, sample) -> {"r_answer","r_faith","r_total"}
config: Optional[GRPOConfig] = None,
ref_model=None, # frozen SFT copy for KL (optional)
):
self.model = model
self.reward_fn = reward_fn
self.cfg = config or GRPOConfig()
self.ref_model = ref_model
self.device = model.device
# ---------------------------------------------------------------- rollouts
@torch.no_grad()
def _sample_rollouts(self, item: dict) -> List[str]:
"""Sample N completions for one prompt via the model's own generate()."""
gen_kwargs = dict(
do_sample=True,
temperature=self.cfg.temperature,
top_p=self.cfg.top_p,
num_return_sequences=self.cfg.num_rollouts,
logits_processor=LogitsProcessorList([SafeLogitsProcessor()]),
)
# generate() batch-decodes self.llm.generate output; with one input item and
# num_return_sequences=N it returns N strings.
completions = self.model.generate(
[item], max_new_tokens=self.cfg.max_new_tokens, **gen_kwargs
)
return completions
# ------------------------------------------------------------- log-probs
def _sequence_logprob(self, item: dict, completion: str, model=None):
"""
Per-token log-probabilities of `completion` given `item`'s prompt, under `model`
(defaults to the policy). Mirrors OpenTSLMSP.compute_loss's embedding layout:
prompt embeds (from pad_and_apply_batch) followed by the answer/completion embeds.
Returns a 1-D tensor of length A (one logprob per completion token).
"""
model = model or self.model
inputs_embeds, attention_mask = model.pad_and_apply_batch([item]) # [1,L,H],[1,L]
L = inputs_embeds.size(1)
tok = model.tokenizer(
[completion], return_tensors="pt", padding=False, truncation=True
)
ans_ids = tok.input_ids.to(self.device) # [1,A]
A = ans_ids.size(1)
ans_emb = model.llm.get_input_embeddings()(ans_ids) # [1,A,H]
full_embeds = torch.cat([inputs_embeds, ans_emb], dim=1) # [1,L+A,H]
full_mask = torch.cat(
[attention_mask, torch.ones_like(ans_ids)], dim=1
) # [1,L+A]
out = model.llm(inputs_embeds=full_embeds, attention_mask=full_mask, return_dict=True)
logits = out.logits # [1,L+A,V]
# logits at position t predict token t+1; completion tokens sit at [L, L+A).
pred_logits = logits[:, L - 1 : L + A - 1, :] # [1,A,V]
logprobs = F.log_softmax(pred_logits, dim=-1)
token_lp = logprobs.gather(-1, ans_ids.unsqueeze(-1)).squeeze(-1) # [1,A]
return token_lp.squeeze(0) # [A]
# ------------------------------------------------------------------- loss
def grpo_loss(self, batch: List[dict]):
"""
batch: list of dataset items (each a PromptWithAnswer.to_dict() augmented with
the fields the reward needs: 'gold_label' and 'facts' — see GRPO_INTEGRATION.md).
Returns (scalar loss, stats dict).
"""
cfg = self.cfg
total_loss = torch.zeros((), device=self.device)
n_groups = 0
stat_r, stat_ans, stat_faith = [], [], []
for item in batch:
completions = self._sample_rollouts(item)
rewards, r_ans, r_fai = [], [], []
for c in completions:
r = self.reward_fn(c, item)
rewards.append(r["r_total"]); r_ans.append(r["r_answer"]); r_fai.append(r["r_faith"])
rew = torch.tensor(rewards, device=self.device, dtype=torch.float32)
# Group-relative advantage: (r - mean) / (std + eps). Zero if no spread.
adv = (rew - rew.mean()) / (rew.std() + cfg.adv_eps)
group_loss = torch.zeros((), device=self.device)
for c, a in zip(completions, adv):
token_lp = self._sequence_logprob(item, c) # [A], grad-on
seq_lp = token_lp.sum()
pg = -a.detach() * seq_lp # policy gradient
if cfg.kl_coef > 0 and self.ref_model is not None:
with torch.no_grad():
ref_lp = self._sequence_logprob(item, c, model=self.ref_model)
# k3 KL estimator (per-token), summed
diff = (ref_lp - token_lp.detach())
kl = (torch.exp(diff) - diff - 1.0).sum()
pg = pg + cfg.kl_coef * kl
group_loss = group_loss + pg
total_loss = total_loss + group_loss / max(len(completions), 1)
n_groups += 1
stat_r.append(rew.mean().item())
stat_ans.append(sum(r_ans) / len(r_ans))
stat_faith.append(sum(r_fai) / len(r_fai))
loss = total_loss / max(n_groups, 1)
stats = {
"reward_mean": sum(stat_r) / max(len(stat_r), 1),
"answer_reward": sum(stat_ans) / max(len(stat_ans), 1),
"faith_reward": sum(stat_faith) / max(len(stat_faith), 1),
}
return loss, stats
|