NVIDIA / NVIDIA/Model-Optimizer
NVFP4 quantized (shared) lm_head fails fake-quant in training forward: ScaledE4M3Function 'Only support E=4 & M=3'
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.8k
- Forks
- 604
- Avg merge
- 2d 8h
- Merged PRs (30d)
- 142
Description
Summary
When an NVFP4-quantized output_layer/lm_head (weight quantizer num_bits: e2m1, block scale e4m3) is exercised in the training forward (not just PTQ/export), _fake_quantize routes into ScaledE4M3Function.forward (modelopt/torch/quantization/tensor_quant.py), which raises NotImplementedError("Only support E=4 & M=3 for now.") because the NVFP4 weight format is E2M1, not E4M3.
Repro context
QAD (KD) on a hybrid-MoE model that has an MTP head, with the MTP loss enabled (#1805). The main LM head is skipped via skip_lm_loss, so PTQ/export never hit this path — but the MTP head shares the quantized output_layer and calls it during the MTP loss, triggering fake-quant in the training forward. Traceback: process_mtp_loss -> output_layer -> quant_linear -> tensor_quantizer._fake_quantize -> tensor_quant.py:436.
Impact / workaround
Blocks keeping lm_head/output_layer quantized (NVFP4) when the MTP head is involved during QAD. Workaround: keep lm_head/output_layer in BF16. A proper fix would let the NVFP4 weight quantizer fake-quant correctly in the training forward instead of dispatching to the E4M3-only path.
Verified on modelopt main 973cb09cb. cc @AAnoosheh (#1805)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read modelopt/torch/quantization/tensor_quant.py around line 436, then trace _fake_quantize from quant_linear during process_mtp_loss. Reproduce the QAD/MTP training-forward path described in the issue; done means an NVFP4 output_layer can remain quantized without the E4M3-only NotImplementedError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100