[ci]: add golden-gate lane — single-layer bitwise DiT fingerprints for all SSIM-covered families (#1682)
This commit is contained in:
William Lin
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 = "<not installed>"
|
||||
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", "<unset>"),
|
||||
"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()
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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/<device>/``); 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 = "<not installed>"
|
||||
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", "<unset>"),
|
||||
"fastvideo_fa4": os.environ.get("FASTVIDEO_FA4", "<unset>"),
|
||||
"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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user