From 574bdc44dcfcda3afac3704abeaad7a0c4ac7bcc Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 22 Jul 2026 19:45:52 -0700 Subject: [PATCH] [fix]: fp8 device-placement check; capture shim disables fsdp/offload like the released example; fastwan anchor CLI --- fastvideo2/layers/fp8.py | 2 +- .../wan21/gates/capture_fastvideo_main.py | 20 ++- fastvideo2/wan21/gates/fastwan_anchor.py | 137 ++++++++++++++++++ 3 files changed, 152 insertions(+), 7 deletions(-) create mode 100644 fastvideo2/wan21/gates/fastwan_anchor.py diff --git a/fastvideo2/layers/fp8.py b/fastvideo2/layers/fp8.py index 9898bee89..deec0ed9b 100644 --- a/fastvideo2/layers/fp8.py +++ b/fastvideo2/layers/fp8.py @@ -78,7 +78,7 @@ def make_fp8_linear_class() -> type: def forward(self, x: torch.Tensor) -> torch.Tensor: out_dim = self._fp8_weight.shape[0] original_shape = x.shape - if not _supports_fp8_compute(): + if not (x.is_cuda and _supports_fp8_compute()): # main's ``_apply_dequant`` fallback (pre-sm89 / CPU) — also # what CPU tests exercise. w_scale = self._fp8_weight_scale.to(x.dtype) diff --git a/fastvideo2/wan21/gates/capture_fastvideo_main.py b/fastvideo2/wan21/gates/capture_fastvideo_main.py index c2525181f..6981c1d69 100644 --- a/fastvideo2/wan21/gates/capture_fastvideo_main.py +++ b/fastvideo2/wan21/gates/capture_fastvideo_main.py @@ -81,19 +81,26 @@ def main() -> None: out = _out_dir() maybe_init_distributed_environment_and_model_parallel(1, 1) + # loader args mirror the released QAD example: no fsdp, no offload + _NO_WRAP = dict(use_fsdp_inference=False, dit_cpu_offload=False, + text_encoder_cpu_offload=False, vae_cpu_offload=False, + pin_cpu_memory=False) + # ------------------------------------------------------------- text --- # t5_args = FastVideoArgs(model_path=os.path.join(root, "text_encoder"), - pipeline_config=WanT2V480PConfig(), pin_cpu_memory=False) + pipeline_config=WanT2V480PConfig(), **_NO_WRAP) + t5_args.device = torch.device(device) t5 = TextEncoderLoader().load(os.path.join(root, "text_encoder"), t5_args).eval() - t5 = t5.to(device) # main default precision fp32 — no cast + t5_device = next(t5.parameters()).device # loader-owned placement tok = AutoTokenizer.from_pretrained(os.path.join(root, "tokenizer")) def encode(text: str) -> torch.Tensor: b = tok([text], padding="max_length", max_length=512, truncation=True, add_special_tokens=True, return_attention_mask=True, return_tensors="pt") with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None): - emb = t5(input_ids=b.input_ids.to(device), - attention_mask=b.attention_mask.to(device)).last_hidden_state + emb = t5(input_ids=b.input_ids.to(t5_device), + attention_mask=b.attention_mask.to(t5_device)).last_hidden_state + emb = emb.to(device) emb = emb.to(torch.float32) emb[:, int(b.attention_mask[0].sum()):] = 0 return emb @@ -131,7 +138,8 @@ def main() -> None: if quant: cfg.quant_config = FP8Config(granularity="tensor") args = FastVideoArgs(model_path=dit_path, - pipeline_config=PipelineConfig(dit_config=cfg, dit_precision="bf16")) + pipeline_config=PipelineConfig(dit_config=cfg, dit_precision="bf16"), + **_NO_WRAP) args.device = torch.device(device) dit = TransformerLoader().load(dit_path, args).to(dtype=torch.bfloat16).eval() if quant: @@ -192,7 +200,7 @@ def main() -> None: # ----------------------------------------------------------- decode --- # vae_args = FastVideoArgs(model_path=os.path.join(root, "vae"), - pipeline_config=WanT2V480PConfig()) + pipeline_config=WanT2V480PConfig(), **_NO_WRAP) vae_args.device = torch.device(device) vae = VAELoader().load(os.path.join(root, "vae"), vae_args).to( torch.device(device), torch.float32).eval() diff --git a/fastvideo2/wan21/gates/fastwan_anchor.py b/fastvideo2/wan21/gates/fastwan_anchor.py new file mode 100644 index 000000000..5bf0b5e0e --- /dev/null +++ b/fastvideo2/wan21/gates/fastwan_anchor.py @@ -0,0 +1,137 @@ +"""FastWan-QAD anchor — v2.1 vs fastvideo-main goldens, bitwise target. + +Consumes evidence/goldens/fastwan-qad-main/ (see capture_fastvideo_main.py) +and runs v2.1's OWN stack — no fastvideo import here: + + text_encoder our fp32 UMT5 path vs main's embeds (target 0.0*) + dit.bf16 WanModelFV probes, 3 timesteps (target 0.0) + dit.fp8 WanModelFVFP8 probes, 3 timesteps (target 0.0) + e2e.step{i}.x per-step latent chain through OUR engine (target 0.0) + e2e.final final latents (target 0.0) + +*ASCII prompts only. main's text cleaning (ftfy) diverges from official's on +CJK width-folding; the FastWan cards currently reuse the wan21 text stage. +# ponytail: ASCII-equal today; add main's clean fn to the FastWan pipeline +# (and a CJK golden) before serving non-ASCII prompts against these cards. + +Usage (cluster): + python -m fastvideo2.wan21.gates.fastwan_anchor [--skip-e2e] +""" +from __future__ import annotations + +import json +import os +import sys +from typing import Any + +import numpy as np + +GOLD = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), + "..", "..", "evidence", "goldens", "fastwan-qad-main")) +PROBE_TS = (1000, 757, 522) + + +def rel_l2(a: np.ndarray, b: np.ndarray) -> float: + denom = float(np.linalg.norm(b.astype(np.float64))) + if denom == 0.0: + return float(np.linalg.norm(a.astype(np.float64))) + return float(np.linalg.norm((a.astype(np.float64) - b.astype(np.float64)))) / denom + + +def run(skip_e2e: bool = False) -> int: + import torch + + from fastvideo2.engine import Instance, Request + from fastvideo2.engine import run as engine_run + from fastvideo2.registry import resolve + from fastvideo2.wan21.card import FASTWAN_QAD_FP8_1_3B + from fastvideo2.wan21.model_fv import WanModelFV, WanModelFVFP8 + + with open(os.path.join(GOLD, "manifest.json")) as f: + manifest = json.load(f) + device = "cuda" + rows: list[tuple[str, float]] = [] + + from fastvideo2.loading import resolve_weights + root = resolve_weights(FASTWAN_QAD_FP8_1_3B) + + # ------------------------------------------------------------- probes --- # + pi = np.load(os.path.join(GOLD, "probe_inputs.npz")) + px = torch.from_numpy(pi["latent"]) + pc = torch.from_numpy(pi["context"]) + + def probe(dit: Any, tag: str) -> None: + for t in PROBE_TS: + xt = px.to(device, torch.bfloat16) + ct = pc.to(device) # fp32 + tt = torch.tensor([t], dtype=torch.int64, device=device) + with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): + o = dit(xt, ct, tt) + gold = np.load(os.path.join(GOLD, f"dit_{tag}_t{t}.npz"))["out"] + rows.append((f"dit.{tag}.t{t}", rel_l2(o.to(torch.float32).cpu().numpy(), gold))) + + dit = WanModelFV.from_pretrained(root, torch_dtype=torch.bfloat16).to(device) + probe(dit, "bf16") + del dit + torch.cuda.empty_cache() + dit_q = WanModelFVFP8.from_pretrained(root, torch_dtype=torch.bfloat16).to(device) + probe(dit_q, "fp8") + del dit_q + torch.cuda.empty_cache() + + # --------------------------------------------------------------- text --- # + # through the served pipeline's own text path, so the gate measures what + # generate() actually feeds the loop + from fastvideo2.wan21.pipeline import _encode_one + instance = Instance(FASTWAN_QAD_FP8_1_3B, device=device) + tok = instance.component("tokenizer") + enc = instance.component("text_encoder") + gold_text = np.load(os.path.join(GOLD, "text_encoder.npz")) + for name, prompt in (("e2e", manifest["e2e_prompt"]), ("probe", manifest["probe_prompt"])): + emb = _encode_one(tok, enc, prompt) + rows.append((f"text.{name}", rel_l2(emb[0].to(torch.float32).cpu().numpy(), + gold_text[name]))) + + # ---------------------------------------------------------------- e2e --- # + if not skip_e2e: + card, builder = resolve(FASTWAN_QAD_FP8_1_3B.model_id) + req = Request(prompt=manifest["e2e_prompt"], seed=manifest["seed"], + num_frames=81, height=480, width=832, + capture_trajectory=True) + out = engine_run(instance, builder(), req) + den = out.outputs["latents"] # the denoise slot: {latents, trajectory, steps} + final = den["latents"].to(torch.float32).cpu().numpy() + gold_final = np.load(os.path.join(GOLD, "e2e_final_latents.npz"))["latents"] + rows.append(("e2e.final", rel_l2(final, gold_final))) + # per-step input chain: golden step_i.x is what main fed the DiT at + # step i (BCFHW — the hook sees the permuted input); our trajectory + # stores post-advance latents in BTCHW = the input of step i+1. + traj = [t.to(torch.float32).permute(0, 2, 1, 3, 4).numpy() for t in den["trajectory"]] + for i in range(1, len(PROBE_TS)): + g = np.load(os.path.join(GOLD, f"e2e_step{i}.npz"))["x"] + rows.append((f"e2e.step{i}.x", rel_l2(traj[i - 1], g))) + + # -------------------------------------------------------------- report --- # + print(f"\nFastWan-QAD anchor vs {manifest['fastvideo_commit'][:9]} " + f"({manifest['attention_backend']}, fp8 {'' if skip_e2e else '+ e2e'})") + failed = [] + for name, val in rows: + ok = val == 0.0 + if not ok: + failed.append(name) + print(f" {name:16s} rel_l2 {val:.6e} {'OK (bitwise)' if ok else 'DIFF'}") + + from fastvideo2.verify import append_ledger, env_fingerprint + append_ledger({"gate": "anchor.fastwan-qad-main", + "status": "pass" if not failed else "fail", + "model_id": FASTWAN_QAD_FP8_1_3B.model_id, + "card_digest": FASTWAN_QAD_FP8_1_3B.digest(), + "metrics": {n: v for n, v in rows}, + "tolerances": {"all": 0.0}, + "env": env_fingerprint(), + "detail": f"goldens {manifest['fastvideo_commit'][:9]}"}) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(run(skip_e2e="--skip-e2e" in sys.argv))