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

Tighten comments in the image stack tests and scripts

Browse source
This commit is contained in:
Daniel Han 2026-07-12 12:21:14 +00:00
commit 4cc35aa87a
11 changed files with 290 additions and 312 deletions
Download patch file Download diff file Expand all files Collapse all files

25
scripts/build_prequant_checkpoint.py
View file

@ -83,17 +83,16 @@ def main(argv = None) -> int:
args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token
).to("cuda")
print(f" quantising in place ({scheme}) ...", flush = True)
# Mirror the runtime path EXACTLY (the offline == runtime, LPIPS-0 invariant): for int8 also
# skip the M=1 AdaLN-modulation / conditioning-embedder projections, else the saved checkpoint
# bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on
# Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns ().
# Mirror the runtime path EXACTLY (offline == runtime, LPIPS-0 invariant): for int8 also skip
# the M=1 AdaLN-modulation / conditioning-embedder projections, else the checkpoint bakes them
# as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. fp8
# / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns ().
exclude_name_tokens = exclude_tokens_for_scheme(scheme)
# fp8 and mxfp8 assert a bf16 weight, so their filter must skip any non-bf16 Linear the
# transformer keeps: a mixed-precision DiT (Wan / Hunyuan) retains its _keep_in_fp32_modules in
# fp32 even under torch_dtype=bf16, so quantising one would raise inside quantize_ and abort the
# whole pass. nvfp4 quantises fp32 fine, so it is not gated. Runtime quantize_transformer gates
# this on scheme membership; mirror it here so the offline checkpoint quantises the exact same
# layer set (offline == runtime).
# transformer keeps: a mixed-precision DiT (Wan / Hunyuan) keeps its _keep_in_fp32_modules in
# fp32 even under torch_dtype=bf16, so quantising one raises inside quantize_ and aborts the
# pass. nvfp4 quantises fp32 fine, so it isn't gated. Runtime quantize_transformer gates on
# scheme membership; mirror it here so the offline checkpoint quantises the same layer set.
require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
# fp8 bakes the accumulate mode into the saved kernels; record the resolved choice so the
# loader can refuse a checkpoint whose baked value contradicts an explicit runtime request.
@ -118,10 +117,10 @@ def main(argv = None) -> int:
"family": fam.name,
"scheme": scheme,
"min_features": args.min_features,
# The layers skipped for this scheme (int8's M=1 modulation projections; () for
# the scaled_mm schemes), whether non-bf16 Linears were skipped (the scaled_mm
# bf16 gate), and, for fp8, the baked accumulate mode. All let the loader reject a
# checkpoint that would not match the runtime path.
# The layers skipped for this scheme (int8's M=1 modulation projections; () for the
# scaled_mm schemes), whether non-bf16 Linears were skipped (the scaled_mm bf16 gate), and,
# for fp8, the baked accumulate mode. All let the loader reject a checkpoint that wouldn't
# match the runtime path.
"exclude_name_tokens": list(exclude_name_tokens),
"require_bf16": require_bf16,
"fast_accum": fast_accum,