Fix DPO IPO log-prob normalization

[taking-lying-flat]

Summary

• Keep DPO forward outputs as sequence-summed log-probabilities for every loss type.
• Return the valid completion-token count for each chosen and rejected sequence.
• Apply length normalization to both policy and reference log-probabilities only while evaluating an IPO loss.
• Preserve the existing sum-based log-probability metrics and add logps_mean/chosen and logps_mean/rejected when IPO is configured.
• Add regression coverage for padded, padding-free, packing, post-gather sequence-parallel, standalone IPO, and mixed-loss paths.

Root cause

IPO uses length-normalized completion log-probabilities, while standard sigmoid DPO uses sequence sums. The previous implementation normalized the shared per-token tensor before reducing it into per-sequence scores:

if 'ipo' in loss_types:
    size_completion = loss_mask.sum(dim=-1)
    per_token_logps = per_token_logps / size_completion

This caused three related problems.

Padding-free and packing used the whole packed row as the denominator

Padding-free batches have shapes similar to:

per_token_logps: [1, total_packed_tokens]
loss_mask:        [1, total_packed_tokens]

Consequently, loss_mask.sum(-1) produced one count for the entire packed row. Sequence boundaries from cu_seqlens were only applied afterward, so every chosen/rejected sequence was divided by the combined token count.

For example, with two completions:

seq0: sum_logp = -3, length = 2
seq1: sum_logp = -6, length = 6

The correct IPO scores are -3 / 2 = -1.5 and -6 / 6 = -1.0. The previous padding-free path returned -3 / 8 = -0.375 and -6 / 8 = -0.75. This changes the relative preference margin, not only its scale.

Packing uses the same flattened representation, and sequence-parallel gathering restores the full packed row before this reduction, so both paths were affected.

IPO contaminated other losses in a mixed configuration

loss_type can contain multiple weighted losses, for example:

loss_type = ['sigmoid', 'ipo']

The forward pass creates one shared pair of chosen_logps and rejected_logps, while the loss loop evaluates each configured loss afterward. Because the old forward normalization ran whenever IPO appeared anywhere in the list, sigmoid DPO also received normalized log-probabilities instead of sequence sums.

As a result, the mixed sigmoid component did not match standalone sigmoid DPO.

The ordinary padded IPO path could fail broadcasting

Outside padding-free mode, the tensors have shapes [B, L] and [B]. Dividing them directly is invalid for the usual case where B != L, producing a size-mismatch error. Moving normalization after sequence reduction also removes this invalid broadcast.

Fix

concatenated_forward() now always reduces per-token log-probabilities into sequence sums and returns the valid completion-token counts alongside them. In padding-free mode, each count is computed inside its cu_seqlens segment from the corresponding loss_mask slice.

The per-loss loop then:

1. uses the unmodified sums for sigmoid and all existing non-IPO losses;
2. divides policy and reference sequence sums by their own completion-token counts only for IPO;
3. clamps counts to at least one to avoid division by zero.

Reference precomputation continues to store sequence sums, so the same cached values work for standalone and mixed losses. Other loss types retain their existing sum-based behavior.

The existing logps/chosen and logps/rejected metrics remain sequence sums for compatibility. When IPO is present, the trainer additionally reports the mean values actually supplied to IPO as logps_mean/chosen and logps_mean/rejected.

The regression tests use completions with clearly different valid lengths and verify:

• ordinary padded forward aggregation;
• padding-free cu_seqlens segmentation;
• the packing flag path;
• the sequence-parallel post-gather tensor contract;
• standalone sigmoid and standalone IPO results;
• mixed loss equality to the weighted standalone components;
• zero-token count clamping.