From adf87c2a0a44a9223bdd6bbe06779c453e9b351d Mon Sep 17 00:00:00 2001 From: sam-paech <152407511+sam-paech@users.noreply.github.com> Date: Fri, 24 Apr 2026 21:03:24 +1000 Subject: [PATCH] ftpo fixes --- .gitignore | 3 ++- core/ftpo_trainer.py | 47 +++++++++++++++++++++----------------------- 2 files changed, 24 insertions(+), 26 deletions(-) diff --git a/.gitignore b/.gitignore index 08d4de3..a911299 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ experiments *.pyc unsloth_compiled_cache results -bct.txt \ No newline at end of file +bct.txt +.codex \ No newline at end of file diff --git a/core/ftpo_trainer.py b/core/ftpo_trainer.py index 2b0009e..7c981cd 100644 --- a/core/ftpo_trainer.py +++ b/core/ftpo_trainer.py @@ -209,7 +209,7 @@ class FTPOTrainer(DPOTrainer): attn = inputs["attention_mask"].to(device) # [B,L] B, L = ids.shape - seq_len = attn.sum(1) + seq_len = attn.sum(1) pad_off = (L - seq_len).unsqueeze(1) arange_L = torch.arange(L, device=ids.device).unsqueeze(0) pos_full = (arange_L - pad_off).clamp(min=0) @@ -232,27 +232,24 @@ class FTPOTrainer(DPOTrainer): logp_bad = logp_all.gather(-1, rejected.unsqueeze(-1)).squeeze(-1) batch_rows = torch.arange(B, device=logp_all.device).unsqueeze(1) - gathered = logits_last[batch_rows, ch_ids] - logit_bad = logits_last.gather(-1, rejected.unsqueeze(-1)) - margin = gathered - logit_bad - weights = torch.clamp((clip_epsilon_logits - margin) / clip_epsilon_logits, 0.0, 1.0) * ch_mask + delta_tok = logits_last[batch_rows, ch_ids] - logits_last.gather( + -1, rejected.unsqueeze(-1) + ) + weights = ( + torch.clamp( + (clip_epsilon_logits - delta_tok) / clip_epsilon_logits, + 0.0, + 1.0, + ) + * ch_mask + ) - zero_row = weights.sum(dim=-1, keepdim=True) < 1e-12 - weights = torch.where(zero_row, ch_mask.float(), weights) + tau = 1.0 + gap = clip_epsilon_logits - delta_tok + per_tok_loss = F.softplus(gap / tau) - weights_sum = weights.sum(dim=-1, keepdim=True) - batch_rows = torch.arange(B, device=ids.device).unsqueeze(1) - - l_chosen = logits_last[batch_rows, ch_ids] - l_bad = logits_last.gather(-1, rejected.unsqueeze(-1)) - delta_tok = l_chosen - l_bad - - margin = clip_epsilon_logits - tau = 1.0 - gap = margin - delta_tok - per_tok_loss = F.softplus(gap / tau) - - pref_loss = (per_tok_loss * weights).sum() / weights_sum.sum() + chosen_counts = ch_mask.sum(dim=-1).clamp(min=1) + pref_loss = ((per_tok_loss * weights).sum(dim=-1) / chosen_counts).mean() extra_metrics = {} @@ -270,13 +267,13 @@ class FTPOTrainer(DPOTrainer): use_cache=False, return_dict=True, ).logits[:, -1, :] - freeze_mask = torch.ones_like(logits_last, dtype=torch.bool) + tether_mask = torch.ones_like(logits_last, dtype=torch.bool) rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids) - freeze_mask[rows[ch_mask], ch_ids[ch_mask]] = False - freeze_mask.scatter_(1, rejected.unsqueeze(-1), False) + tether_mask[rows[ch_mask], ch_ids[ch_mask]] = False + tether_mask.scatter_(1, rejected.unsqueeze(-1), False) diff = logits_last - ref_logits_last - mse_elem_raw = (freeze_mask * diff.pow(2)).sum() / freeze_mask.sum() + mse_elem_raw = (tether_mask * diff.pow(2)).sum() / tether_mask.sum() tgt_mask = torch.zeros_like(logits_last, dtype=torch.bool) rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids) @@ -328,4 +325,4 @@ class FTPOTrainer(DPOTrainer): # ---------------------------------------------------------- def _prepare_dataset(self, dataset, *args, **_): - return dataset \ No newline at end of file + return dataset