[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

@@ -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"

.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"

.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

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`) |

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).

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"
     )
 
 

fastvideo/tests/train/fixtures/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

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: {}

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: {}

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: {}

fastvideo/tests/train/models/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

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"

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"

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__}")