[ci] layer-0 grad-norm regression for per-method training tests (5a-ii)

[alexzms]

Summary

Phase 2 / PR 5a-ii of the fastvideo.train CI plan. Layers a device-keyed layer-0 grad-norm regression on top of the per-method smoke tests landed in 5a-i (#1343).

5a-i asserts the loss is finite and block-0 grads are finite + non-zero. This PR adds a numeric fingerprint: after single_train_step + backward, the L2 norm of transformer block-0's trainable grads is compared against a per-GPU reference within rtol.

Changes

• fastvideo/tests/train/methods/grad_norm_regression.py — helper:
  • layer0_grad_norm(transformer) — global L2 norm of block-0 trainable grads.
  • resolve_device_key() — maps the GPU name to a reference key (L40S / GB200 / H100 / …).
  • check_grad_norm_regression(test_name, transformer, rtol=0.10) — asserts within rtol; skips (never hard-fails) on an unsupported or not-yet-seeded GPU; with FASTVIDEO_GRADNORM_UPDATE=1 records/refreshes the reference instead.
• test_wan_finetune.py / test_wan_causal_dfsft.py — call check_grad_norm_regression(...) after the existing finite/non-zero assertions.
• fastvideo/tests/train/methods/grad_norm_refs.json — in-repo references, seeded for GB200 (local dev) and L40S (the CI runner).
• fastvideo/tests/modal/pr_test.py — seed_grad_norm_references Modal entrypoint (L40S, FASTVIDEO_GRADNORM_UPDATE=1) to record/refresh a runner's reference.

Determinism & why device-keyed

method.on_train_start() seeds the global RNG + the method's cuda_generator (training.data.seed=42 in the fixture) and the synthetic raw_batch is built after that call, so the grad norm is reproducible — identical to 4 decimals across local runs. Inputs are synthetic (torch.randn text-embed + VAE latent, all-ones mask); model weights are the real Wan-AI/Wan2.1-T2V-1.3B-Diffusers checkpoint. The check is a forward/backward wiring + scale fingerprint, not a real-data quality test.

References are device-keyed because the norm differs across GPU architectures (bf16 accumulation, kernel choice). The causal/dfsft test spreads ~9% between GB200 (2.9781) and L40S (3.2562) — a single golden value couldn't cover both; rtol=0.10.

Grad-alignment coverage after this PR

model ＼ method	finetune	dfsft
Wan (WanModel)	✅
WanCausal (WanCausalModel)		✅

i.e. 2 (model, method) cells: WanModel/FineTuneMethod and WanCausalModel/DiffusionForcingSFTMethod. dmd2 / self_forcing / kd and other models (Hunyuan / Cosmos / LongCat / MatrixGame2) are out of scope here — see plan below.

Test plan

• Local (GB200): pytest fastvideo/tests/train/methods -v → 2 passed (asserting against GB200 references)
• L40S references seeded via modal run fastvideo/tests/modal/pr_test.py::seed_grad_norm_references
• /test train-framework — full GPU CI run on L40S (now asserts the grad-norm, no longer skips)

Plan context

Slice	Scope	Status
5a-i (#1343)	per-method test framework + 2 reference tests	merged
5a-ii (this PR)	layer-0 grad-norm regression + device-keyed refs	—
5b	dmd2 / self_forcing / kd (needs tiny teacher / fake-score fixtures)	—
5c	Hunyuan + finetune, more (model, method) combos	—

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 PR merge requirements

Wonderful, this rule succeeded.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces device-keyed gradient norm regression checks for training smoke tests, allowing validation of layer-0 trainable gradients against reference values pinned per GPU architecture. The feedback suggests several improvements to the implementation: making CUDA device name matching case-insensitive for better robustness, optimizing gradient norm calculation by accumulating squared sums on the GPU to avoid multiple CPU-GPU synchronizations, and explicitly specifying UTF-8 encoding when reading and writing the reference JSON file.

[Davids048]

Most looks good.
Just clearer documentation on the tested configs

[alexzms]

Pushed d64e4e94 addressing the review feedback:

@Davids048

• Documented how the GB200 values were obtained — expanded the module docstring of grad_norm_regression.py to spell out two seeding paths: CI / L40S via seed_grad_norm_references in pr_test.py, vs. local / non-L40S GPUs via FASTVIDEO_GRADNORM_UPDATE=1 pytest ... on the workstation. Also clarified in the Modal function's docstring that it is pinned to L40S:1 and seeds the L40S key only.
• Cleaned up _DEVICE_MAPPINGS — dropped the un-seeded H100 / H200 / A100 entries; only the keys present in grad_norm_refs.json remain (L40S, GB200, plus B200 → GB200 aliasing the same Blackwell arch).

@gemini-code-assist (nits)

• Case-insensitive substring match in resolve_device_key.
• layer0_grad_norm now accumulates squared sums on the GPU and does a single .item() at the end instead of one per parameter. Same math — recorded references still apply.
• encoding="utf-8" on the refs JSON read/write.

Re-requesting review when you have a moment.

[gemini-code-assist]

Code Review

This pull request introduces a device-keyed grad-norm regression check for training smoke tests. It adds a regression harness that compares the L2 norm of transformer block 0's trainable gradients against reference values stored in a JSON file (grad_norm_refs.json), supporting both local and CI environments (such as L40S and GB200 GPUs). The check is integrated into the Wan causal and finetuning tests, and a Modal helper function is added to seed references. Feedback on the changes suggests using a file lock when updating the reference JSON file to prevent race conditions during parallel test execution.

[gemini-code-assist]

When running tests in parallel (e.g., using pytest -n with pytest-xdist) with FASTVIDEO_GRADNORM_UPDATE=1 enabled, concurrent processes will attempt to read, modify, and write to the same grad_norm_refs.json file. This creates a race condition where one process can overwrite the updates made by another, or read a partially written file. \n\nTo prevent this, we can use a file lock to make the read-modify-write operation atomic. Since filelock is a standard dependency of packages like torch and huggingface_hub, it is guaranteed to be available in the environment.

        from filelock import FileLock\n        with FileLock(_REFS_PATH.with_suffix(\".lock\")):\n            refs = _load_refs()\n            refs.setdefault(test_name, {})[device_key] = round(norm, 4)\n            _save_refs(refs)

[Davids048]

ok

[alexzms]

/merge

[alexzms]

ok

Thanks for reviewing!