From 0df1445d0dbfcafc6dec827e7be9ec6c267d2c24 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Tue, 12 May 2026 16:54:37 -0700 Subject: [PATCH] [ci] add GPU model loading tests for fastvideo.train (PR 4/9) (#1274) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/pipeline.yml | 26 ++++++ .buildkite/scripts/pr_test.sh | 4 + .github/workflows/ci-slash-commands.yml | 3 +- docs/contributing/ci_architecture.md | 2 + docs/contributing/testing.md | 3 +- fastvideo/tests/modal/pr_test.py | 16 +++- fastvideo/tests/train/fixtures/__init__.py | 1 + .../tests/train/fixtures/hunyuan_t2v_min.yaml | 17 ++++ .../train/fixtures/wan_causal_t2v_min.yaml | 17 ++++ .../tests/train/fixtures/wan_t2v_min.yaml | 21 +++++ fastvideo/tests/train/models/__init__.py | 1 + .../tests/train/models/test_load_hunyuan.py | 77 +++++++++++++++++ fastvideo/tests/train/models/test_load_wan.py | 82 +++++++++++++++++++ .../train/models/test_load_wan_causal.py | 56 +++++++++++++ 14 files changed, 323 insertions(+), 3 deletions(-) create mode 100644 fastvideo/tests/train/fixtures/__init__.py create mode 100644 fastvideo/tests/train/fixtures/hunyuan_t2v_min.yaml create mode 100644 fastvideo/tests/train/fixtures/wan_causal_t2v_min.yaml create mode 100644 fastvideo/tests/train/fixtures/wan_t2v_min.yaml create mode 100644 fastvideo/tests/train/models/__init__.py create mode 100644 fastvideo/tests/train/models/test_load_hunyuan.py create mode 100644 fastvideo/tests/train/models/test_load_wan.py create mode 100644 fastvideo/tests/train/models/test_load_wan_causal.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index f53830bf1..641938df7 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -195,6 +195,17 @@ steps: limit: 2 agents: queue: "default" + - label: ":test_tube: Train Framework Tests" + if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "train_framework" + command: "timeout 90m .buildkite/scripts/pr_test.sh" + retry: + automatic: + - exit_status: 128 + limit: 3 + - exit_status: -1 + limit: 2 + agents: + queue: "default" # ============================================================ # Fastcheck: Runs on every PR (~10-15 min parallel) @@ -436,3 +447,18 @@ steps: - TEST_TYPE=api_server agents: queue: "default" + - path: + - "fastvideo/train/**" + - "fastvideo/tests/train/models/**" + - "fastvideo/tests/train/fixtures/**" + - "fastvideo/models/dits/**" + - "fastvideo/models/loader/**" + - "pyproject.toml" + - "docker/Dockerfile.python3.12" + config: + command: "timeout 30m .buildkite/scripts/pr_test.sh" + label: ":test_tube: Train Framework Tests" + env: + - TEST_TYPE=train_framework + agents: + queue: "default" diff --git a/.buildkite/scripts/pr_test.sh b/.buildkite/scripts/pr_test.sh index 22817fdf5..ec18ddbdd 100755 --- a/.buildkite/scripts/pr_test.sh +++ b/.buildkite/scripts/pr_test.sh @@ -197,6 +197,10 @@ case "$TEST_TYPE" in log "Running unit tests..." MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_unit_test" ;; + "train_framework") + log "Running fastvideo.train framework tests..." + MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_train_framework_tests" + ;; "lora_extraction") log "Running LoRA extraction tests..." MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_lora_extraction_tests" diff --git a/.github/workflows/ci-slash-commands.yml b/.github/workflows/ci-slash-commands.yml index 9ff729cd2..e8fca04c0 100644 --- a/.github/workflows/ci-slash-commands.yml +++ b/.github/workflows/ci-slash-commands.yml @@ -125,7 +125,7 @@ jobs: set -euo pipefail TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true) - VALID="encoder vae transformer kernel unit ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api full fastcheck pre-commit" + VALID="encoder vae transformer kernel unit ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api train-framework full fastcheck pre-commit" if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then echo "Unknown test: '$TEST_NAME'. Valid: $VALID" exit 1 @@ -139,6 +139,7 @@ jobs: [distillation]=distillation_dmd [self-forcing]=self_forcing [vsa]=training_vsa [vmoba]=inference_vmoba [performance]=performance [api]=api_server + [train-framework]=train_framework ) if [ "$TEST_NAME" = "full" ]; then diff --git a/docs/contributing/ci_architecture.md b/docs/contributing/ci_architecture.md index fe9c07c38..a15e81d0c 100644 --- a/docs/contributing/ci_architecture.md +++ b/docs/contributing/ci_architecture.md @@ -101,6 +101,7 @@ failing test's output. | Inference Tests VMoBA | `inference_vmoba` | 15 min | | Performance Tests | `performance` | 30 min | | API Server Tests | `api_server` | 30 min | +| Train Framework Tests | `train_framework` | 30 min | If a Full Suite test fails, check the Buildkite build log for the failing step's output. Fix the regression, push, and comment `/merge` again to re-trigger. @@ -277,6 +278,7 @@ Triggers a specific Buildkite test or suite on the current PR branch. | `/test vmoba` | VMoBA inference tests | `inference_vmoba` | | `/test performance` | Performance benchmarks | `performance` | | `/test api` | API server integration tests | `api_server` | +| `/test train-framework` | `fastvideo.train` GPU model loading + per-method tests | `train_framework` | | `/test full` | Entire Full Suite | all (with `TEST_SCOPE=full`) | | `/test fastcheck` | Entire Fastcheck suite | fastcheck (with `TEST_SCOPE=fastcheck`) | | `/test pre-commit` | Pre-commit checks on PR code | — (runs `ci-precommit.yml` via `workflow_call`) | diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index c68deb7c3..3a32de875 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -4,8 +4,9 @@ This guide explains how to add and run tests in FastVideo. The testing suite is ## Test Types -* **Unit Tests**: Located in `fastvideo/tests/dataset`, `fastvideo/tests/entrypoints`, and `fastvideo/tests/workflow`. These test individual functions and classes. +* **Unit Tests**: Located in `fastvideo/tests/api`, `fastvideo/tests/dataset`, `fastvideo/tests/entrypoints`, `fastvideo/tests/workflow`, and the CPU-only subset of `fastvideo/tests/train` (callbacks, utils). These test individual functions and classes. * **Component Tests**: Located in `fastvideo/tests/encoders`, `fastvideo/tests/transformers`, and `fastvideo/tests/vaes`. These verify the loading and basic functionality of model components. +* **Train Framework Tests** (GPU): Located in `fastvideo/tests/train/models`. Cover model loading + forward smoke for the new `fastvideo/train/` framework. Triggered via `/test train-framework` or as part of the Full Suite. * **SSIM Tests**: Located in `fastvideo/tests/ssim`. These are regression tests that compare generated videos against reference videos using the Structural Similarity Index Measure (SSIM) to detect quality degradation. * **Training Tests**: Located in `fastvideo/tests/training`. These validate training loops, loss calculations, and specific training techniques like LoRA, Distillation, and VSA. * **Inference Tests**: Located in `fastvideo/tests/inference`. These test specialized inference pipelines and optimizations (e.g., VSA, V-MoBA). diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 58fc63ab4..28c920617 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -214,7 +214,21 @@ def run_self_forcing_tests(): @app.function(gpu="L40S:1", image=image, timeout=900) def run_unit_test(): run_test( - "pytest ./fastvideo/tests/api/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py -vs" + "pytest ./fastvideo/tests/api/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models -vs" + ) + + +@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 run_train_framework_tests(): + run_test( + "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/train/models -vs" ) diff --git a/fastvideo/tests/train/fixtures/__init__.py b/fastvideo/tests/train/fixtures/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/tests/train/fixtures/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/tests/train/fixtures/hunyuan_t2v_min.yaml b/fastvideo/tests/train/fixtures/hunyuan_t2v_min.yaml new file mode 100644 index 000000000..74368da65 --- /dev/null +++ b/fastvideo/tests/train/fixtures/hunyuan_t2v_min.yaml @@ -0,0 +1,17 @@ +# Minimum config to instantiate HunyuanModel for the loading + forward +# smoke test. Real HunyuanVideo checkpoint (~13B) loaded via +# load_module_from_path. + +models: + student: + _target_: fastvideo.train.models.hunyuan.HunyuanModel + init_from: hunyuanvideo-community/HunyuanVideo + trainable: false + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + dit_precision: bf16 + +pipeline: {} diff --git a/fastvideo/tests/train/fixtures/wan_causal_t2v_min.yaml b/fastvideo/tests/train/fixtures/wan_causal_t2v_min.yaml new file mode 100644 index 000000000..4ac0d89cf --- /dev/null +++ b/fastvideo/tests/train/fixtures/wan_causal_t2v_min.yaml @@ -0,0 +1,17 @@ +# Minimum config to instantiate WanCausalModel for the loading + forward +# smoke test. Same checkpoint as Wan; transformer class is overridden to +# CausalWanTransformer3DModel by WanCausalModel itself. + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + +method: + _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod + +training: + dit_precision: bf16 + +pipeline: {} diff --git a/fastvideo/tests/train/fixtures/wan_t2v_min.yaml b/fastvideo/tests/train/fixtures/wan_t2v_min.yaml new file mode 100644 index 000000000..96161ac0f --- /dev/null +++ b/fastvideo/tests/train/fixtures/wan_t2v_min.yaml @@ -0,0 +1,21 @@ +# Minimum config to instantiate WanModel for the loading + forward smoke +# test. Real Wan2.1 1.3B checkpoint loaded via load_module_from_path. +# Defaults cover everything else (see fastvideo/train/utils/config.py). + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + dit_precision: bf16 + +# Empty dict (not omitted) so _parse_pipeline_config resolves the +# Wan-specific dit_config from the model checkpoint instead of falling +# back to a bare DiTConfig (which lacks Wan-specific fields like +# added_kv_proj_dim and breaks the loader). +pipeline: {} diff --git a/fastvideo/tests/train/models/__init__.py b/fastvideo/tests/train/models/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/tests/train/models/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/tests/train/models/test_load_hunyuan.py b/fastvideo/tests/train/models/test_load_hunyuan.py new file mode 100644 index 000000000..348e2113b --- /dev/null +++ b/fastvideo/tests/train/models/test_load_hunyuan.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GPU loading + forward smoke test for ``HunyuanModel``. + +Loads the real HunyuanVideo checkpoint (~13B at bf16) via +``HunyuanModel.__init__`` and runs one transformer forward pass on +synthetic inputs. Hunyuan's transformer takes a slightly different +forward signature than Wan (no ``encoder_attention_mask``, no +``return_dict``); this test mirrors the kwargs in +``HunyuanModel._build_distill_input_kwargs``. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29517") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.forward_context import set_forward_context +from fastvideo.train.models.hunyuan import HunyuanModel +from fastvideo.train.utils.config import load_run_config + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" / + "hunyuan_t2v_min.yaml") + + +@pytest.mark.usefixtures("distributed_setup") +def test_hunyuan_model_loads_and_forwards(): + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + cfg = load_run_config(_FIXTURE) + model = HunyuanModel( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=False, + ) + + transformer = model.transformer + assert isinstance(transformer, torch.nn.Module) + assert sum(p.numel() for p in transformer.parameters()) > 0 + + device = torch.device("cuda:0") + dtype = torch.bfloat16 + transformer = transformer.to(device=device, dtype=dtype).eval() + + # Hunyuan transformer takes [B, C, T, H, W] (in_channels=16, + # patch_size=2 spatial, patch_size_t=1). Small spatial + few + # frames so this fits next to the 13B model on a single L40S. + hidden_states = torch.randn(1, 16, 8, 16, 16, device=device, dtype=dtype) + # Hunyuan splits encoder_hidden_states into a leading global token + # and per-token text embeddings, so length must be >= 2. + encoder_hidden_states = torch.randn(1, 4, 4096, device=device, dtype=dtype) + timestep = torch.tensor([500], device=device, dtype=dtype) + + with torch.no_grad(), set_forward_context( + current_timestep=0, + attn_metadata=None, + ): + out = transformer( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + ) + + if isinstance(out, tuple): + out = out[0] + assert out.shape == hidden_states.shape, ( + f"output shape {tuple(out.shape)} != input shape " + f"{tuple(hidden_states.shape)}") + assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_wan.py b/fastvideo/tests/train/models/test_load_wan.py new file mode 100644 index 000000000..0099c1443 --- /dev/null +++ b/fastvideo/tests/train/models/test_load_wan.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GPU loading + forward smoke test for ``WanModel``. + +Loads the real Wan2.1 1.3B checkpoint via ``WanModel.__init__`` and +runs one transformer forward pass on synthetic inputs. Catches loader +or forward-signature regressions in +``fastvideo.train.models.wan.WanModel`` and the underlying +``WanTransformer3DModel``. +""" + +from __future__ import annotations + +import os + +# Required by the ``distributed_setup`` fixture pulled from +# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29515") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.forward_context import set_forward_context +from fastvideo.train.models.wan import WanModel +from fastvideo.train.utils.config import load_run_config + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" / "wan_t2v_min.yaml") + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan_model_loads_and_forwards(): + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + cfg = load_run_config(_FIXTURE) + model = WanModel( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=False, + ) + + transformer = model.transformer + assert isinstance(transformer, torch.nn.Module) + assert sum(p.numel() for p in transformer.parameters()) > 0 + + device = torch.device("cuda:0") + dtype = torch.bfloat16 + transformer = transformer.to(device=device, dtype=dtype).eval() + + # Tiny synthetic inputs at the smallest patch-aligned resolution + # (Wan patch_size=(1,2,2)). Keeps activations small so the test + # fits comfortably alongside the 1.3B model on a single L40S. + hidden_states = torch.randn(1, 16, 4, 32, 32, device=device, dtype=dtype) + encoder_hidden_states = torch.randn(1, + 16, + 4096, + device=device, + dtype=dtype) + encoder_attention_mask = torch.ones(1, 16, device=device, dtype=dtype) + timestep = torch.tensor([500], device=device, dtype=dtype) + + with torch.no_grad(), set_forward_context( + current_timestep=0, + attn_metadata=None, + ): + out = transformer( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + return_dict=False, + ) + + if isinstance(out, tuple): + out = out[0] + assert out.shape == hidden_states.shape, ( + f"output shape {tuple(out.shape)} != input shape " + f"{tuple(hidden_states.shape)}") + assert torch.isfinite(out).all().item(), "output contains NaN/Inf" diff --git a/fastvideo/tests/train/models/test_load_wan_causal.py b/fastvideo/tests/train/models/test_load_wan_causal.py new file mode 100644 index 000000000..6afb07076 --- /dev/null +++ b/fastvideo/tests/train/models/test_load_wan_causal.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GPU loading smoke test for ``WanCausalModel``. + +Verifies that ``WanCausalModel.__init__`` resolves the +``CausalWanTransformer3DModel`` class override and successfully loads +weights from the regular Wan2.1 1.3B checkpoint. + +A real forward pass is intentionally omitted here: the causal +transformer requires per-frame timesteps, a block-causal attention +mask, and KV cache state that ``WanCausalModel.predict_noise_streaming`` +manages for production callers. PR 5 (per-method tests) exercises that +streaming forward path end-to-end. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29516") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.train.models.wan import WanCausalModel +from fastvideo.train.utils.config import load_run_config + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" / + "wan_causal_t2v_min.yaml") + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan_causal_model_loads(): + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + cfg = load_run_config(_FIXTURE) + model = WanCausalModel( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=False, + ) + + transformer = model.transformer + assert isinstance(transformer, torch.nn.Module) + assert sum(p.numel() for p in transformer.parameters()) > 0 + + # Spot-check that the override pulled the causal class, not the + # plain Wan one. Transformer is a torch.nn.Module wrapping the + # CausalWanTransformer3DModel architecture. + assert "Causal" in type(transformer).__name__, ( + f"expected CausalWanTransformer3DModel-derived class, got " + f"{type(transformer).__name__}")