222 lines
9.0 KiB
Python
222 lines
9.0 KiB
Python
"""QAD FP4 quality A/B/C/D harness — Wan2.1-T2V-1.3B on sm_121 (DGX Spark GB10).
|
|
|
|
Tests whether the quantization-aware-distilled checkpoint
|
|
``FastVideo/FastWan-QAD-1.3B`` recovers FP4 quality on sm_121, using the
|
|
sm_121-enabled ``ATTN_QAT_INFER`` attention kernel. Running FP4 attention on
|
|
*stock* Wan weights gives output below bf16 — expected, because stock weights
|
|
were never trained to tolerate FP4 attention. This harness runs the checkpoint
|
|
that *was* (fake-quant FP4 attention + NVFP4 linear trained into the model).
|
|
|
|
Attention (bf16 vs FP4) and linear (bf16 vs FP4) are fully decoupled at
|
|
inference, so we can isolate each axis:
|
|
|
|
FASTVIDEO_ATTENTION_BACKEND unset -> bf16 attention (SDPA)
|
|
FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER -> FP4 attention
|
|
QAD_LINEAR=0 -> bf16 linear
|
|
QAD_LINEAR=1 -> NVFP4 FP4 linear (flashinfer)
|
|
|
|
arm attention linear selects with
|
|
A bf16 bf16 QAD_LINEAR=0 (no ATTN env) reference
|
|
B FP4 bf16 QAD_LINEAR=0 ATTN_QAT_INFER isolate attn
|
|
C bf16 FP4 QAD_LINEAR=1 (no ATTN env) isolate linear
|
|
D FP4 FP4 QAD_LINEAR=1 ATTN_QAT_INFER full 4-bit
|
|
|
|
All four arms run end-to-end on the GB10 (the full 4-bit path — FP4 linear +
|
|
FP4 attention — works on sm_121). This script still runs exactly ONE arm per
|
|
invocation and dumps a C stack on any hard crash, so a single misbehaving arm
|
|
can never take the others down with it; the runbook loops it four times with
|
|
different env. Quality is the eye/ear on the saved mp4 + a matching-frame still;
|
|
timing is the mean generation_time (full pipeline: text-encode +
|
|
denoise + VAE decode) over the measured runs.
|
|
|
|
On the GB10, expect FP4 attention ~6% faster end-to-end generation vs bf16 and quality-neutral
|
|
by eye on the QAD checkpoint (both share the 3-step distill's quality ceiling).
|
|
FP4 *linear* is roughly break-even at 1.3B/480p (the per-call quantize overhead
|
|
eats the small-GEMM saving in eager mode); its win shows at higher resolution
|
|
with torch.compile.
|
|
|
|
Knobs (env): QAD_MODEL, QAD_DISTILLED, QAD_STEPS (3), QAD_GUIDANCE (1.0),
|
|
QAD_HEIGHT (480), QAD_WIDTH (832), QAD_FRAMES (77), QAD_SEED (42),
|
|
QAD_WARMUP (1), QAD_RUNS (3), QAD_STILL (20), QAD_OUT (qad_fp4_samples).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import faulthandler
|
|
import glob
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
|
|
faulthandler.enable() # dump a C stack if any arm hard-crashes.
|
|
|
|
|
|
def _env(name: str, default: str) -> str:
|
|
return os.environ.get(name, default)
|
|
|
|
|
|
def _env_int(name: str, default: int) -> int:
|
|
return int(os.environ.get(name, str(default)))
|
|
|
|
|
|
def _env_float(name: str, default: float) -> float:
|
|
return float(os.environ.get(name, str(default)))
|
|
|
|
|
|
# Distilled QAD transformer, loaded on top of the base Wan2.1-1.3B pipeline.
|
|
DEFAULT_DISTILLED = "FastVideo/FastWan-QAD-1.3B"
|
|
# The repo is a full diffusers pipeline; we overlay only its transformer onto the
|
|
# base Wan pipeline (vae/text_encoder are Wan-identical).
|
|
DISTILLED_WEIGHTS_FILE = "transformer/diffusion_pytorch_model.safetensors"
|
|
|
|
PROMPT = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
|
|
"wide with interest. The playful yet serene atmosphere is complemented by soft "
|
|
"natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
|
|
|
|
|
|
def resolve_distilled_weights(hf_id: str) -> str:
|
|
"""Return a local path to the distilled transformer safetensors."""
|
|
if not hf_id:
|
|
return ""
|
|
if os.path.exists(hf_id):
|
|
return hf_id
|
|
from huggingface_hub import hf_hub_download
|
|
return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE)
|
|
|
|
|
|
def build_generator(fp4_linear: bool):
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.configs.pipelines.base import PipelineConfig
|
|
|
|
model_id = _env("QAD_MODEL", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
|
|
|
pipeline_config = PipelineConfig.from_pretrained(model_id)
|
|
pipeline_config.dit_precision = "bf16"
|
|
pipeline_config.vae_precision = "bf16"
|
|
pipeline_config.text_encoder_precisions = ("bf16", )
|
|
|
|
if fp4_linear:
|
|
# Wan-style config: matches to_q/k/v/out + ffn (the plain NVFP4 config is
|
|
# LTX2-specific and would quantize nothing on Wan).
|
|
from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig
|
|
pipeline_config.dit_config.quant_config = NVFP4QATConfig()
|
|
|
|
extra_kwargs = {}
|
|
distilled = resolve_distilled_weights(_env("QAD_DISTILLED", DEFAULT_DISTILLED))
|
|
if distilled:
|
|
print(f"[qad] distilled weights: {distilled}")
|
|
extra_kwargs["init_weights_from_safetensors"] = distilled
|
|
|
|
# Keep everything resident (the 1.3B QAD model + FP4 fits the GB10's unified
|
|
# memory); real Wan VAE decode for a faithful quality read (no TAEHV).
|
|
return VideoGenerator.from_pretrained(
|
|
model_id,
|
|
pipeline_config=pipeline_config,
|
|
num_gpus=1,
|
|
use_fsdp_inference=False,
|
|
dit_cpu_offload=False,
|
|
dit_layerwise_offload=False,
|
|
vae_cpu_offload=False,
|
|
text_encoder_cpu_offload=False,
|
|
pin_cpu_memory=False,
|
|
enable_torch_compile=False, # eager: isolate the FP4 effect, no compile noise
|
|
**extra_kwargs,
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
if not torch.cuda.is_available():
|
|
raise SystemExit("CUDA is required.")
|
|
|
|
fp4_linear = _env_int("QAD_LINEAR", 0) == 1
|
|
attn = _env("FASTVIDEO_ATTENTION_BACKEND", "default")
|
|
steps = _env_int("QAD_STEPS", 3)
|
|
guidance = _env_float("QAD_GUIDANCE", 1.0)
|
|
height = _env_int("QAD_HEIGHT", 480)
|
|
width = _env_int("QAD_WIDTH", 832)
|
|
frames = _env_int("QAD_FRAMES", 77)
|
|
seed = _env_int("QAD_SEED", 42)
|
|
warmup = _env_int("QAD_WARMUP", 1)
|
|
runs = _env_int("QAD_RUNS", 3)
|
|
still_idx = _env_int("QAD_STILL", 20)
|
|
out_dir = _env("QAD_OUT", "qad_fp4_samples")
|
|
prompt = _env("QAD_PROMPT", PROMPT)
|
|
|
|
tag = f"lin-{'fp4' if fp4_linear else 'bf16'}_attn-{attn.lower()}"
|
|
cap = torch.cuda.get_device_capability()
|
|
print(f"[qad] GPU {torch.cuda.get_device_name()} (cc {cap[0]}.{cap[1]})")
|
|
print(f"[qad] ARM {tag}: linear={'FP4' if fp4_linear else 'bf16'}, "
|
|
f"attention={attn}, {steps} steps, guidance {guidance}, "
|
|
f"{height}x{width}x{frames}, seed {seed}")
|
|
print(f"[qad] prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}")
|
|
|
|
arm_dir = os.path.join(out_dir, tag)
|
|
os.makedirs(arm_dir, exist_ok=True)
|
|
generator = build_generator(fp4_linear)
|
|
|
|
def _generate():
|
|
# seed + frame dims live under `sampling` (SamplingConfig); `output`
|
|
# only takes output_path/save_video/return_frames (OutputConfig).
|
|
return generator.generate(
|
|
request={
|
|
"prompt": prompt,
|
|
"sampling": {
|
|
"seed": seed,
|
|
"num_inference_steps": steps,
|
|
"guidance_scale": guidance,
|
|
"height": height,
|
|
"width": width,
|
|
"num_frames": frames,
|
|
},
|
|
"output": {
|
|
"save_video": True,
|
|
"output_path": arm_dir,
|
|
"return_frames": True
|
|
},
|
|
})
|
|
|
|
for _ in range(warmup):
|
|
_generate()
|
|
|
|
denoise_times: list[float] = []
|
|
last = None
|
|
for i in range(runs):
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
last = _generate()
|
|
torch.cuda.synchronize()
|
|
wall = time.perf_counter() - t0
|
|
# generate_video returns a plain dict; attribute access would always
|
|
# fall back to wall time.
|
|
gen_t = last.get("generation_time") if isinstance(last, dict) else None
|
|
denoise_times.append(gen_t if gen_t is not None else wall)
|
|
print(f"[qad] {tag} run {i + 1}/{runs}: {wall:.2f}s wall "
|
|
f"(gen {denoise_times[-1]:.2f}s)")
|
|
|
|
# The pipeline wrote the mp4 (full known-good encode) into arm_dir; report
|
|
# it and pull a matching-frame still from the [b,c,t,h,w] samples tensor
|
|
# using the same recipe the pipeline's frame builder uses.
|
|
mp4s = sorted(glob.glob(os.path.join(arm_dir, "*.mp4")), key=os.path.getmtime)
|
|
if mp4s:
|
|
print(f"[qad] video: {mp4s[-1]}")
|
|
samples = getattr(last, "samples", None) if last is not None else None
|
|
if samples is not None and getattr(samples, "ndim", 0) == 5:
|
|
import imageio
|
|
f = min(still_idx, samples.shape[2] - 1) # samples: [b, c, t, h, w]
|
|
still = (samples[0, :, f].permute(1, 2, 0).clamp(0, 1) * 255)
|
|
still = still.to(torch.uint8).cpu().numpy()
|
|
png = os.path.join(arm_dir, f"raccoon_{tag}_f{f}.png")
|
|
imageio.imwrite(png, still)
|
|
print(f"[qad] still: {png}")
|
|
else:
|
|
print("[qad] note: no 5-D samples tensor; grab a frame from the mp4 above")
|
|
|
|
mean = sum(denoise_times) / len(denoise_times)
|
|
print(f"\n[qad][{tag}] generation mean {mean:.2f}s over {runs} runs "
|
|
f"({warmup} warmup, {steps} steps)")
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|