36df317293
Comment-only pass over the Python this PR touches: drop what the code already says, collapse multi-line explanations that still read on one line, and keep the reasoning that is not recoverable from the code. No code, docstring semantics or behaviour changes; verified with an AST comparison against the previous revision, and the backend suite is unchanged (same 37 environment failures as before: the API integration tests that need a live keyed server, the flash-attn install hooks, and the GPU memory fields).
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
# 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 surfaces. Run eager: forward
|
|
# hooks don't trace through torch.compile, and accumulation is identical either way.
|
|
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())
|