NVIDIA / NVIDIA/Model-Optimizer

NVFP4 quantized (shared) lm_head fails fake-quant in training forward: ScaledE4M3Function 'Only support E=4 & M=3'

Open
#1,875 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug investigating torch.quantization
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.