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 🙂

Qwen org
Qwen org

currently, the transformers implementation is not fully optimized (especially prefill), so we recommend using other repos e.g. llama.cpp / mlx-vlm

@JJJYmmm

Thanks for the PR — that's a much cleaner fix for issue 2 than my gather-time patch ❤️

Issue 1 still reproduces on today's main (@42ca97014c85). The trigger is the loading class: AutoModelForCausalLM instantiates modules as model.layers.*, while the exclusion list has model.language_model.layers.*, so the layer-path patterns never match (loading via the multimodal wrapper class matches fine — likely why it didn't repro on your side).

import torch
from transformers import AutoConfig, AutoModelForCausalLM, FineGrainedFP8Config
from transformers.integrations.finegrained_fp8 import replace_with_fp8_linear
cfg = AutoConfig.from_pretrained("Qwen/Qwen3.8-Flash-Next-FP8")
mods = cfg.quantization_config["modules_to_not_convert"]
delattr(cfg, "quantization_config")
with torch.device("meta"):
    m = AutoModelForCausalLM.from_config(cfg)
m = replace_with_fp8_linear(m, modules_to_not_convert=mods,
      quantization_config=FineGrainedFP8Config(modules_to_not_convert=mods))
print(type(dict(m.named_modules())["model.layers.1.ple.key_proj"]))
# -> FP8Linear  (checkpoint stores this tensor in BF16, so it silently mis-converts)

Agreed on llama.cpp for serving 🙂 That said, I do think this is worth fixing — it corrupts silently, and transformers is what people validate against.

Correction to my comment above — my repro was flawed. I called replace_with_fp8_linear directly, which bypasses _normalize_modules_to_not_convert: on current main the quantizer normalizes the wrapper prefixes before conversion, so through the real from_pretrained path issue 1 does not reproduce anymore. You were right, and credit to 22elix3r's detailed analysis on huggingface/transformers#48349. Sorry for the noise — both issues addressed 🙂

hellohazime changed discussion status to closed

Sign up or log in to comment