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.

Out of curiosity, how was performance with the CPU offload?

Good question - two very different answers depending on the stack:

transformers + this FP8 checkpoint (what this report is about): ~0.02 tok/s.
Not a serving path. The fp8 kernels are CUDA-only, so for the layers that
device_map=auto puts on CPU we had to bounce every call to the GPU, re-uploading the
cpu-resident expert weights each forward (~45GB/token). We only used it to verify
correctness and to instrument routing / n-gram table access - fine for science,
useless for chat.

llama.cpp (PR ggml-org/llama.cpp#27742) + unsloth UD-Q4_K_XL GGUF: 74.9 tok/s decode,
~2900 tok/s prefill (pp4096, llama-bench)
on the same single RTX PRO 6000 (96GB) with
plain dual-channel DDR4:

config VRAM decode
all experts on GPU, 51B n-gram table in host RAM 80 GB 74.9 tok/s
VRAM capped to 32GB (5090-sized) --n-cpu-moe 36 ~28 GB 24.2 tok/s

The key property is the one the model card advertises: the n-gram table costs a few KB of
row gathers per token, so keeping it CPU-side is essentially free - "CPU offload" of the
PLE table is not a compromise at all. It's offloading the experts that hurts (see the
32GB row), and even that stays usable.

Setup notes + more numbers: https://github.com/01554/llama.cpp/blob/expert-tier/EXPERT_TIER.md

Cheers, and thanks for sharing your work! This n00b greatly appreciates it 🙂

Sign up or log in to comment