diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 1890d30334..e35c589413 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -286,6 +286,35 @@ def run_train_framework_tests(): ) +@app.function(gpu="L40S:1", + image=image, + timeout=1800, + secrets=[ + modal.Secret.from_dict( + {"HF_API_KEY": os.environ.get("HF_API_KEY", "")}) + ], + volumes={"/root/data": model_vol}) +def seed_grad_norm_references(): + """Record the per-method grad-norm reference for the **CI GPU (L40S only)**. + + Phase 2 / 5a-ii one-off seeding entrypoint. Pinned to ``gpu="L40S:1"`` (the + Modal CI runner), so this function only seeds the ``L40S`` key in + ``fastvideo/tests/train/methods/grad_norm_refs.json``. + + ``FASTVIDEO_GRADNORM_UPDATE=1`` makes ``check_grad_norm_regression`` record + the measured norm instead of asserting; ``-rs`` surfaces the recorded value + in the log so it can be copied into the JSON. + + To seed any other device (e.g. our local Blackwell dev box → ``GB200`` + key), run the same env-var + pytest invocation directly on that + workstation — see the module docstring of ``grad_norm_regression.py`` for + the local command and the ``_DEVICE_MAPPINGS`` table. + """ + run_test( + "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs" + ) + + @app.function(gpu="L40S:1", image=image, timeout=3600, diff --git a/fastvideo/tests/train/methods/grad_norm_refs.json b/fastvideo/tests/train/methods/grad_norm_refs.json new file mode 100644 index 0000000000..f3c88fcedd --- /dev/null +++ b/fastvideo/tests/train/methods/grad_norm_refs.json @@ -0,0 +1,10 @@ +{ + "test_wan_causal_dfsft": { + "GB200": 2.9781, + "L40S": 3.2562 + }, + "test_wan_finetune": { + "GB200": 1.6486, + "L40S": 1.6467 + } +} diff --git a/fastvideo/tests/train/methods/grad_norm_regression.py b/fastvideo/tests/train/methods/grad_norm_regression.py new file mode 100644 index 0000000000..4ec5384586 --- /dev/null +++ b/fastvideo/tests/train/methods/grad_norm_regression.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Layer-0 grad-norm regression for the per-method training smoke tests. + +Phase 2 / 5a-ii: layers a device-keyed grad-norm check on top of the +finite/non-zero grad assertions established in 5a-i. After one +``single_train_step`` + ``backward``, the L2 norm of transformer block 0's +trainable gradients is compared against a reference value pinned per GPU in +``grad_norm_refs.json`` (next to this module). + +Determinism: the harness seeds both the global RNG and the method's +``cuda_generator`` via ``method.on_train_start()`` (``training.data.seed`` in the +fixture), and the synthetic ``raw_batch`` is built *after* that call, so the +forward/backward is reproducible within bf16 reduction noise on a given GPU. + +Why device-keyed: grad norms differ across GPU architectures (kernels, +accumulation order), so a single golden value can't cover every runner. The +JSON currently carries refs for the two GPUs we actually run on — ``L40S`` (CI) +and ``GB200`` (our Blackwell dev box; ``B200`` maps to the same key). + +Seeding a reference for the current device: + +- **CI / L40S** — invoke ``modal run`` against ``seed_grad_norm_references`` in + ``fastvideo/tests/modal/pr_test.py`` (pinned to ``gpu="L40S:1"``), then copy + the recorded value from the log into ``grad_norm_refs.json``. +- **Local / non-L40S GPUs** — on that workstation:: + + FASTVIDEO_GRADNORM_UPDATE=1 \\ + pytest fastvideo/tests/train/methods -vs -rs + + The harness writes the measured norm into ``grad_norm_refs.json`` under the + device's key and skips the assertion for that run. Append a new substring + entry to ``_DEVICE_MAPPINGS`` first for any device not already listed. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +import torch + +_REFS_PATH = Path(__file__).resolve().parent / "grad_norm_refs.json" +_UPDATE_ENV = "FASTVIDEO_GRADNORM_UPDATE" + +# bf16 single-step smoke: catch gross breakage (wrong wiring, dead grads, +# scale regressions), not micro-drift from reduction nondeterminism. +_DEFAULT_RTOL = 0.10 + +# GPU-name substring -> reference key. First match wins. Only devices with +# seeded references in ``grad_norm_refs.json`` are listed here — to add a new +# GPU, append an entry, then seed the reference (see module docstring). +_DEVICE_MAPPINGS: tuple[tuple[str, str], ...] = ( + ("L40S", "L40S"), + ("GB200", "GB200"), + ("B200", "GB200"), # same Blackwell arch as GB200 +) + + +def _device_name() -> str: + if not torch.cuda.is_available(): + return "CPU" + return torch.cuda.get_device_name(0) + + +def resolve_device_key(device_name: str | None = None) -> str | None: + """Map a CUDA device name to its reference key, or None if unsupported. + + The substring match is case-insensitive so it survives driver/environment + differences in how ``torch.cuda.get_device_name`` capitalizes the model. + """ + name = device_name if device_name is not None else _device_name() + name_lower = name.lower() + for pattern, key in _DEVICE_MAPPINGS: + if pattern.lower() in name_lower: + return key + return None + + +def layer0_grad_norm(transformer) -> float: + """Global L2 norm of transformer block 0's trainable gradients. + + Block 0 is the reference surface 5a-i already isolates: its grad is the + *last* one produced during backprop, so a healthy value implies the whole + forward + chain-rule path is intact. + + Accumulates the squared sums on the GPU and does a single CPU-GPU sync + (``.item()``) at the end, rather than one per parameter. + """ + blocks = getattr(transformer, "blocks", None) + assert blocks is not None and len(blocks) > 0, ( + "transformer is expected to expose a non-empty ``.blocks``") + grads = [ + p.grad for p in blocks[0].parameters() + if p.requires_grad and p.grad is not None + ] + if not grads: + return 0.0 + sq_sum = torch.zeros((), device=grads[0].device, dtype=torch.float32) + for g in grads: + sq_sum += g.detach().float().pow(2).sum() + return sq_sum.sqrt().item() + + +def _load_refs() -> dict[str, dict[str, float]]: + if _REFS_PATH.exists(): + return json.loads(_REFS_PATH.read_text(encoding="utf-8")) + return {} + + +def _save_refs(refs: dict[str, dict[str, float]]) -> None: + _REFS_PATH.write_text( + json.dumps(refs, indent=2, sort_keys=True) + "\n", + encoding="utf-8") + + +def check_grad_norm_regression( + test_name: str, + transformer, + *, + rtol: float = _DEFAULT_RTOL, +) -> None: + """Assert block-0 grad norm matches the device-keyed reference within rtol. + + - Skips when the current GPU has no reference (unsupported device, or not + yet seeded) so a new runner never hard-fails before its golden exists. + - With ``FASTVIDEO_GRADNORM_UPDATE=1`` records/updates the reference for the + current device instead of asserting. + """ + norm = layer0_grad_norm(transformer) + device_key = resolve_device_key() + + if os.environ.get(_UPDATE_ENV) == "1": + if device_key is None: + pytest.skip( + f"{_UPDATE_ENV}=1 but GPU '{_device_name()}' has no reference " + "key; add it to _DEVICE_MAPPINGS first") + refs = _load_refs() + refs.setdefault(test_name, {})[device_key] = round(norm, 4) + _save_refs(refs) + pytest.skip( + f"recorded grad-norm reference {test_name}[{device_key}] = " + f"{norm:.4f} (assertion skipped under {_UPDATE_ENV}=1)") + + ref = _load_refs().get(test_name, {}).get(device_key) \ + if device_key is not None else None + if ref is None: + pytest.skip( + f"no grad-norm reference for {test_name} on '{_device_name()}' " + f"(device_key={device_key}); run with {_UPDATE_ENV}=1 to seed it") + + rel = abs(norm - ref) / (abs(ref) + 1e-12) + assert rel <= rtol, ( + f"{test_name}[{device_key}] grad-norm regression: got {norm:.4f}, " + f"reference {ref:.4f}, relative error {rel:.3%} exceeds rtol " + f"{rtol:.0%}. If this is an intentional change, refresh the reference " + f"with {_UPDATE_ENV}=1 and explain why in the PR.") diff --git a/fastvideo/tests/train/methods/test_wan_causal_dfsft.py b/fastvideo/tests/train/methods/test_wan_causal_dfsft.py index 429c98316e..2da186e688 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_dfsft.py +++ b/fastvideo/tests/train/methods/test_wan_causal_dfsft.py @@ -28,6 +28,8 @@ from fastvideo.train.models.wan import WanCausalModel from fastvideo.train.utils.config import load_run_config +from .grad_norm_regression import check_grad_norm_regression + _FIXTURE = str( Path(__file__).resolve().parent.parent / "fixtures" @@ -122,3 +124,7 @@ def test_wan_causal_dfsft_single_train_step( assert any_nonzero, ( "all layer-0 grads are exactly zero; backward did not " "reach the first transformer block") + + # 5a-ii: device-keyed grad-norm regression on top of the same harness. + # Skips when the current GPU has no seeded reference. + check_grad_norm_regression("test_wan_causal_dfsft", model.transformer) diff --git a/fastvideo/tests/train/methods/test_wan_finetune.py b/fastvideo/tests/train/methods/test_wan_finetune.py index 59a4c2876a..f5b25414dc 100644 --- a/fastvideo/tests/train/methods/test_wan_finetune.py +++ b/fastvideo/tests/train/methods/test_wan_finetune.py @@ -36,6 +36,8 @@ from fastvideo.train.models.wan import WanModel from fastvideo.train.utils.config import load_run_config +from .grad_norm_regression import check_grad_norm_regression + _FIXTURE = str( Path(__file__).resolve().parent.parent / "fixtures" @@ -139,3 +141,7 @@ def test_wan_finetune_single_train_step( assert any_nonzero, ( "all layer-0 grads are exactly zero; backward did not " "reach the first transformer block") + + # 5a-ii: device-keyed grad-norm regression on top of the same harness. + # Skips when the current GPU has no seeded reference. + check_grad_norm_regression("test_wan_finetune", model.transformer)