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

Studio diffusion (Phase 14) review round 2: align helper name with the stack

Browse source 
Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.
This commit is contained in:
Daniel Han 2026-06-29 11:03:37 +00:00
commit d127be20a5
3 changed files with 19 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

12
scripts/build_prequant_checkpoint.py
View file

@ -57,7 +57,7 @@ def main(argv = None) -> int:
from core.inference.diffusion_transformer_quant import (
TQ_SCHEMES,
_make_quant_config,
int8_exclude_name_tokens,
exclude_tokens_for_scheme,
make_filter_fn,
)
from torchao.quantization import quantize_
@ -79,15 +79,15 @@ 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)
# Use the SAME int8 M=1 exclusion as the runtime quantiser (single source of truth):
# otherwise an int8 prequant checkpoint quantises the AdaLN modulation / conditioning
# embedders and reintroduces the torch._int_mm M=1 crash when loaded via
# transformer_prequant_path. fp8/fp4/mx get an empty exclusion (artifacts unchanged).
# 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 ().
quantize_(
transformer,
_make_quant_config(scheme),
filter_fn = make_filter_fn(
args.min_features, exclude_name_tokens = int8_exclude_name_tokens(scheme)
args.min_features, exclude_name_tokens = exclude_tokens_for_scheme(scheme)
),
)