[Quantization] Support NVFP4-quantized lm_head and embed_tokens via modelopt

[lucaspirola]

Summary

Add support for serving models where lm_head and/or embed_tokens are quantized to NVFP4 by nvidia-modelopt.

lm_head support:

• ParallelLMHead extends VocabParallelEmbedding (not LinearBase), so ModelOptQuantConfigBase.get_quant_method() previously returned None for it, falling back to UnquantizedEmbeddingMethod. This caused weight loading to fail with ValueError: There is no module or parameter named 'lm_head.input_scale' when NVFP4 scale tensors were present in the checkpoint.
• Add ModelOptNvFp4LMHeadMethod (extends ModelOptNvFp4LinearMethod) with a custom weight_loader that handles PerTensorScaleParameter scalars and vocab-parallel sharding for NVFP4 packed weights on VocabParallelEmbedding-based layers.
• When lm_head is excluded from quantization (the default), behavior is unchanged.

embed_tokens support:

• Add ModelOptNvFp4EmbeddingMethod that stores embeddings in NVFP4 packed format and dequantizes per-row during lookup via an FP4 E2M1 lookup table (no GEMM kernel needed).
• Enabled by "nvfp4_embed": true in hf_quant_config.json or config.json quantization_config.
• For a 131K-vocab x 5120-dim embedding, reduces storage from 1280 MB (BF16) to ~360 MB (NVFP4).
• Config parsing handles both hf_quant_config.json format ({"quantization": {"nvfp4_embed": true}}) and compressed-tensors format ({"nvfp4_embed": true}).

Motivation: On VRAM-constrained GPUs (e.g. 16 GB), quantizing both lm_head and embed_tokens to NVFP4 saves ~1880 MB total for a 131K-vocab model. On RTX 5080 with Devstral-24B, this enables 32K tokens context (up from 20K with BF16 embeddings, +57.9%).

How to produce an NVFP4 checkpoint

lm_head - include in modelopt quantization:

import copy
import modelopt.torch.quantization as mtq

nvfp4_cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG)
del nvfp4_cfg["quant_cfg"]["*lm_head*"]  # include lm_head in quantization
model = mtq.quantize(model, nvfp4_cfg, forward_loop)

embed_tokens - quantize offline after modelopt export:

# convert_embed_nvfp4.py quantizes embed_tokens to NVFP4 packed format
# and sets "nvfp4_embed": true in hf_quant_config.json
python convert_embed_nvfp4.py <src_model_dir> <dst_model_dir>

Test plan

• Tested with Devstral-24B quantized to NVFP4 (lm_head + embed_tokens) on RTX 5080 16 GB
• Quality verified: correct math (2+2=4), valid JSON, factual recall (capital of France)
• Binary-searched max context: 2022 blocks = 32,352 tokens (2.468 GiB KV cache)
• Backward compatible: models without nvfp4_embed flag are unaffected
• Tensor-parallel sharding handled via VocabParallelEmbedding.shard_indices
• All pre-commit checks pass (ruff, mypy, typos, DCO)
• nvfp4_embed config works with both hf_quant_config.json and config.json formats

🤖 Generated with Claude Code

[mergify]

Hi @lucaspirola, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[gemini-code-assist]

Code Review

This pull request introduces support for NVFP4 quantization for the lm_head module. This is achieved by adding a specific quantization method, ModelOptNvFp4LMHeadMethod, for ParallelLMHead layers. This method includes a custom weight loader to correctly handle vocabulary-parallel sharding and scalar parameters associated with NVFP4 quantization. The changes appear to be correct and well-integrated with the existing quantization framework.

[mergify]

Hi @lucaspirola, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @lucaspirola.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @lucaspirola.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!