ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-06-26 07:38:23 +00:00
commit 244124d740
1 changed files with 36 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

57
scripts/fp8_overflow_check.py
View file

@ -22,9 +22,11 @@ PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly
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)
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
)
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
pipe.to("cuda")
return pipe
@ -37,13 +39,16 @@ def _run(fast_accum, steps, res, seed, mf):
pipe = _load_dense()
def filt(mod, fqn=""):
return (isinstance(mod, nn.Linear)
and mod.in_features >= mf and mod.out_features >= mf)
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)
quantize_(
pipe.transformer,
Float8DynamicActivationFloat8WeightConfig(
mm_config = Float8MMConfig(use_fast_accum = fast_accum)
),
filter_fn = filt,
)
stats = {"max_abs": 0.0, "nonfinite": 0, "hooked": 0}
@ -66,10 +71,17 @@ def _run(fast_accum, steps, res, seed, mf):
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]
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
@ -77,21 +89,24 @@ def _run(fast_accum, steps, res, seed, mf):
return stats, img_finite
def main(argv=None) -> int:
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)
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)
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)
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