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

[yeyu-nvidia]

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)