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

Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder

Browse source 
Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.

Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).

Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.
This commit is contained in:
Daniel Han 2026-06-29 10:47:47 +00:00
commit 7098f1b363
3 changed files with 36 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

13
scripts/build_prequant_checkpoint.py
View file

@ -57,6 +57,7 @@ def main(argv = None) -> int:
from core.inference.diffusion_transformer_quant import (
TQ_SCHEMES,
_make_quant_config,
int8_exclude_name_tokens,
make_filter_fn,
)
from torchao.quantization import quantize_
@ -78,7 +79,17 @@ 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)
quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(args.min_features))
# 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).
quantize_(
transformer,
_make_quant_config(scheme),
filter_fn = make_filter_fn(
args.min_features, exclude_name_tokens = int8_exclude_name_tokens(scheme)
),
)
# Move the state dict to CPU for a portable, GPU-free artifact.
state_dict = {