From fb7be2fe2cf726cf6ffee92f4f96082a77056d7e Mon Sep 17 00:00:00 2001 From: William Lin Date: Wed, 5 Aug 2026 12:47:03 -0700 Subject: [PATCH] =?UTF-8?q?[ci]:=20add=20golden-gate=20lane=20=E2=80=94=20?= =?UTF-8?q?single-layer=20bitwise=20DiT=20fingerprints=20for=20all=20SSIM-?= =?UTF-8?q?covered=20families=20(#1682)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .buildkite/pipeline.yml | 11 + .buildkite/scripts/pr_test.sh | 4 + .github/workflows/ci-slash-commands.yml | 4 +- fastvideo/tests/golden_gate/_harness.py | 220 +++++++++++++ fastvideo/tests/golden_gate/test_dreamx.py | 123 ++++++++ fastvideo/tests/golden_gate/test_flux.py | 83 +++++ .../tests/golden_gate/test_flux2_klein.py | 101 ++++++ fastvideo/tests/golden_gate/test_gamecraft.py | 85 +++++ fastvideo/tests/golden_gate/test_gen3c.py | 100 ++++++ fastvideo/tests/golden_gate/test_glm_image.py | 86 ++++++ .../tests/golden_gate/test_kandinsky5.py | 104 +++++++ fastvideo/tests/golden_gate/test_lingbot.py | 108 +++++++ fastvideo/tests/golden_gate/test_longcat.py | 78 +++++ fastvideo/tests/golden_gate/test_ltx2.py | 165 ++++++++++ .../tests/golden_gate/test_matrixgame.py | 222 +++++++++++++ .../tests/golden_gate/test_minimax_h3_t2v.py | 292 ++++++++++++++++++ fastvideo/tests/golden_gate/test_sd35.py | 78 +++++ .../tests/golden_gate/test_stable_audio.py | 73 +++++ fastvideo/tests/golden_gate/test_wan_t2v.py | 97 ++++++ fastvideo/tests/golden_gate/test_zimage.py | 87 ++++++ fastvideo/tests/modal/pr_test.py | 15 + 21 files changed, 2134 insertions(+), 2 deletions(-) create mode 100644 fastvideo/tests/golden_gate/_harness.py create mode 100644 fastvideo/tests/golden_gate/test_dreamx.py create mode 100644 fastvideo/tests/golden_gate/test_flux.py create mode 100644 fastvideo/tests/golden_gate/test_flux2_klein.py create mode 100644 fastvideo/tests/golden_gate/test_gamecraft.py create mode 100644 fastvideo/tests/golden_gate/test_gen3c.py create mode 100644 fastvideo/tests/golden_gate/test_glm_image.py create mode 100644 fastvideo/tests/golden_gate/test_kandinsky5.py create mode 100644 fastvideo/tests/golden_gate/test_lingbot.py create mode 100644 fastvideo/tests/golden_gate/test_longcat.py create mode 100644 fastvideo/tests/golden_gate/test_ltx2.py create mode 100644 fastvideo/tests/golden_gate/test_matrixgame.py create mode 100644 fastvideo/tests/golden_gate/test_minimax_h3_t2v.py create mode 100644 fastvideo/tests/golden_gate/test_sd35.py create mode 100644 fastvideo/tests/golden_gate/test_stable_audio.py create mode 100644 fastvideo/tests/golden_gate/test_wan_t2v.py create mode 100644 fastvideo/tests/golden_gate/test_zimage.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index a3e56aa47..fd7d59371 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -68,6 +68,17 @@ steps: limit: 2 agents: queue: "default" + - label: ":vertical_traffic_light: Golden-Gate Tests" + if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "golden_gate" + command: "timeout 90m .buildkite/scripts/pr_test.sh" + retry: + automatic: + - exit_status: 128 + limit: 3 + - exit_status: -1 + limit: 2 + agents: + queue: "default" - label: ":microscope: Unit Tests" if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "unit_test" command: "timeout 90m .buildkite/scripts/pr_test.sh" diff --git a/.buildkite/scripts/pr_test.sh b/.buildkite/scripts/pr_test.sh index fd46e1d4d..b97b3757c 100755 --- a/.buildkite/scripts/pr_test.sh +++ b/.buildkite/scripts/pr_test.sh @@ -187,6 +187,10 @@ case "$TEST_TYPE" in log "Running transformer tests..." MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests" ;; + "golden_gate") + log "Running golden-gate tests..." + MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_golden_gate_tests" + ;; "ssim") log "Running SSIM tests..." SSIM_BOOTSTRAP_ARGS=$(ssim_bootstrap_args) diff --git a/.github/workflows/ci-slash-commands.yml b/.github/workflows/ci-slash-commands.yml index e8de276a1..06600df9f 100644 --- a/.github/workflows/ci-slash-commands.yml +++ b/.github/workflows/ci-slash-commands.yml @@ -129,7 +129,7 @@ jobs: set -euo pipefail TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true) - VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit" + VALID="encoder vae transformer kernel unit dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit" if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then echo "Unknown test: '$TEST_NAME'. Valid: $VALID" exit 1 @@ -138,7 +138,7 @@ jobs: declare -A MAP=( [encoder]=encoder [vae]=vae [transformer]=transformer [kernel]=kernel_tests [unit]=unit_test [dreamverse]=dreamverse_app - [ssim]=ssim [training]=training + [ssim]=ssim [golden-gate]=golden_gate [training]=training [lora-inference]=inference_lora [lora-training]=training_lora [lora-extraction]=lora_extraction [distillation]=distillation_dmd [self-forcing]=self_forcing diff --git a/fastvideo/tests/golden_gate/_harness.py b/fastvideo/tests/golden_gate/_harness.py new file mode 100644 index 000000000..deb500024 --- /dev/null +++ b/fastvideo/tests/golden_gate/_harness.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared machinery for golden-gate tests: single-layer bitwise DiT fingerprints. + +A gate loads ONE transformer block of a DiT (selective safetensors read via the +checkpoint index when present), feeds a fixed seeded batch, and compares the +output bitwise against a device-keyed golden latent. Zero tolerance is the +point: on one device these blocks reproduce bit-identically, so any drift is a +real change to the compute path — and the golden's environment fingerprint +turns legitimate env changes into named failures instead of mystery drift. + +Per-model tests provide a :class:`GateSpec`; everything else lives here. +""" + +from __future__ import annotations + +import json +import os +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable + +import pytest +import torch +from torch.testing import assert_close + +GOLDEN_REPO_ID = os.environ.get("FASTVIDEO_SSIM_REFERENCE_HF_REPO", "FastVideo/ssim-reference-videos") +GOLDEN_ROOT = Path(os.environ.get("FASTVIDEO_GOLDEN_GATE_DIR", Path(__file__).resolve().parent / "goldens")) +DEFAULT_SEED = 20260805 + +# Must be set before cuBLAS initializes for use_deterministic_algorithms. +os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") +os.environ.setdefault("MASTER_ADDR", "127.0.0.1") +os.environ.setdefault("MASTER_PORT", "29661") +os.environ.setdefault("RANK", "0") +os.environ.setdefault("WORLD_SIZE", "1") +os.environ.setdefault("LOCAL_RANK", "0") + + +@dataclass +class GateSpec: + """Everything model-specific about one golden gate.""" + name: str # golden filename stem, e.g. "wan_t2v" + repo_id: str # HF checkpoint repo + build_block: Callable[[int], torch.nn.Module] + make_inputs: Callable[[torch.device, int], dict[str, Any]] + subfolder: str = "transformer" + layer: int = 0 + prefix_template: str = "blocks.{N}." + renames: tuple[tuple[str, str], ...] = () # (regex, replacement) on ".{stripped}" names + attention_backend: str = "FLASH_ATTN" + seed: int = DEFAULT_SEED + # compute dtype for the block; matrixgame2's int64 null-modulation promotes + # activations to fp32 mid-block, so that DiT runs fp32 in production + dtype: torch.dtype = torch.bfloat16 + # some blocks return tuples (e.g. double-stream img/txt) — reduce to one tensor + postprocess: Callable[[Any], torch.Tensor] = field(default=lambda out: out) + model_root_env: str | None = None # env var naming a local checkpoint root + weight_file: str = "diffusion_pytorch_model.safetensors" # single-file fallback + + +def env_fingerprint(spec: GateSpec) -> dict[str, str]: + try: + import flash_attn + fa_version = str(getattr(flash_attn, "__version__", "?")) + except ImportError: + fa_version = "" + return { + "torch": str(torch.__version__), + "cuda": str(torch.version.cuda), + "cudnn": str(torch.backends.cudnn.version()), + "flash_attn": fa_version, + "device_name": torch.cuda.get_device_name(0), + "attention_backend": spec.attention_backend, + "fastvideo_fa4": os.environ.get("FASTVIDEO_FA4", ""), + "tf32_matmul": str(torch.backends.cuda.matmul.allow_tf32), + "tf32_cudnn": str(torch.backends.cudnn.allow_tf32), + } + + +def _component_dir(spec: GateSpec) -> Path: + """Local checkout wins; otherwise fetch the index + only the shards holding + this layer (or the single weight file for index-less checkpoints).""" + if spec.model_root_env: + root = os.environ.get(spec.model_root_env) + if root: + return Path(root) / spec.subfolder + + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + + index_name = f"{spec.subfolder}/{spec.weight_file}.index.json" + prefix = spec.prefix_template.format(N=spec.layer) + try: + index_path = Path(hf_hub_download(spec.repo_id, index_name)) + except EntryNotFoundError: + single = Path(hf_hub_download(spec.repo_id, f"{spec.subfolder}/{spec.weight_file}")) + return single.parent + weight_map = json.loads(index_path.read_text())["weight_map"] + shards = sorted({shard for name, shard in weight_map.items() if name.startswith(prefix)}) + if not shards: + pytest.fail(f"no shards found for prefix {prefix} in {spec.repo_id}", pytrace=False) + for shard in shards: + hf_hub_download(spec.repo_id, f"{spec.subfolder}/{shard}") + return index_path.parent + + +def load_layer_state(spec: GateSpec, component_dir: Path) -> dict[str, torch.Tensor]: + from safetensors import safe_open + + prefix = spec.prefix_template.format(N=spec.layer) + renames = [(re.compile(pat), repl) for pat, repl in spec.renames] + index_file = component_dir / f"{spec.weight_file}.index.json" + + by_shard: dict[str, list[str]] = {} + if index_file.is_file(): + for name, shard in json.loads(index_file.read_text())["weight_map"].items(): + if name.startswith(prefix): + by_shard.setdefault(shard, []).append(name) + else: + with safe_open(component_dir / spec.weight_file, framework="pt", device="cpu") as f: + by_shard[spec.weight_file] = [n for n in f.keys() if n.startswith(prefix)] + assert by_shard and any(by_shard.values()), f"no parameters found for prefix {prefix}" + + state: dict[str, torch.Tensor] = {} + for shard, names in sorted(by_shard.items()): + with safe_open(component_dir / shard, framework="pt", device="cpu") as f: + for name in names: + mapped = f".{name[len(prefix):]}" + for pattern, repl in renames: + mapped = pattern.sub(repl, mapped) + state[mapped[1:]] = f.get_tensor(name) + return state + + +def _golden_relpath(spec: GateSpec) -> str: + device_slug = re.sub(r"[^A-Za-z0-9]+", "_", torch.cuda.get_device_name(0)).strip("_") + return f"{device_slug}/{spec.name}_layer{spec.layer}_{spec.attention_backend}_seed{spec.seed}.pt" + + +def _resolve_golden(spec: GateSpec) -> tuple[Path, bool]: + local = GOLDEN_ROOT / _golden_relpath(spec) + if local.exists(): + return local, True + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + try: + remote = hf_hub_download(GOLDEN_REPO_ID, f"golden_gates/{_golden_relpath(spec)}", repo_type="dataset") + return Path(remote), True + except EntryNotFoundError: + return local, False + + +def run_gate(spec: GateSpec) -> None: + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("golden-gate requires a bf16-capable CUDA GPU") + + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = spec.attention_backend + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.benchmark = False + + device = torch.device("cuda:0") + component_dir = _component_dir(spec) + + from fastvideo.forward_context import set_forward_context + + block = spec.build_block(spec.layer) + state = load_layer_state(spec, component_dir) + missing, unexpected = block.load_state_dict(state, strict=True) + assert not missing and not unexpected, ( + f"state mismatch for {spec.name}: missing={missing[:6]} unexpected={unexpected[:6]}") + block = block.to(device=device, dtype=spec.dtype).eval() + + inputs = spec.make_inputs(device, spec.seed) + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + raw = block(**inputs) + out = spec.postprocess(raw) + out = out.detach().float().cpu() + + golden_path, golden_exists = _resolve_golden(spec) + meta = {"name": spec.name, "layer": spec.layer, "seed": spec.seed, + "shape": list(out.shape), "env": env_fingerprint(spec)} + if not golden_exists: + golden_path.parent.mkdir(parents=True, exist_ok=True) + torch.save({"output": out, "metadata": meta}, golden_path) + pytest.fail( + f"golden latent seeded at {golden_path} ({meta}). Verify with a second run, " + f"then upload to {GOLDEN_REPO_ID} at golden_gates/{_golden_relpath(spec)}", + pytrace=False, + ) + + golden = torch.load(golden_path, weights_only=True) + mismatches = { + key: (value, meta["env"].get(key)) + for key, value in golden["metadata"]["env"].items() + if meta["env"].get(key) != value + } + assert not mismatches, ( + f"environment changed since golden was minted — re-mint deliberately rather " + f"than chasing bit drift. golden -> current: {mismatches}") + + drift = (out - golden["output"]).abs() + print(f"golden-gate[{spec.name}] drift: max_abs={drift.max().item():.10f} " + f"mean_abs={drift.mean().item():.10f}", flush=True) + assert_close(out, golden["output"], atol=0.0, rtol=0.0) + + +@pytest.fixture(scope="module") +def distributed_runtime(): + """Single-process distributed init — required by DistributedAttention.""" + if not torch.cuda.is_available(): + yield + return + from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel, + ) + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + cleanup_dist_env_and_memory() diff --git a/fastvideo/tests/golden_gate/test_dreamx.py b/fastvideo/tests/golden_gate/test_dreamx.py new file mode 100644 index 000000000..21238aa91 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_dreamx.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: DreamX-World-5B-Cam transformer block 0 (Wan2.2-TI2V-5B derivative +with PRoPE camera control, bidirectional diffusion variant). + +Spec notes: DreamXWorldArchConfig class defaults inherit Wan-14B dims and do NOT +match the 5B checkpoint — geometry comes from make_dreamx_world_5b_cam_dit_config() +(configs/pipelines/dreamx_world.py:19-34), verified equal to the mirror repo's +transformer/config.json. Checkpoint index is model.safetensors.index.json (saved +via save_torch_state_dict), hence weight_file="model.safetensors". temb is the 4D +[B, seq, 6, inner_dim] per-token shape (expand_timesteps=True pipeline path, +dreamx_world.py:273-282). y_camera dtype must equal hidden_states dtype (mixed +dtype einsum in _dreamx_prope_qkv errors); the camera builder is deterministic +(camera_conditioning.py:201-228). The ssim test runs TORCH_SDPA only. Forward +returns a plain [1, seq, 3072] tensor — no postprocess. The AR variant +(DreamX-World-5B, CausalWanAttentionBlock) is a different implementation and +needs its own gate recipe; not covered by this spec. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +INNER_DIM = 3072 # 24 heads x 128, DreamX-World-5B-Cam transformer/config.json +TEXT_LEN = 512 + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.pipelines.dreamx_world import ( + make_dreamx_world_5b_cam_dit_config) + from fastvideo.models.dits.dreamx_world import DreamXWorldTransformerBlock + + arch = make_dreamx_world_5b_cam_dit_config().arch_config + # exactly as DreamXWorldTransformer3DModel.__init__ builds each block + # (dreamx_world.py:374-391) + return DreamXWorldTransformerBlock( + arch.num_attention_heads * arch.attention_head_dim, # 3072 + arch.ffn_dim, # 14336 + arch.num_attention_heads, # 24 + arch.qk_norm, # "rms_norm_across_heads" + arch.cross_attn_norm, # True + arch.eps, # 1e-6 + arch.added_kv_proj_dim, # None -> T2V cross-attn + arch._supported_attention_backends, + quant_config=None, + prefix=f"Wan.blocks.{layer}", + add_control_adapter=arch.add_control_adapter, # True + cam_method=arch.cam_method, # "prope" + attn_compress=arch.attn_compress, # 1 + cam_self_attn_layers=arch.cam_self_attn_layers, # None -> cam attn everywhere + layer_idx=layer, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + """Fixed batch mirroring the SSIM defaults (9 frames @ 64x64): + latent (3, 4, 4) -> patch (1,2,2) -> seq = 3*2*2 = 12, cameras = 3 latent frames.""" + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import ( + build_dreamx_camera_condition) + + generator = torch.Generator(device="cpu").manual_seed(seed) + num_frames, height, width = 9, 64, 64 + post_t = (num_frames - 1) // 4 + 1 # VAE 4x temporal + post_h, post_w = height // 16 // 2, width // 16 // 2 # VAE 16x spatial, patch (1,2,2) + seq = post_t * post_h * post_w # 12 + + # RoPE exactly as the model forward builds it (dreamx_world.py:429-439) + d = INNER_DIM // 24 # head_dim 128 + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] # [44, 42, 42] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_t, post_h, post_w), INNER_DIM, 24, rope_dim_list, + rope_theta=10000, dtype=torch.float64) + freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) + + hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, TEXT_LEN, INNER_DIM, generator=generator, dtype=torch.float32) + # expand_timesteps=True pipeline -> per-token timesteps -> 4D temb (B, seq, 6, dim) + temb = torch.randn(1, seq, 6, INNER_DIM, generator=generator, dtype=torch.float32) + + # Deterministic camera condition, same builder + keys as the pipeline stage + # (stages.py:55-64). dtype MUST match hidden_states. + camera = build_dreamx_camera_condition( + ["w"], [4.0], num_frames=num_frames, height=height, width=width, + dtype=torch.bfloat16, device=device) + y_camera = {k: v.unsqueeze(0) for k, v in camera.items()} # viewmats (1,3,4,4), K (1,3,3,3) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "freqs_cis": freqs_cis, + "original_seq_len": seq, + "y_camera": y_camera, + } + + +SPEC = GateSpec( + name="dreamx_world_5b_cam", + repo_id="FastVideo/DreamX-World-5B-Cam-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="blocks.{N}.", + renames=( + # to_out rename must run before the generic attn1 strip + (r"\.attn1\.to_out\.0\.", ".to_out."), + (r"\.attn1\.", "."), # to_q/to_k/to_v/norm_q/norm_k -> block level + (r"\.attn2\.to_out\.0\.", ".attn2.to_out."), + (r"\.ffn\.net\.0\.proj\.", ".ffn.fc_in."), + (r"\.ffn\.net\.2\.", ".ffn.fc_out."), + (r"\.norm2\.", ".self_attn_residual_norm.norm."), + # identity: attn2.{to_q,to_k,to_v,norm_q,norm_k}, cam_self_attn.*, scale_shift_table + ), + attention_backend="TORCH_SDPA", + weight_file="model.safetensors", +) + + +def test_dreamx_world_cam_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_flux.py b/fastvideo/tests/golden_gate/test_flux.py new file mode 100644 index 000000000..99449b3e7 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_flux.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: FLUX.1-dev double-stream transformer block 0. + +Spec notes: dim = 24 heads x 128 = 3072 (configs/models/dits/flux.py:16-17). +Checkpoint names load 1:1 (param_names_mapping is empty, flux.py:391-393) so +no renames. RoPE is built outside the block by FluxPosEmbed(theta=10000, +axes_dim=[16, 56, 56]) with text ids FIRST (flux.py:414-415, 536-537); cos/sin +are float64 [S, 128]. temb is [B, 3072] (SD3AdaLayerNormZero 6-way chunk, +sd3.py:394-395). Forward returns the TUPLE (encoder_hidden_states, +hidden_states) (flux.py:329) — postprocess cats both streams along seq into +one golden. The flux SSIM test pins TORCH_SDPA. black-forest-labs/FLUX.1-dev +is a GATED HF repo; FLUX_T2I_MODEL_DIR can point at a local checkout. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +DIM = 3072 # 24 heads x 128, FluxTransformer2DArchConfig defaults + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.flux import FluxTransformer2DArchConfig + from fastvideo.models.dits.flux import FluxTransformer2DModel, FluxTransformerBlock + + arch = FluxTransformer2DArchConfig() + return FluxTransformerBlock( + dim=arch.num_attention_heads * arch.attention_head_dim, # 3072 + num_attention_heads=arch.num_attention_heads, # 24 + attention_head_dim=arch.attention_head_dim, # 128 + supported_attention_backends=FluxTransformer2DModel._supported_attention_backends, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.flux import FluxPosEmbed + + generator = torch.Generator(device="cpu").manual_seed(seed) + n_txt, h, w = 32, 16, 16 + n_img = h * w + + # Diffusers-style ids: txt ids all zeros; img ids carry (0, row, col). + txt_ids = torch.zeros(n_txt, 3, dtype=torch.float64) + img_ids = torch.zeros(n_img, 3, dtype=torch.float64) + img_ids[:, 1] = torch.arange(h, dtype=torch.float64).repeat_interleave(w) + img_ids[:, 2] = torch.arange(w, dtype=torch.float64).repeat(h) + + rope = FluxPosEmbed(theta=10000, axes_dim=[16, 56, 56]) + cos, sin = rope(torch.cat([txt_ids, img_ids], dim=0)) # text ids first + + hidden = torch.randn(1, n_img, DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, n_txt, DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(1, DIM, generator=generator, dtype=torch.float32) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "image_rotary_emb": (cos.to(device), sin.to(device)), + # joint_attention_kwargs is accepted but immediately deleted — omit. + } + + +SPEC = GateSpec( + name="flux1_dev", + repo_id="black-forest-labs/FLUX.1-dev", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer_blocks.{N}.", + renames=(), # HF weight names already match this module layout (flux.py:391-393) + attention_backend="TORCH_SDPA", + # forward returns (encoder_hidden_states, hidden_states) — both [1, S, 3072] + postprocess=lambda out: torch.cat(out, dim=1), + model_root_env="FLUX_T2I_MODEL_DIR", +) + + +def test_flux_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_flux2_klein.py b/fastvideo/tests/golden_gate/test_flux2_klein.py new file mode 100644 index 000000000..01f70ef56 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_flux2_klein.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: FLUX.2-klein-4B double-stream transformer block 0. + +Spec notes: Klein-4B block dims equal the Flux2ArchConfig defaults (24 heads x +128 head_dim -> 3072, mlp_ratio 3.0, eps 1e-6, bias=False); the model-level +config diffs (num_layers=5, joint_attention_dim=7680, ...) never enter the +block ctor. Checkpoint is SINGLE-FILE (transformer/diffusion_pytorch_model. +safetensors, no index.json) and its keys match FastVideo names exactly, so no +renames. RoPE lives outside the block and is built exactly as +Flux2Transformer2DModel.forward does (flux_2.py:1043-1054); cos/sin stay fp32 +— the block never casts them to bf16. Modulation params arrive as two +(shift, scale, gate) 3-tuples per stream, each tensor [1, 1, 3072] bf16. +Block construction needs the distributed fixture (Flux2FeedForward reads +get_tp_world_size() in __init__). forward returns (encoder_hidden_states, +hidden_states) — postprocess concats them on dim=1. Backend TORCH_SDPA per +the flux2 ssim lane (test_flux2_similarity.py:86). +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +DIM = 3072 # 24 heads x 128, Flux2ArchConfig defaults == Klein-4B block dims +TEXT_LEN = 32 +IMG_H, IMG_W = 8, 16 # 128 image tokens + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.flux_2 import Flux2ArchConfig + from fastvideo.models.dits.flux_2 import Flux2TransformerBlock + + arch = Flux2ArchConfig() + return Flux2TransformerBlock( + dim=arch.num_attention_heads * arch.attention_head_dim, + num_attention_heads=arch.num_attention_heads, + attention_head_dim=arch.attention_head_dim, + mlp_ratio=arch.mlp_ratio, + eps=arch.eps, + bias=False, + supported_attention_backends=arch._supported_attention_backends, + ff_context_swiglu_fp32=arch.ff_context_swiglu_fp32, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.flux_2 import ( + Flux2PosEmbed, + compute_flux2_freqs_cis_from_ids, + ) + + generator = torch.Generator(device="cpu").manual_seed(seed) + n_img = IMG_H * IMG_W + + hidden = torch.randn(1, n_img, DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, TEXT_LEN, DIM, generator=generator, dtype=torch.float32) + + def _mod_sets(n_sets: int): + # (shift, scale, gate) per set, [1, 1, DIM] each, in the DiT dtype. + return tuple( + tuple( + torch.randn(1, 1, DIM, generator=generator, dtype=torch.float32) + .to(device=device, dtype=torch.bfloat16) + for _ in range(3)) + for _ in range(n_sets)) + + # ids built as in Flux2Transformer2DModel.forward (flux_2.py:1044-1051) + txt_ids = torch.cartesian_prod( + torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(TEXT_LEN)) + img_ids = torch.cartesian_prod( + torch.arange(1), torch.arange(IMG_H), torch.arange(IMG_W), torch.arange(1)) + rope = Flux2PosEmbed(theta=2000, axes_dim=[32, 32, 32, 32]) + freqs_cis = compute_flux2_freqs_cis_from_ids(rope, txt_ids, img_ids, device=device) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb_mod_params_img": _mod_sets(2), + "temb_mod_params_txt": _mod_sets(2), + "freqs_cis": freqs_cis, # fp32 (cos, sin) [S_txt + S_img, 128] — keep fp32 + } + + +SPEC = GateSpec( + name="flux2_klein_4b", + repo_id="black-forest-labs/FLUX.2-klein-4B", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer_blocks.{N}.", + renames=(), + attention_backend="TORCH_SDPA", + # (txt, img) tuple -> one tensor; both are [1, S, DIM] so cat on dim=1. + postprocess=lambda out: torch.cat(out, dim=1), +) + + +def test_flux2_klein_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_gamecraft.py b/fastvideo/tests/golden_gate/test_gamecraft.py new file mode 100644 index 000000000..2bf8a42fc --- /dev/null +++ b/fastvideo/tests/golden_gate/test_gamecraft.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: HunyuanGameCraft double-stream block 0 (Hunyuan-video arch). + +Spec notes: the repeated block is MMDoubleStreamBlock from hunyuanvideo.py +(imported by hunyuangamecraft.py:24-28); ctor mirrors +HunyuanGameCraftTransformer3DModel.__init__ (hunyuangamecraft.py:252-261) with +hidden_size 3072 = 24 heads x 128. Checkpoint is a SINGLE FILE (no index.json) +under transformer/ and its keys are already in FastVideo form, so renames is +empty. RoPE is built over the video grid only, rope_dim_list (16,56,56), +rope_theta 256 (hunyuangamecraft.py:340-348) — applied inside +DistributedAttention to the img qkv, so freqs_cis length must equal the img +sequence length. Forward returns an (img, txt) TUPLE (hunyuanvideo.py:280); +postprocess concatenates along the sequence dim into one tensor. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +HIDDEN = 3072 # 24 heads x 128, configs/models/dits/hunyuangamecraft.py:125 + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftArchConfig + from fastvideo.models.dits.hunyuanvideo import MMDoubleStreamBlock + + arch = HunyuanGameCraftArchConfig() + return MMDoubleStreamBlock( + hidden_size=arch.hidden_size, # 3072 + num_attention_heads=arch.num_attention_heads, # 24 + mlp_ratio=arch.mlp_ratio, # 4.0 + dtype=arch.dtype, # None + supported_attention_backends=arch._supported_attention_backends, + prefix=f"hunyuan_gamecraft.double_blocks.{layer}", + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + """3x4x5 = 60 video tokens + 32 text tokens, hidden 3072. + + Shapes per MMDoubleStreamBlock.forward (hunyuanvideo.py:193-200): + img [B, S_img, 3072], txt [B, S_txt, 3072], vec [B, 3072], + freqs_cis = (cos, sin) each [S_img, 128]. + """ + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + + generator = torch.Generator(device="cpu").manual_seed(seed) + tt, th, tw = 3, 4, 5 + img_seq, txt_seq = tt * th * tw, 32 + + freqs_cos, freqs_sin = get_rotary_pos_embed( + (tt, th, tw), HIDDEN, 24, [16, 56, 56], 256) # -> [60, 128] fp32 each + + img = torch.randn(1, img_seq, HIDDEN, generator=generator, dtype=torch.float32) + txt = torch.randn(1, txt_seq, HIDDEN, generator=generator, dtype=torch.float32) + vec = torch.randn(1, HIDDEN, generator=generator, dtype=torch.float32) + + return { + "img": img.to(device=device, dtype=torch.bfloat16), + "txt": txt.to(device=device, dtype=torch.bfloat16), + "vec": vec.to(device=device, dtype=torch.bfloat16), + "freqs_cis": (freqs_cos.to(device=device, dtype=torch.bfloat16), + freqs_sin.to(device=device, dtype=torch.bfloat16)), + } + + +SPEC = GateSpec( + name="gamecraft", + repo_id="FastVideo/HunyuanGameCraft-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="double_blocks.{N}.", + renames=(), # checkpoint keys already match the block's state_dict 1:1 + attention_backend="FLASH_ATTN", + postprocess=lambda out: torch.cat(out, dim=1), # (img, txt) tuple -> one tensor + model_root_env="GAMECRAFT_MODEL_PATH", +) + + +def test_gamecraft_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_gen3c.py b/fastvideo/tests/golden_gate/test_gen3c.py new file mode 100644 index 000000000..498426f2b --- /dev/null +++ b/fastvideo/tests/golden_gate/test_gen3c.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: GEN3C-Cosmos-7B transformer block 0 (Cosmos-family DiT). + +Spec notes: block ctor takes no prefix arg (gen3c.py:393-403); forward returns +a plain tensor (1, 240, 4096), not a tuple (gen3c.py:547). RoPE is built +outside the block by the model with hidden_size=attention_head_dim (128), and +cos/sin stay fp32 while hidden is bf16 (gen3c.py:800-806, 620-626). The +FastVideo mirror checkpoint is a SINGLE file transformer/model.safetensors +with no index.json and already stores FastVideo-native names, so renames are +empty — the harness only strips the transformer_blocks.{N}. prefix. The gen3c +SSIM test parametrizes only TORCH_SDPA (test_gen3c_similarity.py:109,115). +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +HIDDEN_SIZE = 4096 # 32 heads x 128, Gen3CArchConfig defaults + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.gen3c import Gen3CArchConfig + from fastvideo.models.dits.gen3c import Gen3CTransformerBlock + + # Mirrors Gen3CTransformer3DModel.__init__ (gen3c.py:840-852). + arch = Gen3CArchConfig() + return Gen3CTransformerBlock( + num_attention_heads=arch.num_attention_heads, # 32 + attention_head_dim=arch.attention_head_dim, # 128 -> hidden 4096 + cross_attention_dim=arch.text_embed_dim, # 1024 + mlp_ratio=arch.mlp_ratio, # 4.0 -> mlp inner 16384 + adaln_lora_dim=arch.adaln_lora_dim, # 256 + use_adaln_lora=arch.use_adaln_lora, # True + qk_norm=(arch.qk_norm == "rms_norm"), # True + supported_attention_backends=arch._supported_attention_backends, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + """Fixed seeded batch: 3x8x10 post-patch latent grid (S=240), 128 text tokens. + + Shapes per Gen3CTransformerBlock.forward (gen3c.py:461-480): + hidden_states (B,S,4096), encoder_hidden_states (B,N,1024), + affine_emb (B,4096), adaln_lora (B,12288), rope_emb ((S,128),(S,128)) fp32, + extra_pos_emb (B,S,4096), original_seq_len int. + """ + from fastvideo.models.dits.gen3c import Gen3CRotaryPosEmbed + + generator = torch.Generator(device="cpu").manual_seed(seed) + t, h, w = 3, 8, 10 + seq = t * h * w # 240 + + # RoPE exactly as Gen3CTransformer3DModel.__init__ builds it (gen3c.py:800-806): + # hidden_size = attention_head_dim (128), NOT the model hidden size. + rope = Gen3CRotaryPosEmbed( + hidden_size=128, + max_size=(128, 240, 240), + patch_size=(1, 2, 2), + rope_scale=(2.0, 1.0, 1.0), + enable_fps_modulation=True, + ) + # rope.forward() reads only shape/device of its input; fps=24 matches the + # SSIM test (test_gen3c_similarity.py:83). + cos, sin = rope(torch.empty(1, t, h, w, 1, device=device), fps=24) # each (240, 128) fp32 + + hidden = torch.randn(1, seq, HIDDEN_SIZE, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, 128, 1024, generator=generator, dtype=torch.float32) + affine = torch.randn(1, HIDDEN_SIZE, generator=generator, dtype=torch.float32) + adaln_lora = torch.randn(1, 3 * HIDDEN_SIZE, generator=generator, dtype=torch.float32) + extra_pos = torch.randn(1, seq, HIDDEN_SIZE, generator=generator, dtype=torch.float32) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "affine_emb": affine.to(device=device, dtype=torch.bfloat16), + "adaln_lora": adaln_lora.to(device=device, dtype=torch.bfloat16), + "rope_emb": (cos, sin), # keep fp32 — the full model passes them fp32 + "extra_pos_emb": extra_pos.to(device=device, dtype=torch.bfloat16), + "original_seq_len": seq, + } + + +SPEC = GateSpec( + name="gen3c_cosmos_7b", + repo_id="FastVideo/GEN3C-Cosmos-7B-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer_blocks.{N}.", + renames=(), # mirror checkpoint already uses FastVideo-native names + attention_backend="TORCH_SDPA", + weight_file="model.safetensors", # single file, no index.json +) + + +def test_gen3c_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_glm_image.py b/fastvideo/tests/golden_gate/test_glm_image.py new file mode 100644 index 000000000..92ad32032 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_glm_image.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: GLM-Image (GlmImageTransformer2DModel) transformer block 0. + +Spec notes: inner_dim = 32 heads x 128 = 4096 (glm_image.py:631, arch config +defaults). RoPE is built outside the block by GlmImageRotaryPosEmbed +(glm_image.py:727-729) from only the shape/device of the pre-patchify latent +(B, 16, H, W); cos/sin are fp32 (image_seq, 64). temb is passed post-SiLU in +the hidden dtype (glm_image.py:747-750), so bf16 here. The block forward +returns a TUPLE (hidden_states, encoder_hidden_states) (glm_image.py:502) — +postprocess cats encoder first along seq, matching the attention layout +(glm_image.py:343). attn1.to_out is nn.ModuleList so "attn1.to_out.0.*" +matches natively; only the two ff.net renames apply. The SSIM test pins +TORCH_SDPA and loads local weights via GLM_IMAGE_MODEL_DIR. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +INNER_DIM = 4096 # 32 heads x 128, GlmImageDiTArchConfig defaults +TEXT_LEN = 32 +LATENT_H = LATENT_W = 64 # patch_size=2 -> 32x32 = 1024 image tokens + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.glm_image import GlmImageDiTArchConfig + from fastvideo.models.dits.glm_image import (GlmImageTransformer2DModel, + GlmImageTransformerBlock) + + arch = GlmImageDiTArchConfig() + return GlmImageTransformerBlock( + arch.num_attention_heads * arch.attention_head_dim, # dim = 4096 + arch.num_attention_heads, # 32 + arch.attention_head_dim, # 128 + arch.time_embed_dim, # 512 + supported_attention_backends=GlmImageTransformer2DModel._supported_attention_backends, + prefix=f"transformer_blocks.{layer}", + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.glm_image import GlmImageRotaryPosEmbed + + generator = torch.Generator(device="cpu").manual_seed(seed) + n_image = (LATENT_H // 2) * (LATENT_W // 2) # 1024 tokens + + # rope consumes only shape+device of the raw latent (glm_image.py:518-522) + rope = GlmImageRotaryPosEmbed(dim=128, patch_size=2, theta=10000.0) + cos, sin = rope(torch.empty(1, 16, LATENT_H, LATENT_W, device=device)) + + hidden = torch.randn(1, n_image, INNER_DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, TEXT_LEN, INNER_DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(1, 512, generator=generator, dtype=torch.float32) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "image_rotary_emb": (cos, sin), # fp32, already on device + } + + +SPEC = GateSpec( + name="glm_image", + repo_id="zai-org/GLM-Image", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer_blocks.{N}.", + renames=( + (r"\.ff\.net\.0\.proj\.", ".ff.fc_in."), + (r"\.ff\.net\.2\.", ".ff.fc_out."), + ), + attention_backend="TORCH_SDPA", + # forward returns (hidden_states, encoder_hidden_states); encoder first, + # like the attention concat layout (glm_image.py:343) + postprocess=lambda out: torch.cat([out[1], out[0]], dim=1), + model_root_env="GLM_IMAGE_MODEL_DIR", +) + + +def test_glm_image_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_kandinsky5.py b/fastvideo/tests/golden_gate/test_kandinsky5.py new file mode 100644 index 000000000..5f68c27d7 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_kandinsky5.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: Kandinsky-5.0-T2V-Lite visual decoder block 0. + +Spec notes: Kandinsky5ArchConfig defaults (model_dim=2048, ff_dim=5120) do NOT +match the Lite checkpoint — production overrides them from +transformer/config.json via update_model_arch, so the geometry is pinned here +explicitly from that config.json (model_dim=1792, ff_dim=7168, time_dim=512, +axes_dims=(16, 24, 24), attention_type="regular" -> use_nabla=False). RoPE is +built outside the block exactly as Kandinsky5Transformer3DModel.forward does +(kandinsky5.py:762-770): Kandinsky5RoPE3D returns real 2x2 rotation matrices +[B, T, H, W, 1, head_dim/2, 2, 2] fp32, flattened over (T, H, W) for the +dense path. time_embed reaches the block as fp32 (modulation computes in fp32 +under autocast); sparse_params is a required positional -> None for the dense +"regular" path. Single-file checkpoint (no index.json) — the harness's +index-less fallback handles it. The ssim test pins FASTVIDEO_FA4=0 alongside +FLASH_ATTN; mirrored here before any fastvideo import. +""" + +from __future__ import annotations + +import os + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +os.environ.setdefault("FASTVIDEO_FA4", "0") # test_kandinsky5_similarity.py:84 + +MODEL_DIM = 1792 # Kandinsky-5.0-T2V-Lite transformer/config.json +TIME_DIM = 512 +FF_DIM = 7168 +AXES_DIMS = (16, 24, 24) # head_dim 64, 28 heads + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5ArchConfig + from fastvideo.models.dits.kandinsky5 import Kandinsky5TransformerDecoderBlock + + arch = Kandinsky5ArchConfig( + model_dim=MODEL_DIM, + ff_dim=FF_DIM, + time_dim=TIME_DIM, + axes_dims=AXES_DIMS, + attention_type="regular", + ) + return Kandinsky5TransformerDecoderBlock( + arch.model_dim, + arch.time_dim, + arch.ff_dim, + sum(arch.axes_dims), + arch._supported_attention_backends, + prefix=f"Kandinsky5.visual_transformer_blocks.{layer}", + use_nabla=arch.attention_type == "nabla", + quant_config=None, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.kandinsky5 import Kandinsky5RoPE3D + + generator = torch.Generator(device="cpu").manual_seed(seed) + b, t, h, w = 1, 2, 4, 4 # post-patchify grid (parity test uses the same) + seq = t * h * w # 32 + n_text = 8 + + visual = torch.randn(b, seq, MODEL_DIM, generator=generator, dtype=torch.float32) + text = torch.randn(b, n_text, MODEL_DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(b, TIME_DIM, generator=generator, dtype=torch.float32) + + # Built exactly like model.forward (kandinsky5.py:762-770): rotation + # matrices [B, T, H, W, 1, head_dim/2, 2, 2] fp32, then the dense + # fractal_flatten(block_mask=False) is rope.flatten(1, 3) (line 109). + rope3d = Kandinsky5RoPE3D(AXES_DIMS) + rope = rope3d((b, t, h, w), + [torch.arange(t), torch.arange(h), torch.arange(w)]) + rope = rope.flatten(1, 3) + + return { + "visual_embed": visual.to(device=device, dtype=torch.bfloat16), + "text_embed": text.to(device=device, dtype=torch.bfloat16), + "time_embed": temb.to(device=device), # stays fp32 (kandinsky5.py:745-746) + "rope": rope.to(device=device), # fp32 2x2 rotation matrices + "sparse_params": None, # required positional; dense "regular" path + } + + +SPEC = GateSpec( + name="kandinsky5_t2v_lite", + repo_id="kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="visual_transformer_blocks.{N}.", + renames=( + (r"\.feed_forward\.in_layer\.", ".feed_forward.mlp.fc_in."), + (r"\.feed_forward\.out_layer\.", ".feed_forward.mlp.fc_out."), + ), + attention_backend="FLASH_ATTN", +) + + +def test_kandinsky5_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_lingbot.py b/fastvideo/tests/golden_gate/test_lingbot.py new file mode 100644 index 000000000..6d5cda55b --- /dev/null +++ b/fastvideo/tests/golden_gate/test_lingbot.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: LingBot-World camera-conditioned I2V DiT, transformer block 0. + +Spec notes: arch-config defaults (lingbotworld.py:57-71) match the checkpoint's +transformer/config.json exactly (heads=40, head_dim=128, ffn_dim=13824, +qk_norm="rms_norm_across_heads", cross_attn_norm=True, eps=1e-6, +added_kv_proj_dim=None -> WanT2VCrossAttention). Two-expert MoE checkpoint +(transformer + transformer_2, identical configs); we gate "transformer" (the +high-noise expert). Checkpoint tensors are fp32; casting the block to bf16 +mirrors serving dtype. RoPE is built outside exactly as +LingBotWorldTransformer3DModel.forward does (model.py:325-336, float64) and +applied inside attn1 via freqs_cis. c2ws_plucker_emb must match hidden_states +shape (model.py:52-55 assert) and is passed as a real tensor so all 8 +cam_conditioner tensors are exercised. Block returns a single tensor. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +INNER_DIM = 5120 # 40 heads x 128, LingBot-World transformer/config.json +TEXT_LEN = 512 + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.lingbotworld import LingBotWorldArchConfig + from fastvideo.models.dits.lingbotworld.model import LingBotWorldTransformerBlock + + arch = LingBotWorldArchConfig() + return LingBotWorldTransformerBlock( + INNER_DIM, + arch.ffn_dim, + arch.num_attention_heads, + arch.qk_norm, + arch.cross_attn_norm, + arch.eps, + arch.added_kv_proj_dim, + arch._supported_attention_backends, + prefix=f"Wan.blocks.{layer}", + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + + generator = torch.Generator(device="cpu").manual_seed(seed) + t, h, w = 3, 8, 14 + seq = t * h * w + + d = INNER_DIM // 40 + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (t, h, w), INNER_DIM, 40, rope_dim_list, rope_theta=10000, dtype=torch.float64) + freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) + + hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, TEXT_LEN, INNER_DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(1, 6, INNER_DIM, generator=generator, dtype=torch.float32) + cam = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "freqs_cis": freqs_cis, + "original_seq_len": seq, + "c2ws_plucker_emb": cam.to(device=device, dtype=torch.bfloat16), + } + + +SPEC = GateSpec( + name="lingbot_world_cam", + repo_id="FastVideo/LingBot-World-Base-Cam-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="blocks.{N}.", + renames=( + (r"^\.modulation$", ".scale_shift_table"), + (r"\.self_attn\.q\.", ".to_q."), + (r"\.self_attn\.k\.", ".to_k."), + (r"\.self_attn\.v\.", ".to_v."), + (r"\.self_attn\.o\.", ".to_out."), + (r"\.self_attn\.norm_q\.", ".norm_q."), + (r"\.self_attn\.norm_k\.", ".norm_k."), + (r"\.norm3\.", ".self_attn_residual_norm.norm."), + (r"\.cross_attn\.q\.", ".attn2.to_q."), + (r"\.cross_attn\.k\.", ".attn2.to_k."), + (r"\.cross_attn\.v\.", ".attn2.to_v."), + (r"\.cross_attn\.o\.", ".attn2.to_out."), + (r"\.cross_attn\.norm_q\.", ".attn2.norm_q."), + (r"\.cross_attn\.norm_k\.", ".attn2.norm_k."), + (r"\.ffn\.0\.", ".ffn.fc_in."), + (r"\.ffn\.2\.", ".ffn.fc_out."), + (r"\.cam_injector_layer1\.", ".cam_conditioner.cam_injector.fc_in."), + (r"\.cam_injector_layer2\.", ".cam_conditioner.cam_injector.fc_out."), + (r"\.cam_scale_layer\.", ".cam_conditioner.cam_scale_layer."), + (r"\.cam_shift_layer\.", ".cam_conditioner.cam_shift_layer."), + ), + attention_backend="FLASH_ATTN", +) + + +def test_lingbot_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_longcat.py b/fastvideo/tests/golden_gate/test_longcat.py new file mode 100644 index 000000000..1da7486e0 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_longcat.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: LongCat-Video (native FastVideo impl) transformer block 0. + +Spec notes: LongCatVideoConfig() arch defaults match the mirror's +transformer/config.json (hidden 4096, 32 heads, mlp_ratio 4, adaln 512). +The block __init__ needs the FULL LongCatVideoConfig, not the arch config — +self_attn reads config._supported_attention_backends (longcat.py:273) while +cross_attn reads config.arch_config._supported_attention_backends (:548). +Checkpoint is a SINGLE fp32 file (transformer/model.safetensors, no +index.json) with zero renames — the mirror is pre-converted, keys match +module names. 3D RoPE is built inside the block from latent_shape; t is +per-latent-frame temb [B, T, 512]. Defaults (return_kv=False) return one +plain tensor, so no postprocess. Backend FLASH_ATTN per the ssim test +(tests/ssim/test_longcat_similarity.py:200). +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +HIDDEN_SIZE = 4096 # configs/models/dits/longcat.py:77 +ADALN_DIM = 512 # configs/models/dits/longcat.py:93 + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.longcat import LongCatVideoConfig + from fastvideo.models.dits.longcat import LongCatTransformerBlock + + config = LongCatVideoConfig() + arch = config.arch_config + return LongCatTransformerBlock( + hidden_size=arch.hidden_size, + num_heads=arch.num_attention_heads, + mlp_ratio=arch.mlp_ratio, + adaln_tembed_dim=arch.adaln_tembed_dim, + config=config, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + # forward signature (longcat.py:773-783): x [B, N, 4096], context + # [B, N_text, 4096], t [B, T, 512], latent_shape (T, H, W), N == T*H*W. + generator = torch.Generator(device="cpu").manual_seed(seed) + t_frames, h, w = 3, 8, 8 + seq = t_frames * h * w + n_text = 32 # production pads to 512; any length works + + hidden = torch.randn(1, seq, HIDDEN_SIZE, generator=generator, dtype=torch.float32) + context = torch.randn(1, n_text, HIDDEN_SIZE, generator=generator, dtype=torch.float32) + temb = torch.randn(1, t_frames, ADALN_DIM, generator=generator, dtype=torch.float32) + + return { + "x": hidden.to(device=device, dtype=torch.bfloat16), + "context": context.to(device=device, dtype=torch.bfloat16), + "t": temb.to(device=device, dtype=torch.bfloat16), + "latent_shape": (t_frames, h, w), + } + + +SPEC = GateSpec( + name="longcat_t2v", + repo_id="FastVideo/LongCat-Video-T2V-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="blocks.{N}.", + renames=(), + attention_backend="FLASH_ATTN", + model_root_env="LONGCAT_MODEL_ROOT", + weight_file="model.safetensors", # single 54.3 GB fp32 file, no index.json +) + + +def test_longcat_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_ltx2.py b/fastvideo/tests/golden_gate/test_ltx2.py new file mode 100644 index 000000000..32da8aca8 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_ltx2.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: LTX-2 (LTX2-Distilled-Diffusers) BasicAVTransformerBlock 0. + +Spec notes: block geometry comes from LTX2VideoArchConfig defaults, which +match the checkpoint's transformer/config.json (video 32x128=4096, audio +32x64=2048, cross 4096/2048, rope "split", stg_block_idx 29). The block +consumes preprocessor outputs (TransformerArgs): AdaLN timestep tensors are +seeded randn with the widths get_ada_values expects (6*dim timesteps, +4*dim + dim AV-CA scale/shift + gate), and RoPE is built outside the block +exactly as TransformerArgsPreprocessor does (double_precision_rope -> +generate_ltx_freq_grid_float64), with cross-PE from positions[:, 0:1, :]. +The SSIM run uses sp_size=2 (LTXDistributedSelfAttention); the gate uses +use_distributed_attention=False — same weights, local attention. The +checkpoint is single-file (transformer/model.safetensors, no index.json). +context_mask=None keeps all four attentions on FLASH_ATTN; production text +cross-attn passes an additive mask that reroutes attn2 through the +TORCH_SDPA attn_masked instance — accepted minimality tradeoff. FLASH_ATTN +must be the backend: VIDEO_SPARSE_ATTN would create extra to_gate_compress +params and break strict load. Forward returns (video, audio) +TransformerArgs; postprocess flattens/concats both .x tensors. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +VIDEO_DIM = 4096 # 32 heads x 128, LTX2-Distilled transformer/config.json +AUDIO_DIM = 2048 # 32 heads x 64 + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.ltx2 import LTX2VideoArchConfig + from fastvideo.models.dits.ltx2 import ( + BasicAVTransformerBlock, LTXRopeType, TransformerConfig) + + arch = LTX2VideoArchConfig() + video_cfg = TransformerConfig( + dim=arch.num_attention_heads * arch.attention_head_dim, + heads=arch.num_attention_heads, + d_head=arch.attention_head_dim, + context_dim=arch.cross_attention_dim, + apply_gated_attention=arch.apply_gated_attention, + cross_attention_adaln=arch.cross_attention_adaln, + ) + audio_cfg = TransformerConfig( + dim=arch.audio_num_attention_heads * arch.audio_attention_head_dim, + heads=arch.audio_num_attention_heads, + d_head=arch.audio_attention_head_dim, + context_dim=arch.audio_cross_attention_dim, + apply_gated_attention=arch.apply_gated_attention, + cross_attention_adaln=arch.cross_attention_adaln, + ) + return BasicAVTransformerBlock( + idx=layer, + video=video_cfg, + audio=audio_cfg, + rope_type=LTXRopeType(arch.rope_type), + norm_eps=arch.norm_eps, + use_distributed_attention=False, + cross_attention_adaln=arch.cross_attention_adaln, + stg_block_idx=arch.stg_block_idx, + quant_config=None, + prefix="ltx2", + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.ltx2 import ( + AudioLatentPatchifier, AudioLatentShape, DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + DEFAULT_LTX2_AUDIO_HOP_LENGTH, DEFAULT_LTX2_AUDIO_MEL_BINS, + DEFAULT_LTX2_AUDIO_SAMPLE_RATE, DEFAULT_LTX2_SCALE_FACTORS, LTXRopeType, + TransformerArgs, VideoLatentPatchifier, VideoLatentShape, + _get_pixel_coords, generate_ltx_freq_grid_float64, + precompute_ltx_freqs_cis) + + generator = torch.Generator(device="cpu").manual_seed(seed) + B, frames, h, w = 1, 2, 4, 6 + seq_v = frames * h * w # 48 + n_text, seq_a = 32, 12 + + # Video positions: patch grid bounds -> pixel coords (scale 8,32,32; + # fps 24; causal fix), cast to hidden dtype like ltx2.py:3044. + v_pos = VideoLatentPatchifier(patch_size=1).get_patch_grid_bounds( + VideoLatentShape((B, 128, frames, h, w)), device=torch.device("cpu")) + v_pos = _get_pixel_coords(v_pos, DEFAULT_LTX2_SCALE_FACTORS, fps=24.0, + causal_fix=True).to(torch.bfloat16) + # Audio positions stay float32 [B, 1, T, 2] (ltx2.py:3099-3101). + a_pos = AudioLatentPatchifier( + patch_size=DEFAULT_LTX2_AUDIO_MEL_BINS, + sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE, + hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH, + audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE, + is_causal=True, shift=0, + ).get_patch_grid_bounds( + AudioLatentShape((B, 8, seq_a, DEFAULT_LTX2_AUDIO_MEL_BINS)), + device=torch.device("cpu")) + + rope_kw = dict(out_dtype=torch.bfloat16, theta=10000.0, + use_middle_indices_grid=True, num_attention_heads=32, + rope_type=LTXRopeType.SPLIT, + freq_grid_generator=generate_ltx_freq_grid_float64) + v_pe = precompute_ltx_freqs_cis(v_pos, dim=VIDEO_DIM, + max_pos=[20, 2048, 2048], **rope_kw) + v_cross_pe = precompute_ltx_freqs_cis(v_pos[:, 0:1, :], dim=AUDIO_DIM, + max_pos=[20], **rope_kw) + a_pe = precompute_ltx_freqs_cis(a_pos, dim=AUDIO_DIM, max_pos=[20], + **rope_kw) + a_cross_pe = precompute_ltx_freqs_cis(a_pos[:, 0:1, :], dim=AUDIO_DIM, + max_pos=[20], **rope_kw) + + def rnd(*shape): + return torch.randn(*shape, generator=generator, + dtype=torch.float32).to(device=device, + dtype=torch.bfloat16) + + def to_dev(pe): + return tuple(t.to(device) for t in pe) + + video = TransformerArgs( + x=rnd(B, seq_v, VIDEO_DIM), + context=rnd(B, n_text, VIDEO_DIM), + context_mask=None, + timesteps=rnd(B, seq_v, 6 * VIDEO_DIM), + embedded_timestep=rnd(B, seq_v, VIDEO_DIM), + positional_embeddings=to_dev(v_pe), + cross_positional_embeddings=to_dev(v_cross_pe), + cross_scale_shift_timestep=rnd(B, seq_v, 4 * VIDEO_DIM), + cross_gate_timestep=rnd(B, seq_v, VIDEO_DIM), + enabled=True, prompt_timestep=None) + audio = TransformerArgs( + x=rnd(B, seq_a, AUDIO_DIM), + context=rnd(B, n_text, AUDIO_DIM), + context_mask=None, + timesteps=rnd(B, seq_a, 6 * AUDIO_DIM), + embedded_timestep=rnd(B, seq_a, AUDIO_DIM), + positional_embeddings=to_dev(a_pe), + cross_positional_embeddings=to_dev(a_cross_pe), + cross_scale_shift_timestep=rnd(B, seq_a, 4 * AUDIO_DIM), + cross_gate_timestep=rnd(B, seq_a, AUDIO_DIM), + enabled=True, prompt_timestep=None) + + return {"video": video, "audio": audio, + "video_original_seq_len": seq_v, "audio_original_seq_len": seq_a} + + +SPEC = GateSpec( + name="ltx2_distilled", + repo_id="FastVideo/LTX2-Distilled-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer_blocks.{N}.", + renames=(), + attention_backend="FLASH_ATTN", + weight_file="model.safetensors", # single-file checkpoint, no index.json + # forward returns (video, audio) TransformerArgs — reduce to one tensor + postprocess=lambda out: torch.cat([out[0].x.flatten(), out[1].x.flatten()]), +) + + +def test_ltx2_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_matrixgame.py b/fastvideo/tests/golden_gate/test_matrixgame.py new file mode 100644 index 000000000..0832b1f5b --- /dev/null +++ b/fastvideo/tests/golden_gate/test_matrixgame.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: MatrixGame family — TWO distinct DiTs, one gate each. + +MG2 = CausalMatrixGame2TransformerBlock (matrixgame2/causal_model.py:385), the +CAUSAL block the checkpoint's _class_name (CausalMatrixGame2WanModel) loads — +NOT model.py's MatrixGame2TransformerBlock. MG3 = MatrixGame3TransformerBlock +(matrixgame3/model.py:199). + +Spec notes: MG2 arch-config defaults are Wan-14B sized and WRONG for this +checkpoint; real dims (1536/8960/12 heads/local_attn_size 6) and action_config +come from the checkpoint's transformer/config.json, fetched at build time. MG3 +arch-class defaults DO match its checkpoint. Both checkpoints are single-file +safetensors (no index.json) in FastVideo-native (MG2) / self_attn-style (MG3) +names. Rope is handled INSIDE both blocks: MG2's freqs_cis argument is dead +(attn1 rebuilds rope internally, causal_model.py:200-211) so pass None; MG3 +takes a caller-built per-head complex table via _build_rope_freqs +(use_memory=True -> sigma_theta 0.8, max_seq_len 2048, model.py:640-646). +Geometry (3, 22, 40) is load-bearing for MG2: seq % 128 != 0 and f % 32 != 0 +or the flex-attention unpad slices `[:, :, :-0]` go empty +(causal_model.py:224/266/558). temb is 4-dim for both (MG2 asserts it, +causal_model.py:521; MG3 per-token path, model.py:309). Layer 0 covers the +ActionModule (blocks 0-14 only); mouse/keyboard conds use pixel length +4*(f-1)+1. MG2's cache-less self-attn runs torch.compile'd flex_attention +(max-autotune) — first call compiles, budget accordingly. Both blocks return a +single tensor; no postprocess needed. Both ssim tests use FLASH_ATTN. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +MG2_REPO = "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers" +MG3_REPO = "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers" + +MG2_INNER_DIM = 1536 # 12 heads x 128, MG2 transformer/config.json +MG3_INNER_DIM = 3072 # 24 heads x 128, matrixgame3.py:42-43 + +F, H, W = 3, 22, 40 # post-patch grid; seq=2640 (% 128 != 0, see docstring) +N_PIX = 4 * (F - 1) + 1 # pixel-frame action length (action_module.py:511-521) + + +# ===== Matrix-Game-2.0 (causal) ===== + + +def _build_block_mg2(layer: int) -> torch.nn.Module: + import json + from pathlib import Path + + from huggingface_hub import hf_hub_download + + from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoArchConfig + from fastvideo.models.dits.matrixgame2 import CausalMatrixGame2TransformerBlock + + arch = MatrixGame2WanVideoArchConfig() + # ponytail: arch-config defaults are Wan-14B and wrong here; the checkpoint's + # config.json is the source of truth (hidden 1536, ffn 8960, 12 heads). + hf = json.loads(Path(hf_hub_download(MG2_REPO, "transformer/config.json")).read_text()) + inner_dim = hf["num_attention_heads"] * hf["attention_head_dim"] + assert inner_dim == MG2_INNER_DIM + return CausalMatrixGame2TransformerBlock( + inner_dim, + hf["ffn_dim"], # 8960 + hf["num_attention_heads"], # 12 + hf["local_attn_size"], # 6 + hf["sink_size"], # 0 + hf["qk_norm"], # "rms_norm_across_heads" + arch.cross_attn_norm, # True (block asserts it, causal_model.py:436) + hf["eps"], # 1e-6 + arch.added_kv_proj_dim, # None -> attn2 = WanT2VCrossAttention + arch._supported_attention_backends, + prefix=f"Wan.blocks.{layer}", + action_config=hf["action_config"], # ActionModule on blocks 0-14 only + block_idx=layer, + rope_cache_policy=arch.rope_cache_policy, # "absolute" + ) + + +def _make_inputs_mg2(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.matrixgame2 import CausalMatrixGame2WanModel + + generator = torch.Generator(device="cpu").manual_seed(seed) + seq = F * H * W # 2640 + + hidden = torch.randn(1, seq, MG2_INNER_DIM, generator=generator, dtype=torch.float32) + # CLIP image embeds post WanImageEmbedding + context = torch.randn(1, 257, MG2_INNER_DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(1, F, 6, MG2_INNER_DIM, generator=generator, dtype=torch.float32) + mouse = torch.randn(1, N_PIX, 2, generator=generator, dtype=torch.float32) # mouse_dim_in=2 + keyboard = torch.randn(1, N_PIX, 4, generator=generator, dtype=torch.float32) # keyboard_dim_in=4 + + # masks exactly as _forward_inference builds them (causal_model.py:1129-1162) + block_mask = CausalMatrixGame2WanModel._prepare_blockwise_causal_attn_mask( + device=device, num_frames=F, frame_seqlen=H * W, + num_frame_per_block=3, local_attn_size=6, sink_size=0) + block_mask_action = CausalMatrixGame2WanModel._prepare_blockwise_causal_attn_mask_action( + device=device, num_frames=F, frame_seqlen=1, + num_frame_per_block=3, local_attn_size=6, sink_size=0) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.float32), + "encoder_hidden_states": context.to(device=device, dtype=torch.float32), + "temb": temb.to(device=device, dtype=torch.float32), + "freqs_cis": None, # dead arg: attn1 rebuilds rope internally + "block_mask": block_mask, + "grid_sizes": (F, H, W), + "mouse_cond": mouse.to(device=device, dtype=torch.float32), + "keyboard_cond": keyboard.to(device=device, dtype=torch.float32), + "block_mask_mouse": block_mask_action, + "block_mask_keyboard": block_mask_action, # use_rope_keyboard=True path + "num_frame_per_block": 3, + "use_rope_keyboard": True, # asserted in ActionModule.forward + # kv_cache* / crossattn_cache None, current_start=0 -> cache-less flex path + } + + +MG2_SPEC = GateSpec( + name="matrixgame2_causal", + dtype=torch.float32, + repo_id=MG2_REPO, + build_block=_build_block_mg2, + make_inputs=_make_inputs_mg2, + prefix_template="blocks.{N}.", + renames=(), # mirror checkpoint is already in FastVideo-native names + attention_backend="FLASH_ATTN", +) + + +def test_matrixgame2_golden_gate(distributed_runtime) -> None: + run_gate(MG2_SPEC) + + +# ===== Matrix-Game-3.0 ===== + + +def _build_block_mg3(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.matrixgame3 import MatrixGame3WanVideoArchConfig + from fastvideo.models.dits.matrixgame3 import MatrixGame3TransformerBlock + + arch = MatrixGame3WanVideoArchConfig() + inner_dim = arch.num_attention_heads * arch.attention_head_dim # 24*128 = 3072 + assert inner_dim == MG3_INNER_DIM + return MatrixGame3TransformerBlock( + inner_dim, + arch.ffn_dim, # 14336 + arch.num_attention_heads, # 24 + arch.qk_norm, # "rms_norm_across_heads" + arch.cross_attn_norm, # True (asserted at model.py:246) + arch.eps, # 1e-6 + arch._supported_attention_backends, + prefix=f"Wan.blocks.{layer}", + action_config=arch.action_config, # blocks 0-14 only + block_id=layer, + use_memory=arch.use_memory, # True -> cam_* Linears on EVERY block + ) + + +def _make_inputs_mg3(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.matrixgame3.model import _build_rope_freqs + + generator = torch.Generator(device="cpu").manual_seed(seed) + seq = F * H * W # 2640 + + hidden = torch.randn(1, seq, MG3_INNER_DIM, generator=generator, dtype=torch.float32) + # text tokens AFTER model-level text_embedding proj (model.py:715-726) + context = torch.randn(1, 512, MG3_INNER_DIM, generator=generator, dtype=torch.float32) + # per-token 4-D temb (bs, seq, 6, 3072), real path (model.py:687-713) + temb = torch.randn(1, seq, 6, MG3_INNER_DIM, generator=generator, dtype=torch.float32) + mouse = torch.randn(1, N_PIX, 2, generator=generator, dtype=torch.float32) # mouse_dim_in=2 + keyboard = torch.randn(1, N_PIX, 6, generator=generator, dtype=torch.float32) # keyboard_dim_in=6 + + # exactly as model.py:640-646: use_memory=True -> per-head 3-D complex table + freqs = _build_rope_freqs(max_seq_len=2048, head_dim=128, num_heads=24, + sigma_theta=0.8, device=device) # (24, 2048, 64) complex128 + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": context.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "freqs": freqs, + "grid_sizes": (F, H, W), + "mouse_cond": mouse.to(device=device, dtype=torch.bfloat16), + "keyboard_cond": keyboard.to(device=device, dtype=torch.bfloat16), + # plucker_emb=None -> cam_* layers loaded but skipped; + # memory_length=0 / memory_latent_idx=None -> arange fallback + } + + +MG3_SPEC = GateSpec( + name="matrixgame3", + repo_id=MG3_REPO, + build_block=_build_block_mg3, + make_inputs=_make_inputs_mg3, + prefix_template="blocks.{N}.", + renames=( + (r"\.self_attn\.q\.", ".to_q."), + (r"\.self_attn\.k\.", ".to_k."), + (r"\.self_attn\.v\.", ".to_v."), + (r"\.self_attn\.o\.", ".to_out."), + (r"\.self_attn\.norm_q\.", ".norm_q."), + (r"\.self_attn\.norm_k\.", ".norm_k."), + (r"\.cross_attn\.q\.", ".attn2.to_q."), + (r"\.cross_attn\.k\.", ".attn2.to_k."), + (r"\.cross_attn\.v\.", ".attn2.to_v."), + (r"\.cross_attn\.o\.", ".attn2.to_out."), + (r"\.cross_attn\.norm_q\.", ".attn2.norm_q."), + (r"\.cross_attn\.norm_k\.", ".attn2.norm_k."), + (r"\.ffn\.0\.", ".ffn.fc_in."), + (r"\.ffn\.2\.", ".ffn.fc_out."), + (r"\.norm3\.", ".self_attn_residual_norm.norm."), + (r"\.modulation$", ".scale_shift_table"), + # cam_* and action_model.* keys already match module names + ), + attention_backend="FLASH_ATTN", +) + + +def test_matrixgame3_golden_gate(distributed_runtime) -> None: + run_gate(MG3_SPEC) diff --git a/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py b/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py new file mode 100644 index 000000000..00b6c828a --- /dev/null +++ b/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate for the MiniMax H3 T2V DiT: single-layer bitwise fingerprint. + +Loads ONE transformer block of the 62 GB checkpoint (the safetensors index +locates the ~1.2 GB of layer tensors; only the shards containing them are +fetched), feeds a fixed seeded packed batch exercising all three modalities, +and compares the output **bitwise** against a device-keyed golden latent. + +Zero tolerance is deliberate, not optimistic. Verified on GB200 2026-08-05: + * this block reproduces bit-identically across independent processes, + * the full 50-layer forward holds ``atol=0`` against pinned Diffusers + (tests/local_tests/transformers/test_minimax_h3_transformer_parity.py), + * two end-to-end seeded T2V generations produced md5-identical videos. +Any drift here is therefore a real change to the compute path — the gate +stands in front of the ~13-minute full H3 SSIM generation at ~40 seconds. + +The golden's environment fingerprint (torch/CUDA/cuDNN/flash-attn/TF32/ +backend/FASTVIDEO_FA4) is asserted before the tensor compare, so a legitimate +environment change fails as "environment changed — re-mint deliberately", +never as unexplained bit drift. + +Goldens are device-keyed like SSIM references and live in the same HF dataset +(``FastVideo/ssim-reference-videos`` under ``golden_gates//``); a +local ``goldens/`` directory or ``MINIMAX_H3_GATE_GOLDEN_DIR`` wins over the +download. A missing golden is seeded locally and the test fails with upload +instructions, mirroring the SSIM missing-reference convention. +""" + +from __future__ import annotations + +import json +import os +import re +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +GOLDEN_REPO_ID = os.environ.get("FASTVIDEO_SSIM_REFERENCE_HF_REPO", "FastVideo/ssim-reference-videos") +MODEL_REPO_ID = "MiniMaxAI/MiniMax-H3" +LAYER = int(os.environ.get("MINIMAX_H3_GATE_LAYER", "0")) +SEED = 20260805 + +# Pin the same attention path the H3 SSIM test runs, before any fastvideo +# import can cache a backend decision. CUBLAS workspace must be set before +# cuBLAS initializes for use_deterministic_algorithms to hold. +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") +os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") +os.environ.setdefault("MASTER_ADDR", "127.0.0.1") +os.environ.setdefault("MASTER_PORT", "29661") +os.environ.setdefault("RANK", "0") +os.environ.setdefault("WORLD_SIZE", "1") +os.environ.setdefault("LOCAL_RANK", "0") + +# Checkpoint -> FastVideo renames, the block-level subset of +# MiniMaxH3Config.param_names_mapping (fastvideo/configs/models/dits/minimax_h3.py). +_RENAMES = ( + (re.compile(r"\.ff\.net\.0\.proj\."), ".ff.fc_in."), + (re.compile(r"\.ff\.net\.2\."), ".ff.fc_out."), + (re.compile(r"\.attn\.to_out\.0\."), ".attn.to_out."), +) + + +@pytest.fixture(scope="module", autouse=True) +def _distributed_runtime(): + if not torch.cuda.is_available(): + yield + return + from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel, + ) + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + cleanup_dist_env_and_memory() + + +def _env_fingerprint() -> dict[str, str]: + """Everything that legitimately changes numerics. A golden minted under a + different fingerprint must fail as 'environment changed', never as + unexplained bit drift.""" + try: + import flash_attn + fa_version = str(getattr(flash_attn, "__version__", "?")) + except ImportError: + fa_version = "" + return { + "torch": str(torch.__version__), + "cuda": str(torch.version.cuda), + "cudnn": str(torch.backends.cudnn.version()), + "flash_attn": fa_version, + "device_name": torch.cuda.get_device_name(0), + "attention_backend": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", ""), + "fastvideo_fa4": os.environ.get("FASTVIDEO_FA4", ""), + "tf32_matmul": str(torch.backends.cuda.matmul.allow_tf32), + "tf32_cudnn": str(torch.backends.cudnn.allow_tf32), + } + + +def _component_dir() -> Path: + """Local checkout wins; otherwise fetch only the index + the shards that + hold this layer's tensors (~4.5 GB of the 62 GB checkpoint).""" + root = os.environ.get("MINIMAX_H3_MODEL_ROOT") + if root: + component = Path(root) / "transformer" + if not (component / "diffusion_pytorch_model.safetensors.index.json").is_file(): + pytest.fail(f"MINIMAX_H3_MODEL_ROOT set but index missing under {component}", pytrace=False) + return component + + from huggingface_hub import hf_hub_download + + index_path = Path(hf_hub_download( + MODEL_REPO_ID, "transformer/diffusion_pytorch_model.safetensors.index.json")) + weight_map = json.loads(index_path.read_text())["weight_map"] + prefix = f"transformer_blocks.{LAYER}." + shards = sorted({shard for name, shard in weight_map.items() if name.startswith(prefix)}) + if not shards: + pytest.fail(f"no shards found for prefix {prefix} in {MODEL_REPO_ID}", pytrace=False) + for shard in shards: + hf_hub_download(MODEL_REPO_ID, f"transformer/{shard}") + return index_path.parent + + +def _load_layer_state(component_dir: Path, layer: int) -> dict[str, torch.Tensor]: + """Load only ``transformer_blocks.{layer}.*`` tensors via the index.""" + from safetensors import safe_open + + index = json.loads((component_dir / "diffusion_pytorch_model.safetensors.index.json").read_text()) + prefix = f"transformer_blocks.{layer}." + by_shard: dict[str, list[str]] = {} + for name, shard in index["weight_map"].items(): + if name.startswith(prefix): + by_shard.setdefault(shard, []).append(name) + assert by_shard, f"no parameters found for prefix {prefix}" + + state: dict[str, torch.Tensor] = {} + for shard, names in sorted(by_shard.items()): + with safe_open(component_dir / shard, framework="pt", device="cpu") as f: + for name in names: + mapped = f".{name[len(prefix):]}" + for pattern, repl in _RENAMES: + mapped = pattern.sub(repl, mapped) + state[mapped[1:]] = f.get_tensor(name) + return state + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig + from fastvideo.models.dits.minimax_h3 import MiniMaxH3TransformerBlock + + arch = MiniMaxH3ArchConfig() + return MiniMaxH3TransformerBlock( + arch.hidden_size, + arch.num_attention_heads, + arch.attention_head_dim, + arch.ffn_dim, + arch.time_embed_dim, + arch.norm_eps, + arch.qk_norm_eps, + arch._supported_attention_backends, + None, + prefix=f"minimax_h3.transformer_blocks.{layer}", + ) + + +def _make_inputs(device: torch.device) -> dict: + """Fixed packed batch: 32 text + 24 audio + 200 video rows, 3 timesteps.""" + from fastvideo.models.dits.minimax_h3 import MiniMaxH3RotaryPosEmbed + + generator = torch.Generator(device="cpu").manual_seed(SEED) + n_text, n_audio, n_video = 32, 24, 200 + seq = n_text + n_audio + n_video + + token_tags = torch.cat([ + torch.full((n_text,), 1, dtype=torch.long), + torch.full((n_audio,), 2, dtype=torch.long), + torch.full((n_video,), 0, dtype=torch.long), + ]) + timestep_indices = torch.zeros(seq, dtype=torch.long) + timestep_indices[n_text:n_text + n_audio] = 1 + timestep_indices[n_text + n_audio:n_text + n_audio + 100] = 2 + adaln_indices = timestep_indices * 3 + token_tags # MINIMAX_H3_MODALITY_NUM = 3 + + position_ids = torch.zeros(seq, 3, dtype=torch.float64) + position_ids[:, 0] = torch.arange(seq, dtype=torch.float64) + vid = torch.arange(n_video, dtype=torch.float64) + position_ids[n_text + n_audio:, 1] = vid % 10 + position_ids[n_text + n_audio:, 2] = vid // 10 + + rope = MiniMaxH3RotaryPosEmbed(rope_freq_dim=16, rope_theta=10000.0) + cos, sin = rope(position_ids.to(torch.float32)) + + hidden = torch.randn(1, seq, 5376, generator=generator, dtype=torch.float32) + temb = torch.randn(9, 2688, generator=generator, dtype=torch.float32) # 3 timesteps x 3 modalities + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "adaln_indices": adaln_indices.to(device), + "rotary_emb": (cos.to(device), sin.to(device)), + "original_seq_len": seq, + } + + +def _device_slug() -> str: + return re.sub(r"[^A-Za-z0-9]+", "_", torch.cuda.get_device_name(0)).strip("_") + + +def _golden_filename() -> str: + backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "default") + return f"minimax_h3_t2v_layer{LAYER}_{backend}_seed{SEED}.pt" + + +def _resolve_golden() -> tuple[Path, bool]: + """Return (path, exists). Local dir wins; falls back to the HF dataset.""" + local_root = Path(os.environ.get( + "MINIMAX_H3_GATE_GOLDEN_DIR", Path(__file__).resolve().parent / "goldens")) + local = local_root / _device_slug() / _golden_filename() + if local.exists(): + return local, True + + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + try: + remote = hf_hub_download( + GOLDEN_REPO_ID, + f"golden_gates/{_device_slug()}/{_golden_filename()}", + repo_type="dataset", + ) + return Path(remote), True + except EntryNotFoundError: + return local, False + + +def test_minimax_h3_t2v_golden_gate() -> None: + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("golden-gate requires a bf16-capable CUDA GPU") + + # Tripwire: any nondeterministic op added to the block errors immediately + # instead of flaking the gate. benchmark=False kills cuDNN autotune + # variance (no-op today — the block has no convs — but the gate may grow). + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.benchmark = False + + device = torch.device("cuda:0") + component_dir = _component_dir() + + from fastvideo.forward_context import set_forward_context + + block = _build_block(LAYER) + state = _load_layer_state(component_dir, LAYER) + missing, unexpected = block.load_state_dict(state, strict=True) + assert not missing and not unexpected + block = block.to(device=device, dtype=torch.bfloat16).eval() + + inputs = _make_inputs(device) + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + out = block(**inputs) + out = out.detach().float().cpu() + + golden_path, golden_exists = _resolve_golden() + meta = { + "layer": LAYER, + "seed": SEED, + "shape": list(out.shape), + "env": _env_fingerprint(), + } + if not golden_exists: + golden_path.parent.mkdir(parents=True, exist_ok=True) + torch.save({"output": out, "metadata": meta}, golden_path) + pytest.fail( + f"golden latent seeded at {golden_path} ({meta}). Verify with a second " + f"run, then upload to {GOLDEN_REPO_ID} at " + f"golden_gates/{_device_slug()}/{_golden_filename()}", + pytrace=False, + ) + + golden = torch.load(golden_path, weights_only=True) + mismatches = { + key: (value, meta["env"].get(key)) + for key, value in golden["metadata"]["env"].items() + if meta["env"].get(key) != value + } + assert not mismatches, ( + f"environment changed since golden was minted — re-mint deliberately rather " + f"than chasing bit drift. golden -> current: {mismatches}") + + drift = (out - golden["output"]).abs() + print(f"golden-gate drift: max_abs={drift.max().item():.10f} " + f"mean_abs={drift.mean().item():.10f}", flush=True) + assert_close(out, golden["output"], atol=0.0, rtol=0.0) diff --git a/fastvideo/tests/golden_gate/test_sd35.py b/fastvideo/tests/golden_gate/test_sd35.py new file mode 100644 index 000000000..659f36361 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_sd35.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: SD3.5-medium MMDiT joint transformer block 0. + +Spec notes: arch-config defaults (configs/models/dits/sd3.py:11-23) match the +HF transformer/config.json of stabilityai/stable-diffusion-3.5-medium +field-for-field, so no explicit geometry pins. Layer 0 is dual-attention +(layers 0-12 carry extra attn2.* params; 23 is context_pre_only and returns +(None, hidden) — stay below 23). Checkpoint is SINGLE-FILE (no index.json); +the harness single-file fallback handles it, and param names match diffusers +bit-for-bit (param_names_mapping is empty, sd3.py:878) so renames are empty. +Forward returns a (encoder_hidden_states, hidden_states) TUPLE (sd3.py:868); +postprocess cats the two streams along dim=1. Inputs are the three (B, *, 1536) +tensors only — no RoPE, no mask; sincos pos embed is added before the blocks +(sd3.py:1024). SSIM test pins TORCH_SDPA (test_sd35_similarity.py:93). +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +DIM = 1536 # 24 heads x 64, sd3.py:896 +N_IMG = 256 # 256x256 SSIM run: 32x32 latent, patch 2 +N_TXT = 154 + + +def _build_block(layer: int) -> torch.nn.Module: + # Mirrors SD3Transformer2DModel.__init__ (sd3.py:923-933). + from fastvideo.configs.models.dits.sd3 import SD3Transformer2DArchConfig + from fastvideo.models.dits.sd3 import (SD3JointTransformerBlock, + SD3Transformer2DModel) + + arch = SD3Transformer2DArchConfig() + return SD3JointTransformerBlock( + dim=arch.num_attention_heads * arch.attention_head_dim, + num_attention_heads=arch.num_attention_heads, + attention_head_dim=arch.attention_head_dim, + context_pre_only=layer == arch.num_layers - 1, + qk_norm=arch.qk_norm, + use_dual_attention=layer in tuple(arch.dual_attention_layers), + supported_attention_backends=SD3Transformer2DModel._supported_attention_backends, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + # Block forward signature: sd3.py:771-777. dim=1536 everywhere: encoder + # states enter AFTER context_embedder (4096->1536), temb is the (B, 1536) + # time_text_embed output; AdaLN chunks temb along dim=1 so it stays 2-D. + generator = torch.Generator(device="cpu").manual_seed(seed) + + hidden = torch.randn(1, N_IMG, DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, N_TXT, DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(1, DIM, generator=generator, dtype=torch.float32) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + } + + +SPEC = GateSpec( + name="sd35_medium", + repo_id="stabilityai/stable-diffusion-3.5-medium", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer_blocks.{N}.", + attention_backend="TORCH_SDPA", + # tuple (encoder_hidden_states, hidden_states) -> one tensor, both (B, S, 1536) + postprocess=lambda out: torch.cat([out[0], out[1]], dim=1), +) + + +def test_sd35_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_stable_audio.py b/fastvideo/tests/golden_gate/test_stable_audio.py new file mode 100644 index 000000000..6da62edb7 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_stable_audio.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: Stable Audio Open 1.0 ContinuousTransformer block 0. + +Spec notes: geometry from StableAudioArchConfig defaults (embed_dim=1536, +24 heads -> dim_heads=64, cond_token_dim=768, project_cond_tokens=False so +dim_context=768, qk_norm=None for the 1.0 base). RoPE is built outside the +block exactly as ContinuousTransformer.forward does — RotaryEmbedding(32) +.forward_from_seq_len(seq) gives (freqs [seq, 32*2->partial rot_dim=32 of +head_dim=64], 1.0) applied inside Attention (stable_audio.py:179-194), not +via LocalAttention freqs_cis. Checkpoint is SINGLE-FILE (no index.json); +zero renames — block-relative keys match FastVideo module names exactly. +Backend pinned to TORCH_SDPA to match the SSIM lane. The pipeline runs the +DiT in fp16 but the harness casts blocks to bf16; the golden is minted and +compared in bf16 (fingerprint consistency, not pipeline parity). +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +EMBED_DIM = 1536 # StableAudioArchConfig.embed_dim +COND_DIM = 768 # StableAudioArchConfig.cond_token_dim (project_cond_tokens=False) + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.stable_audio import StableAudioArchConfig + from fastvideo.models.dits.stable_audio import TransformerBlock + + arch = StableAudioArchConfig() + return TransformerBlock( + dim=arch.embed_dim, + dim_heads=arch.embed_dim // arch.num_attention_heads, + cross_attend=True, + dim_context=arch.cond_token_dim, + zero_init_branch_outputs=True, + qk_norm=arch.qk_norm, + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.models.dits.stable_audio import RotaryEmbedding + + generator = torch.Generator(device="cpu").manual_seed(seed) + seq, ctx = 1 + 64, 65 # prepended global-cond token + 64 latent frames; T5 rows + + hidden = torch.randn(1, seq, EMBED_DIM, generator=generator, dtype=torch.float32) + context = torch.randn(1, ctx, COND_DIM, generator=generator, dtype=torch.float32) + freqs, scale = RotaryEmbedding(32).forward_from_seq_len(seq) # [seq, 32] fp32 + + return { + "x": hidden.to(device=device, dtype=torch.bfloat16), + "context": context.to(device=device, dtype=torch.bfloat16), + "rotary_pos_emb": (freqs.to(device), scale), + } + + +SPEC = GateSpec( + name="stable_audio_open_1_0", + repo_id="FastVideo/stable-audio-open-1.0-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="transformer.layers.{N}.", + renames=(), + attention_backend="TORCH_SDPA", +) + + +def test_stable_audio_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_wan_t2v.py b/fastvideo/tests/golden_gate/test_wan_t2v.py new file mode 100644 index 000000000..e811e1901 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_wan_t2v.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: Wan2.1-T2V-1.3B transformer block 0 (also the SFWan/TurboWan arch). + +Spec notes: arch-config defaults are the 14B model, so the 1.3B geometry is +pinned explicitly from the checkpoint's transformer/config.json. RoPE is built +outside the block exactly as WanTransformer3DModel.forward does +(wanvideo.py:707-719). attn1 is DistributedAttention (needs the distributed +fixture + forward context); temb must be the 3-dim [B, 6, inner_dim] wan2.1 +shape — 4-dim takes the wan2.2-ti2v per-token path. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +INNER_DIM = 1536 # 12 heads x 128, Wan2.1-T2V-1.3B transformer/config.json +TEXT_LEN = 512 + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig + from fastvideo.models.dits.wanvideo import WanTransformerBlock + + arch = WanVideoArchConfig( + num_attention_heads=12, + attention_head_dim=128, + ffn_dim=8960, + num_layers=30, + ) + return WanTransformerBlock( + INNER_DIM, + arch.ffn_dim, + arch.num_attention_heads, + arch.qk_norm, + arch.cross_attn_norm, + arch.eps, + arch.added_kv_proj_dim, + arch._supported_attention_backends, + quant_config=None, + prefix=f"Wan.blocks.{layer}", + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + + generator = torch.Generator(device="cpu").manual_seed(seed) + t, h, w = 3, 10, 10 + seq = t * h * w + + d = INNER_DIM // 12 + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (t, h, w), INNER_DIM, 12, rope_dim_list, rope_theta=10000, dtype=torch.float64) + freqs_cis = (freqs_cos.to(device).float(), freqs_sin.to(device).float()) + + hidden = torch.randn(1, seq, INNER_DIM, generator=generator, dtype=torch.float32) + encoder = torch.randn(1, TEXT_LEN, INNER_DIM, generator=generator, dtype=torch.float32) + temb = torch.randn(1, 6, INNER_DIM, generator=generator, dtype=torch.float32) + + return { + "hidden_states": hidden.to(device=device, dtype=torch.bfloat16), + "encoder_hidden_states": encoder.to(device=device, dtype=torch.bfloat16), + "temb": temb.to(device=device, dtype=torch.bfloat16), + "freqs_cis": freqs_cis, + "original_seq_len": seq, + } + + +SPEC = GateSpec( + name="wan_t2v_1_3b", + repo_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="blocks.{N}.", + renames=( + (r"\.attn1\.to_q\.", ".to_q."), + (r"\.attn1\.to_k\.", ".to_k."), + (r"\.attn1\.to_v\.", ".to_v."), + (r"\.attn1\.to_out\.0\.", ".to_out."), + (r"\.attn1\.norm_q\.", ".norm_q."), + (r"\.attn1\.norm_k\.", ".norm_k."), + (r"\.attn2\.to_out\.0\.", ".attn2.to_out."), + (r"\.ffn\.net\.0\.proj\.", ".ffn.fc_in."), + (r"\.ffn\.net\.2\.", ".ffn.fc_out."), + (r"\.norm2\.", ".self_attn_residual_norm.norm."), + ), + attention_backend="FLASH_ATTN", +) + + +def test_wan_t2v_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_zimage.py b/fastvideo/tests/golden_gate/test_zimage.py new file mode 100644 index 000000000..d282744d6 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_zimage.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden-gate: Z-Image-Turbo (Tongyi-MAI) main-stack transformer block 0. + +Spec notes: ZImageDiTArchConfig defaults match the checkpoint's +transformer/config.json (dim=3840, n_heads=n_kv_heads=30, adaln_embed_dim=256). +Checkpoint keys match FastVideo module names exactly — no renames. Rotary is +built OUTSIDE the block: freqs_cis is complex64 from RopeEmbedder(theta=256, +axes_dims=(32,48,48)) and must stay complex64 (rotary math runs fp32 inside +attention, zimage.py:104-108). adaln_input is width min(dim, adaln_embed_dim) += 256, not dim. Attention is plain F.scaled_dot_product_attention with a bool +[B, S] key-padding mask (True=attend); the block ignores fastvideo attention +backends, but TORCH_SDPA matches the ssim test and keys the golden filename. +Block returns a bare tensor [B, S, dim], no postprocess needed. +""" + +from __future__ import annotations + +import torch + +from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate + +__all__ = ["distributed_runtime"] + +DIM = 3840 # 30 heads x 128, Z-Image-Turbo transformer/config.json + + +def _build_block(layer: int) -> torch.nn.Module: + from fastvideo.configs.models.dits.zimage import ZImageDiTArchConfig + from fastvideo.models.dits.zimage import ZImageTransformerBlock + + arch = ZImageDiTArchConfig() + return ZImageTransformerBlock( + layer, # main `layers` stack uses the plain index (zimage.py:337-338) + arch.dim, + arch.n_heads, + arch.n_kv_heads, + arch.norm_eps, + arch.qk_norm, + arch.adaln_embed_dim, + modulation=True, # main-stack blocks are modulated (zimage.py:337-338) + ) + + +def _make_inputs(device: torch.device, seed: int) -> dict: + """Fixed padded batch: 2 sequences over an 8x8 image grid, lengths [64, 48].""" + from fastvideo.models.dits.zimage import RopeEmbedder, ZImageTransformer2DModel + + generator = torch.Generator(device="cpu").manual_seed(seed) + batch, seq = 2, 64 + lengths = [64, 48] + + # freqs_cis exactly as the model builds it (zimage.py:342, 527-531); image + # position ids start at frame index cap_len+1 -> use start=(1, 0, 0). + rope = RopeEmbedder(theta=256.0, axes_dims=(32, 48, 48), axes_lens=(1536, 512, 512)) + pos_ids = ZImageTransformer2DModel.create_coordinate_grid( + (1, 8, 8), start=(1, 0, 0)).flatten(0, 2) # [64, 3] int32 + freqs_cis = rope(pos_ids).unsqueeze(0).expand(batch, seq, 64).contiguous() # complex64 [2, 64, 64] + + attn_mask = torch.zeros(batch, seq, dtype=torch.bool) + for i, length in enumerate(lengths): + attn_mask[i, :length] = True + + x = torch.randn(batch, seq, DIM, generator=generator, dtype=torch.float32) + adaln = torch.randn(batch, 256, generator=generator, dtype=torch.float32) + + return { + "x": x.to(device=device, dtype=torch.bfloat16), + "attn_mask": attn_mask.to(device), + "freqs_cis": freqs_cis.to(device), # stays complex64, never cast to bf16 + "adaln_input": adaln.to(device=device, dtype=torch.bfloat16), + } + + +SPEC = GateSpec( + name="zimage_turbo", + repo_id="Tongyi-MAI/Z-Image-Turbo", + build_block=_build_block, + make_inputs=_make_inputs, + prefix_template="layers.{N}.", + renames=(), + attention_backend="TORCH_SDPA", + model_root_env="ZIMAGE_MODEL_DIR", +) + + +def test_zimage_golden_gate(distributed_runtime) -> None: + run_gate(SPEC) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 3f334b584..4124271f4 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -281,6 +281,21 @@ def run_vae_tests(): ) +@app.function(gpu="L40S:1", + image=image, + timeout=900, + secrets=[hf_secret, ci_env_secret], + volumes={"/root/data": model_vol}) +def run_golden_gate_tests(): + # Single-layer bitwise DiT fingerprints (~40s/model on GPU): a green gate + # means the compute path is bit-identical to the golden, so the expensive + # SSIM generation for that model cannot have regressed. Downloads only the + # shards holding the gated layer, never full checkpoints. + run_test( + "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/golden_gate -vs" + ) + + @app.function(gpu="L40S:1", image=image, timeout=900,