FP8 config: modules_to_not_convert paths never match on text-only loads (silent quality issue) + fp8 PLE table lacks a dequant path

#2
by hellohazime - opened

Thanks for the release! Reporting two day-0 issues we hit while loading
Qwen/Qwen3.8-Flash-Next-FP8 with transformers main (5.16.0.dev0, Qwen4Exp support from
huggingface/transformers#48337) on a single RTX PRO 6000 (96GB) + CPU offload.

1. quantization_config.modules_to_not_convert uses multimodal wrapper paths

The list entries are all of the form model.language_model.layers..., but when the
checkpoint is loaded through AutoModelForCausalLM, modules are named model.layers....
None of the exclusion patterns match, so the FineGrainedFP8 quantizer converts modules that
the config intends to keep unquantized (e.g. ple.key_proj / ple.value_proj,
hyper-connection mixers, linear_attn.conv1d / in_proj_*). Their BF16 checkpoint weights
are then cast to fp8 without scales — no load-time error, just silently degraded
numerics. (We verified the checkpoint side: e.g.
model.language_model.layers.1.ple.key_proj.weight is stored BF16 [10240, 2560].)

Workaround we used:

qc = json.load(open(f"{MODEL}/config.json"))["quantization_config"]
mods = sorted({m.replace("model.language_model.", "model.") for m in qc["modules_to_not_convert"]})
quant = FineGrainedFP8Config(activation_scheme=qc["activation_scheme"],
                             weight_block_size=tuple(qc["weight_block_size"]),
                             modules_to_not_convert=mods)
model = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=quant, ...)

A checkpoint-side fix could be to also list the text-only paths (or suffix patterns).

2. fp8 n-gram (PLE) table rows are gathered without dequantization

ple.ple_embedding.ngram_embedding (320,001,536 x 160) is stored fp8 with a per-tensor
weight_scale ([1], BF16). Since the module is a plain nn.Embedding, transformers drops
weight_scale as an unexpected key and Qwen4ExpTextNGramEmbedding.forward returns raw fp8
rows, which crash (or worse, silently mis-scale, combined with issue 1) at the first
consumer:
RuntimeError: expected m1 and m2 to have the same dtype: Float8_e4m3fn != BFloat16.

We patched the gather to .to(bfloat16) * weight_scale (scale read directly from the
shard) and the model runs. This one is probably a transformers-side fix (we are filing the
corresponding issues there and can link back), but flagging it here since the resolution
may involve the checkpoint layout (e.g. storing the table scale where the loader keeps it).

Happy to provide full repro scripts / patches. Box: RTX PRO 6000 Blackwell 96GB, 128GB RAM,
partial CPU offload; with the two fixes (plus small quantizer guards reported to
transformers) the model loads and generates.

Sign up or log in to comment