[taking-lying-flat]

Summary

• keep one reward-model margin per preference pair by filling missing margins with 0.0
• validate the number of margins before computing the pairwise loss
• reshape margins to match scalar reward outputs in both the Transformers and Megatron reward trainers

Root cause

The RLHF collator previously dropped samples whose margin was missing, which could misalign the remaining margins with their original preference pairs.

For a standard scalar reward head, chosen and rejected rewards have shape [B, 1], while the collator produced margins with shape [B]. Subtracting these tensors broadcasts to [B, B], silently mixing every reward pair with every margin instead of computing one loss term per pair.

Impact

Reward-model training now applies each margin only to its corresponding chosen/rejected pair. Samples without an explicit margin use a zero margin, preserving the original batch ordering.

Validation

• CPU in-memory smoke test verified the loss against an explicit margin[:, None] calculation
• mixed-margin batch [0.1, None, 0.3] produced [0.1, 0.0, 0.3]
• project pre-commit checks passed for all three changed files
• git diff --check passed