transformers
4cb5d9cb - Quantize per module group; split the arms by producer, not by format

Commit
28 days ago
Quantize per module group; split the arms by producer, not by format Two axes that were tangled into one `quant_method`. GROUPS. A checkpoint can be more than one format — DeepSeek-V4 is W4A4 mxfp4 experts over block-FP8 linears — and a flat config cannot say so, which is why its expert format arrives on the MODEL config as `expert_dtype` and is read back out deep inside `FineGrainedExperts.__init__`. `FineGrainedGroup` names a format per module subset; `group_for(name)` resolves one; `replace_with_finegrained_layer` asks per module, so the side-channel read is gone. Nothing about this is DeepSeek-shaped. Verified across three models and three combinations, swap and in-flight quantization both: Qwen3-MoE with nvfp4 experts over mxfp8 dense, and GPT-OSS with nvfp4 attention, static block-FP8 experts and mxfp8 for the rest — three groups, the exotic format on the dense side, and differing activation SCHEMES per group. Backwards compatibility is normalisation, the same move `quant_algo` already gets: the flat fields become a single group, and a legacy `expert_dtype` becomes two (`split_out_experts`). Existing checkpoints load unchanged and nothing downstream learns which spelling it came from. `config_groups` is normalised too — GLM-5.2-NVFP4 ships it and we were dropping it, along with the `input_activations` and `dynamic: false` it carries. A group naming a format we do not serve falls back to the flat fields rather than guessing at a partial translation. Resolution is order-INDEPENDENT: `to_json_string` sorts keys, so a config that round-tripped through `config.json` has lost its declaration order. Targeted groups win over the catch-all, and two targeted groups claiming one module raises. ARMS. Split per PRODUCER key layout rather than per format, which is what actually differs: GPT-OSS's packed `_blocks`, modelopt's two-level scales, the calibrated `input_scale`. MXFP8 ships the plain `weight`/`weight_scale_inv` pair and gets no arm — the base is its handler. The base keeps no `_quant_method() == ...` test at all. Also: a finegrained module holding a full-precision weight now fails the load instead of silently computing in bf16 and handing back a model the caller believes is quantized. Checked once in the post-load pass, not per forward — the dtype cannot change between them. No shipping checkpoint relied on the old fallback: gpt-oss-120b lists attention in `modules_to_not_convert`, so those modules are never swapped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Parents
Loading