Fix cross-rank metric averaging in the experimental self-distillation trainers

[ucalyptus]

Split out of #6237, which grew to cover several unrelated changes. This is the first piece.

Two fixes that SDFTTrainer and SDPOTrainer share:

• The distillation-loss metric averaged each rank's own masked mean with an unweighted gather(...).mean(). When completion lengths vary, ranks hold different numbers of valid tokens, so a rank with few valid tokens counted as much as one with many. The logged value was biased. _global_masked_mean sums the local numerator and denominator and divides once, which is what GRPO and RLOO do.
• local_sum is cast to float32 before torch.stack. values can be bf16 or fp16 under mixed precision while the count is float32, and torch.stack needs matching dtypes. The cast also avoids precision loss when many low-precision values are summed.
• train_dataset is checked for type at construction, so passing a list or a DataFrame fails with a clear message rather than later.

Both trainers get the same code, per the repository's consistency rule for duplicated logic.

Tests

CPU, transformers 5.x, trl-internal-testing/tiny-Qwen2ForCausalLM-2.5:

pytest tests/experimental/test_sdft_trainer.py -q                       18 passed
pytest tests/experimental/test_sdpo_trainer.py -q                       21 passed
pytest tests/experimental/test_self_distillation_trainer_behavior.py -q 12 passed

---

Note

Low Risk
Changes affect logged metrics and constructor validation only; optimization loss computation is untouched.

Overview
Fixes distributed logging for self-distillation metrics in SDFTTrainer and SDPOTrainer so masked means are token-weighted across ranks instead of averaging per-rank means with gather(...).mean().

A local global_masked_mean helper sums masked values and valid-token counts per rank, accelerator.reduces them, and divides once—matching GMPO-style aggregation. This applies to logged distillation loss (local and vLLM server paths) and the server logprob abs-diff diagnostic. Training loss math is unchanged.

Adds early train_dataset type checks: SDFTTrainer requires a Hugging Face Dataset; SDPOTrainer accepts Dataset or IterableDataset, with a clear TypeError for invalid inputs.

^{Reviewed by Cursor Bugbot for commit 17dc011. Bugbot is set up for automated code reviews on this repo. Configure here.}

[qgallouedec]

Thanks! Pushed 7dcf875 to inline global_masked_mean like GRPO.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 1925617. Configure here.}

[cursor]

Mixed-precision metric average dtype crash

High Severity

global_masked_mean stacks local_sum with local_count after casting only the count with .float(). With bf16 or fp16 losses, the masked sum stays low precision, so torch.stack raises and the step dies while logging distillation metrics. The same helper is copied in both trainers, including the server-loss paths.

Additional Locations (2)

• trl/experimental/sdft/sdft_trainer.py#L956-L960
• trl/experimental/sdpo/sdpo_trainer.py#L1263-L1267

 

^{Reviewed by Cursor Bugbot for commit 1925617. Configure here.}