Add training-invariance tests

[qgallouedec]

Motivation

Unit tests don't catch trainer-level numerical drift: grad-accum normalization issue, attention-impl divergence (eager ↔ FA2 / kernels), etc. These silently shift the loss trajectory; users only notice when their run no longer reproduces. This PR adds a small opt-in suite that catches that class of bug.

Some context about issues that could be caught thanks to this kind of tests

• Last year highly visible "gradient accumulation bug": Fix Gradient Accumulation issue transformers#34191
• "We found two bugs in deep speed": https://arxiv.org/abs/2604.23747

Important design decisions.

• Equivalence classes. Configs that should produce the same trajectory are grouped (e.g. sft_default, sft_pdb1_gas8, sft_attn_fa2_kernels all in class sft). Each class has one canonical config that owns the reference snapshot; every other config asserts it matches that one snapshot. This catches actual invariance breakage.
• Hardware pin: H100 80GB: full_determinism is per-device, not cross-device; different GPUs pick different cuBLAS/cuDNN kernels, FlashAttention tile sizes, and reduction orders, producing small-but-consistent loss differences that exceed the snapshot tolerance.
• Training on real model (not tiny): tiny models are randomly initialized, so the systematic O(1e-3) loss shifts these tests are designed to catch get drowned in init noise; a real pretrained model gives a smooth, reproducible loss curve where any drift sticks out.

Scope

sft and dpo, two invariance axes each (grad-accum × attn-impl). 50 steps, fp32, fixed seed. Intentionally minimal coverage grows in follow-ups once the design is settled.

The main concern here are not the tests themselves but the design!

---

Note

Low Risk
Adds a new opt-in pytest suite that runs short end-to-end training via the trl CLI and compares against committed reference trajectories; main risk is test brittleness/maintenance and large snapshot updates rather than runtime behavior.

Overview
Introduces a new opt-in pytest marker invariant and a tests/invariant/ suite that runs short end-to-end trl sft/trl dpo trainings (via subprocess) and asserts their loss/grad_norm trajectories match committed reference snapshots within tolerances.

Adds a small harness to record/load/compare trajectories and checks equivalence classes (e.g., default vs gradient-accumulation variants), plus initial reference JSON snapshots and documentation for recording/updating them.

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

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0e0e29bd30

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[qgallouedec]

Note that currently, the tests are failing, I'm trying to understand why:

FAILED tests/invariant/test_invariant.py::test_invariant[sft_pdb1_gas8] - AssertionError: 'sft_pdb1_gas8' diverges from class 'sft' reference:
FAILED tests/invariant/test_invariant.py::test_invariant[dpo_pdb1_gas8] - AssertionError: 'dpo_pdb1_gas8' diverges from class 'dpo' reference:

Claude suggests that it's because fp32 addition is non-associative, so PDB=8 and PDB=1×GAS=8 reduce gradients in different orders and diverge by O(1e-7) at step 1, and from there training's chaotic dynamics amplify that drift exponentially across optimizer steps, producing the kind of grad-norm spike we saw at step 6.

I'll check if it's true by hand, but it's a bit out of the scope of this PR

[albertvillanova]

Thanks. Increasing test coverage is always a good direction. Additionally, we currently don't have any insight about the training: current tests treat trainers as black boxes.

[cursor]

Cursor Bugbot has reviewed your changes 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 97a0161. Configure here.}