142 lines
5.6 KiB
Python
142 lines
5.6 KiB
Python
# core/dpo_trainer.py
|
|
|
|
from __future__ import annotations
|
|
from typing import List, Optional
|
|
|
|
from core.finetuning import DPOTrainer, torch, F
|
|
|
|
|
|
class DPOTrainerWithChoiceWin(DPOTrainer):
|
|
"""
|
|
Drop-in DPO trainer that adds a token-level `chosen_win` metric for
|
|
dpo_final_token batches without altering DPO loss/behavior.
|
|
|
|
Metric (per-row):
|
|
Given the prompt context, compare the model's next-token probabilities:
|
|
win = 1{ log p(chosen_first | prompt) > log p(rejected_first | prompt) }
|
|
chosen_win = mean over rows of win.
|
|
|
|
Expected batch keys (dpo_final_token):
|
|
- chosen_input_ids, chosen_attention_mask
|
|
- rejected_input_ids, rejected_attention_mask
|
|
- prompt_input_ids (+ prompt_attention_mask) OR
|
|
prompt_ids (+ attention_mask)
|
|
"""
|
|
|
|
def compute_loss(self, model, inputs, return_outputs: bool = False, **kwargs):
|
|
# 1) standard DPO loss (unchanged)
|
|
loss = super().compute_loss(model, inputs, return_outputs=False, **kwargs)
|
|
|
|
# 2) metric (only if dpo_final_token fields are present)
|
|
needed = {
|
|
"chosen_input_ids", "chosen_attention_mask",
|
|
"rejected_input_ids", "rejected_attention_mask",
|
|
}
|
|
if needed.issubset(inputs.keys()):
|
|
chosen_win = self._metric_dpo_final_token(model, inputs)
|
|
if chosen_win is not None:
|
|
self.store_metrics({"chosen_win": chosen_win}, train_eval="train")
|
|
|
|
if return_outputs:
|
|
return loss, {}
|
|
return loss
|
|
|
|
# ----------------------------- helpers ---------------------------------
|
|
|
|
@staticmethod
|
|
def _first_real_index(mask_row: torch.Tensor) -> Optional[int]:
|
|
"""Index of the first non-pad token according to an attention mask row."""
|
|
nz = mask_row.nonzero(as_tuple=False)
|
|
return int(nz[0].item()) if nz.numel() else None
|
|
|
|
# ---------------------- core metric (dpo_final_token) -------------------
|
|
|
|
def _metric_dpo_final_token(self, model, inputs) -> Optional[torch.Tensor]:
|
|
device = next(model.parameters()).device
|
|
tok = getattr(self, "tokenizer", getattr(self, "processing_class", None))
|
|
pad_id = getattr(tok, "pad_token_id", 0)
|
|
|
|
# Continuations
|
|
ch_ids = inputs["chosen_input_ids"].to(device) # [B, Lc] (continuation only)
|
|
ch_am = inputs["chosen_attention_mask"].to(device) # [B, Lc]
|
|
rj_ids = inputs["rejected_input_ids"].to(device) # [B, Lr]
|
|
rj_am = inputs["rejected_attention_mask"].to(device) # [B, Lr]
|
|
|
|
# Prompt (prefer explicit prompt_input_ids → fall back to prompt_ids)
|
|
pr_ids_key = "prompt_input_ids" if "prompt_input_ids" in inputs else (
|
|
"prompt_ids" if "prompt_ids" in inputs else None
|
|
)
|
|
pr_am_key = "prompt_attention_mask" if "prompt_attention_mask" in inputs else (
|
|
"attention_mask" if "attention_mask" in inputs else None
|
|
)
|
|
|
|
if pr_ids_key is None:
|
|
# No prompt in the batch → cannot compute the metric reliably
|
|
return None
|
|
|
|
pr_ids_full = inputs[pr_ids_key].to(device) # [B, Lp]
|
|
if pr_am_key in inputs:
|
|
pr_am_full = inputs[pr_am_key].to(device) # [B, Lp]
|
|
else:
|
|
# derive attention mask from pad id
|
|
pr_am_full = pr_ids_full.ne(pad_id).to(pr_ids_full.dtype)
|
|
|
|
B = ch_ids.size(0)
|
|
|
|
# Assemble per-row prompt and first continuation tokens
|
|
prompts: List[torch.Tensor] = []
|
|
chosen_first: List[int] = []
|
|
rejected_first: List[int] = []
|
|
last_idx: List[int] = []
|
|
|
|
for b in range(B):
|
|
# first token of each continuation (continuations are typically 1 token + EOS)
|
|
k_ch = self._first_real_index(ch_am[b])
|
|
k_rj = self._first_real_index(rj_am[b])
|
|
if k_ch is None or k_rj is None:
|
|
continue
|
|
|
|
ch_first = int(ch_ids[b, k_ch].item())
|
|
rj_first = int(rj_ids[b, k_rj].item())
|
|
|
|
# prompt (all real tokens)
|
|
pr_mask = pr_am_full[b].bool()
|
|
pr_seq = pr_ids_full[b][pr_mask] # 1D tensor with real prompt ids
|
|
if pr_seq.numel() == 0:
|
|
continue
|
|
|
|
prompts.append(pr_seq)
|
|
chosen_first.append(ch_first)
|
|
rejected_first.append(rj_first)
|
|
last_idx.append(pr_seq.numel() - 1)
|
|
|
|
n = len(prompts)
|
|
if n == 0:
|
|
return None
|
|
|
|
# Right-pad prompts into a dense batch for one forward pass
|
|
max_pr = max(p.numel() for p in prompts)
|
|
prompt_ids = pr_ids_full.new_full((n, max_pr), pad_id)
|
|
prompt_am = pr_am_full.new_zeros((n, max_pr))
|
|
|
|
for i, p in enumerate(prompts):
|
|
Lp = p.numel()
|
|
prompt_ids[i, :Lp] = p
|
|
prompt_am [i, :Lp] = 1
|
|
|
|
last_idx_t = torch.tensor(last_idx, device=device, dtype=torch.long)
|
|
chosen_tok = torch.tensor(chosen_first, device=device, dtype=torch.long)
|
|
reject_tok = torch.tensor(rejected_first, device=device, dtype=torch.long)
|
|
|
|
# One no-grad forward on prompts; evaluate next-token distribution at last prompt token
|
|
with torch.no_grad():
|
|
out = model(prompt_ids, attention_mask=prompt_am, use_cache=False, return_dict=True)
|
|
logits_last = out.logits[torch.arange(n, device=device), last_idx_t, :] # [n, V]
|
|
logp_last = F.log_softmax(logits_last, dim=-1)
|
|
|
|
lp_good = logp_last.gather(1, chosen_tok.unsqueeze(1)).squeeze(1) # [n]
|
|
lp_bad = logp_last.gather(1, reject_tok.unsqueeze(1)).squeeze(1) # [n]
|
|
wins = (lp_good > lp_bad).float()
|
|
|
|
return wins.mean().detach()
|