From ee5e2732a9515bcb62a2d4ce507c626a30d5e396 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Tue, 26 May 2026 00:38:18 +0000 Subject: [PATCH 1/2] [ci]: layer-0 grad-norm regression for per-method train tests (5a-ii) Phase 2 / 5a-ii: layers a device-keyed grad-norm check on top of the finite/non-zero grad assertions from 5a-i. After single_train_step + backward, block-0's trainable grad L2 norm is compared against a per-GPU reference in grad_norm_refs.json (in-repo). - grad_norm_regression.py: capture + device-key + compare/record helper. FASTVIDEO_GRADNORM_UPDATE=1 records the reference for the current GPU; unseeded/unsupported GPUs skip rather than fail. - Seeded references for GB200 (local dev) and L40S (the CI runner). The norm is stable to 4 decimals across runs; rtol=0.10. The cross-GPU spread on the causal/dfsft test (~9%) is why references are device-keyed. - pr_test.py: seed_grad_norm_references Modal entrypoint (L40S, UPDATE=1) to record/refresh a runner's reference. --- fastvideo/tests/modal/pr_test.py | 22 +++ .../tests/train/methods/grad_norm_refs.json | 10 ++ .../train/methods/grad_norm_regression.py | 134 ++++++++++++++++++ .../train/methods/test_wan_causal_dfsft.py | 6 + .../tests/train/methods/test_wan_finetune.py | 6 + 5 files changed, 178 insertions(+) create mode 100644 fastvideo/tests/train/methods/grad_norm_refs.json create mode 100644 fastvideo/tests/train/methods/grad_norm_regression.py diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index a932616e70..9b4792daa5 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -280,6 +280,28 @@ 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). + + Phase 2 / 5a-ii one-off seeding entrypoint. ``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 ``fastvideo/tests/train/methods/grad_norm_refs.json``. Re-run on + any new runner GPU to seed its key. + """ + 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..5a28fe1cc6 --- /dev/null +++ b/fastvideo/tests/train/methods/grad_norm_regression.py @@ -0,0 +1,134 @@ +# 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. CI runs +this suite on L40S; B200/GB200 is our local dev GPU. + +To seed (or refresh) the reference for the current GPU, run the test with +``FASTVIDEO_GRADNORM_UPDATE=1`` — it records the measured norm into +``grad_norm_refs.json`` and skips the assertion for that run. +""" + +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. +_DEVICE_MAPPINGS: tuple[tuple[str, str], ...] = ( + ("L40S", "L40S"), + ("GB200", "GB200"), + ("B200", "GB200"), + ("H100", "H100"), + ("H200", "H200"), + ("A100", "A100"), +) + + +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.""" + name = device_name if device_name is not None else _device_name() + for pattern, key in _DEVICE_MAPPINGS: + if pattern in name: + 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. + """ + blocks = getattr(transformer, "blocks", None) + assert blocks is not None and len(blocks) > 0, ( + "transformer is expected to expose a non-empty ``.blocks``") + sq_sum = 0.0 + for p in blocks[0].parameters(): + if p.requires_grad and p.grad is not None: + sq_sum += p.grad.detach().float().pow(2).sum().item() + return sq_sum**0.5 + + +def _load_refs() -> dict[str, dict[str, float]]: + if _REFS_PATH.exists(): + return json.loads(_REFS_PATH.read_text()) + return {} + + +def _save_refs(refs: dict[str, dict[str, float]]) -> None: + _REFS_PATH.write_text( + json.dumps(refs, indent=2, sort_keys=True) + "\n") + + +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) From d64e4e942ee8d277d3e1a70e2ba78d5198566948 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Fri, 29 May 2026 00:01:59 +0000 Subject: [PATCH 2/2] [ci]: address #1396 review feedback for grad-norm regression MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Document the two seeding paths in grad_norm_regression.py's module docstring (CI / L40S via Modal vs local / non-L40S via pytest + FASTVIDEO_GRADNORM_UPDATE=1) and clarify that seed_grad_norm_references in pr_test.py is L40S-only — answers Davids048's question about how the GB200 values were obtained. - Drop unused H100/H200/A100 mappings from _DEVICE_MAPPINGS; only devices with seeded references in grad_norm_refs.json remain. - Make resolve_device_key's substring match case-insensitive. - Accumulate layer0_grad_norm's squared sums on the GPU and sync once at the end instead of one .item() per parameter. - Specify encoding="utf-8" on _REFS_PATH read/write for cross-platform safety. --- fastvideo/tests/modal/pr_test.py | 19 ++++-- .../train/methods/grad_norm_regression.py | 62 +++++++++++++------ 2 files changed, 56 insertions(+), 25 deletions(-) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 9b4792daa5..7ea5923008 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -289,13 +289,20 @@ def run_train_framework_tests(): ], volumes={"/root/data": model_vol}) def seed_grad_norm_references(): - """Record the per-method grad-norm reference for the CI GPU (L40S). + """Record the per-method grad-norm reference for the **CI GPU (L40S only)**. - Phase 2 / 5a-ii one-off seeding entrypoint. ``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 ``fastvideo/tests/train/methods/grad_norm_refs.json``. Re-run on - any new runner GPU to seed its key. + 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" diff --git a/fastvideo/tests/train/methods/grad_norm_regression.py b/fastvideo/tests/train/methods/grad_norm_regression.py index 5a28fe1cc6..4ec5384586 100644 --- a/fastvideo/tests/train/methods/grad_norm_regression.py +++ b/fastvideo/tests/train/methods/grad_norm_regression.py @@ -13,12 +13,23 @@ 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. CI runs -this suite on L40S; B200/GB200 is our local dev GPU. +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). -To seed (or refresh) the reference for the current GPU, run the test with -``FASTVIDEO_GRADNORM_UPDATE=1`` — it records the measured norm into -``grad_norm_refs.json`` and skips the assertion for that run. +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 @@ -37,14 +48,13 @@ # scale regressions), not micro-drift from reduction nondeterminism. _DEFAULT_RTOL = 0.10 -# GPU-name substring -> reference key. First match wins. +# 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"), - ("H100", "H100"), - ("H200", "H200"), - ("A100", "A100"), + ("B200", "GB200"), # same Blackwell arch as GB200 ) @@ -55,10 +65,15 @@ def _device_name() -> str: def resolve_device_key(device_name: str | None = None) -> str | None: - """Map a CUDA device name to its reference key, or None if unsupported.""" + """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 in name: + if pattern.lower() in name_lower: return key return None @@ -69,26 +84,35 @@ def layer0_grad_norm(transformer) -> float: 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``") - sq_sum = 0.0 - for p in blocks[0].parameters(): - if p.requires_grad and p.grad is not None: - sq_sum += p.grad.detach().float().pow(2).sum().item() - return sq_sum**0.5 + 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()) + 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") + json.dumps(refs, indent=2, sort_keys=True) + "\n", + encoding="utf-8") def check_grad_norm_regression(