Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe
scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes.
This commit is contained in:
parent
bdb8eb9185
commit
e966a6d7ef
1 changed files with 100 additions and 0 deletions
100
scripts/fp8_overflow_check.py
Normal file
100
scripts/fp8_overflow_check.py
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Empirically check that fp8 dynamic quant with fast accumulation does not overflow.
|
||||
|
||||
Hooks every quantised Linear's output during a real Z-Image generation and reports the
|
||||
global max |output| and any non-finite (Inf/NaN) count, for use_fast_accum True vs False.
|
||||
The concern fast_accum raises is accumulation *precision*, not overflow (the accumulator
|
||||
stays FP32-range and torchao's dynamic per-row scale keeps FP8 inputs <= 448); this proves
|
||||
it on the real model, including Z-Image's large (~9e5) activation peaks. Run on one CUDA GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
BASE = "Tongyi-MAI/Z-Image-Turbo"
|
||||
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
|
||||
|
||||
|
||||
def _load_dense():
|
||||
import torch, diffusers
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder="transformer", torch_dtype=torch.bfloat16)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, transformer=t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _run(fast_accum, steps, res, seed, mf):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
|
||||
from torchao.float8 import Float8MMConfig
|
||||
|
||||
pipe = _load_dense()
|
||||
|
||||
def filt(mod, fqn=""):
|
||||
return (isinstance(mod, nn.Linear)
|
||||
and mod.in_features >= mf and mod.out_features >= mf)
|
||||
|
||||
quantize_(pipe.transformer,
|
||||
Float8DynamicActivationFloat8WeightConfig(mm_config=Float8MMConfig(use_fast_accum=fast_accum)),
|
||||
filter_fn=filt)
|
||||
|
||||
stats = {"max_abs": 0.0, "nonfinite": 0, "hooked": 0}
|
||||
|
||||
def hook(mod, inp, out):
|
||||
t = out[0] if isinstance(out, tuple) else out
|
||||
if not torch.is_tensor(t):
|
||||
return
|
||||
finite = torch.isfinite(t)
|
||||
nf = int((~finite).sum().item())
|
||||
stats["nonfinite"] += nf
|
||||
m = float(t[finite].abs().max().item()) if finite.any() else float("inf")
|
||||
if m > stats["max_abs"]:
|
||||
stats["max_abs"] = m
|
||||
|
||||
# Hook the quantised linears (where an fp8-accumulation overflow would surface).
|
||||
# Run EAGER: forward hooks don't trace through torch.compile, and the fp8 fast-accum
|
||||
# accumulation is identical compiled or eager -- compile only changes scheduling.
|
||||
for m in pipe.transformer.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
m.register_forward_hook(hook)
|
||||
stats["hooked"] += 1
|
||||
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
img = pipe(prompt=PROMPT, width=res, height=res, num_inference_steps=steps,
|
||||
guidance_scale=0.0, generator=g).images[0]
|
||||
import numpy as np
|
||||
arr = np.array(img)
|
||||
img_finite = bool(np.isfinite(arr).all())
|
||||
del pipe
|
||||
torch.cuda.empty_cache()
|
||||
return stats, img_finite
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type=int, default=4)
|
||||
p.add_argument("--res", type=int, default=512)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--min-feat", type=int, default=512)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
print(f"== fp8 overflow check (Z-Image dense, {args.res}px, {args.steps} steps) ==", flush=True)
|
||||
for fast in (True, False):
|
||||
stats, img_finite = _run(fast, args.steps, args.res, args.seed, args.min_feat)
|
||||
print(f" fast_accum={str(fast):5s} hooked_linears={stats['hooked']:3d} "
|
||||
f"max|linear_out|={stats['max_abs']:.1f} nonfinite_elems={stats['nonfinite']} "
|
||||
f"image_all_finite={img_finite}", flush=True)
|
||||
print("FP8-OVERFLOW-CHECK-DONE", flush=True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
|
||||
sys.exit(main())
|
||||
Loading…
Add table
Add a link
Reference in a new issue