[OMNIML-4730] Support quantized nn.Embedding

CHANGELOG.rst

@@ -27,6 +27,7 @@ Changelog
 - Add NVFP4 W4A16 weight-only quantization (``w4a16_nvfp4``): FP4 weights with group_size=16, BF16 activations, no calibration forward pass required. Use ``mtq.W4A16_NVFP4_CFG`` or ``--qformat w4a16_nvfp4`` in ``hf_ptq.py``. vLLM deployment support is in progress.
 - Add ``DATASET_COMBOS`` to ``modelopt.torch.utils.dataset_utils`` — single ``--dataset`` tokens that fan out to multiple registered datasets; per-entry ``num_samples`` is split evenly across the members. Initial combos: ``cnn_nemotron_v2_mix`` (``cnn_dailymail`` + ``nemotron-post-training-dataset-v2``, used by ``hf_ptq.py`` when no ``--dataset`` is provided) and ``nemotron-post-training-v3`` (the seven ``nvidia/Nemotron-*`` SFT datasets added in #1498, mirroring the `nemotron-post-training-v3 collection <https://huggingface.co/collections/nvidia/nemotron-post-training-v3>`_). Combo names are listed by ``get_supported_datasets()`` and surfaced in ``--dataset`` help. ``get_dataset_dataloader`` rejects inputs that mix a combo with one of its member datasets (e.g. ``cnn_dailymail,cnn_nemotron_v2_mix``) to avoid double-sampling, and ``get_dataset_samples`` rejects combo names so callers route through the dataloader. ``hf_ptq.py`` default ``--calib_size`` is bumped from ``512`` to ``1024`` so the total calibration sample count under the new default combo matches the previous two-dataset fallback.
 - The ``nemotron-sft-agentic-v2`` registered dataset (added in #1498) now uses only the ``search`` split. The previously configured ``interactive_agent`` and ``tool_calling`` splits contain content-level defects (heterogeneous schema and a malformed JSON row, respectively) that cause pyarrow's streaming JSON reader to fail deterministically.
+- Add quantized ``nn.Embedding`` support. ``nn.Embedding`` is now registered in ``QuantModuleRegistry`` and exposes ``weight_quantizer`` (embedding table), ``output_quantizer`` (lookup activations), and a permanently disabled ``input_quantizer`` placeholder — embedding inputs are integer indices and cannot be fake-quantized, so direct ``enable*()`` calls raise. ``export_hf_checkpoint`` packs quantized embedding weights alongside Linear layers. Embedding quantizers are opt-in (``parent_class: nn.Embedding`` disabled by default).
 
 **Bug Fixes**
 

modelopt/torch/export/unified_export_hf.py

@@ -679,6 +679,34 @@ def _process_quantized_modules(
                     raise AssertionError(
                         f"Failed to export module '{name}' (type={type(sub_module).__name__}): {e}"
                     ) from e
+            elif isinstance(sub_module, nn.Embedding) and hasattr(sub_module, "weight_quantizer"):
+                # Quantized nn.Embedding: pack the embedding table the same way as Linear
+                # weights so downstream loaders see the NVFP4/FP8/INT-packed bytes + scales.
+                # Skip packing when the embedding's weight is tied to another module
+                # (e.g. tied_word_embeddings → lm_head): _export_quantized_weight reassigns
+                # the .weight attribute to a new uint8 Parameter, which severs the Python-
+                # level tie and leaves the other module pointing at a stale float Parameter.
+                tied_to = [
+                    other_name
+                    for other_name, other_module in model.named_modules()
+                    if other_module is not sub_module
+                    and getattr(other_module, "weight", None) is sub_module.weight
+                ]
+                if tied_to:
+                    warnings.warn(
+                        f"Skipping quantized weight packing for embedding '{name}': its "
+                        f"weight Parameter is shared with {tied_to} (weight tying). Packing "
+                        "would break the tie and produce stale weights in the tied module(s). "
+                        "The embedding will be exported as its fake-quantized float weight."
+                    )
+                else:
+                    try:
+                        with fsdp2_aware_weight_update(model, sub_module, reshard=False):
+                            _export_quantized_weight(sub_module, dtype)
+                    except AssertionError as e:
+                        raise AssertionError(
+                            f"Failed to export embedding '{name}' (type={type(sub_module).__name__}): {e}"
+                        ) from e
             elif (
                 "Llama4TextExperts" in type(sub_module).__name__
                 or "GptOssExperts" in type(sub_module).__name__

modelopt/torch/quantization/nn/__init__.py

@@ -18,6 +18,7 @@
 from .modules.quant_activations import *
 from .modules.quant_batchnorm import *
 from .modules.quant_conv import *
+from .modules.quant_embedding import *
 from .modules.quant_instancenorm import *
 from .modules.quant_layernorm import *
 from .modules.quant_linear import *

modelopt/torch/quantization/nn/modules/quant_embedding.py

 # TODO: Add a registry for TensorQuantizers and avoid this manual conversion. 
 def test_save_restore(skip_on_windows, model_cls, quant_config):  # Flaky on Windows 
 # TODO: Add a registry for TensorQuantizers and avoid this manual conversion. 
@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Quantized Embedding."""
+
+import contextlib
+
+import torch
+import torch.nn as nn
+
+from ...tensor_quant import QUANT_DESC_8BIT_PER_TENSOR
+from ...utils import is_torch_export_mode
+from .quant_module import QuantModule, QuantModuleRegistry
+from .tensor_quantizer import HardDisabledTensorQuantizer, SequentialQuantizer, TensorQuantizer
+
+__all__ = ["QuantEmbedding"]
+
+
+@QuantModuleRegistry.register({nn.Embedding: "nn.Embedding"})
+class _QuantEmbedding(QuantModule):
+    """Quantized version of ``nn.Embedding``.
+
+    The literal input to ``nn.Embedding`` is integer indices, which cannot be
+    fake-quantized. The ``input_quantizer`` attribute is kept (for symmetry with
+    other quant modules and for introspection by export/calibration code) but is
+    a :class:`HardDisabledTensorQuantizer` — direct ``enable*()`` calls raise and
+    wildcard configs are silently absorbed-then-force-disabled. Only the embedding
+    table (weight) and the lookup output (an activation feeding downstream layers)
+    are quantizable.
+
+    Quantizer roles:
+        - ``weight_quantizer``: quantizes the embedding table (``self.weight``).
+        - ``input_quantizer``: permanently disabled placeholder
+          (:class:`HardDisabledTensorQuantizer`).
+        - ``output_quantizer``: optional activation quantizer for the lookup output,
+          disabled by default.
+    """
+
+    weight_quantizer: TensorQuantizer | SequentialQuantizer
+    input_quantizer: HardDisabledTensorQuantizer
+    output_quantizer: TensorQuantizer
+    _enable_weight_quantization: bool
+    default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_output = QUANT_DESC_8BIT_PER_TENSOR
+
+    @contextlib.contextmanager
+    def quantize_weight(self):
+        """Context in which ``self.weight`` is quantized via the dynamic attribute."""
+        self._enable_weight_quantization = True
+        try:
+            yield
+        finally:
+            self._enable_weight_quantization = False
+
+    @staticmethod
+    def _get_quantized_weight(module: "_QuantEmbedding", weight: torch.Tensor) -> torch.Tensor:
+        if module._enable_weight_quantization or is_torch_export_mode():
+            return module.weight_quantizer(weight)
+        return weight
+
+    def _setup(self):
+        """Register weight, (locked) input, and output quantizers."""
+        self._register_temp_attribute(
+            "weight_quantizer", TensorQuantizer(self.default_quant_desc_weight)
+        )
+        # Build the input quantizer disabled. HardDisabledTensorQuantizer's mutators raise,
+        # so we disable it once at construction via direct attribute assignment.
+        input_quantizer = HardDisabledTensorQuantizer(self.default_quant_desc_input)
+        input_quantizer._disabled = True
+        self._register_temp_attribute("input_quantizer", input_quantizer)
+        self._register_temp_attribute(
+            "output_quantizer", TensorQuantizer(self.default_quant_desc_output)
+        )
+        self.output_quantizer.disable()
+        self._register_temp_attribute("_enable_weight_quantization", False)
+        self._register_dynamic_attribute("weight", self._get_quantized_weight)
+
+    def forward(self, input, *args, **kwargs):
+        """Quantize the embedding table, look up, then optionally quantize the output.
+
+        ``input_quantizer`` is intentionally never invoked — embedding inputs are
+        integer indices. ``HardDisabledTensorQuantizer.set_from_attribute_config``
+        keeps that quantizer disabled regardless of what configs target it, so a
+        runtime check in this hot path is unnecessary.
+        """
+        if is_torch_export_mode():
+            # quantize_weight()'s attribute write is not allowed under torch.export;
+            # weight quantization is still applied inline via _get_quantized_weight's
+            # is_torch_export_mode() branch. Apply output_quantizer in this path too
+            # so users who opt into output activation quantization don't silently
+            # lose it during export — matches QuantInputBase.forward's behavior.
+            output = super().forward(input, *args, **kwargs)
+        else:
+            with self.quantize_weight():
+                output = super().forward(input, *args, **kwargs)
+        return self.output_quantizer(output)
+
+
+# Public alias consistent with quant_linear / quant_conv naming.
+QuantEmbedding = _QuantEmbedding

modelopt/torch/quantization/nn/modules/tensor_quantizer.py

@@ -64,6 +64,7 @@
 from ..functional import normalized_hadamard_transform
 
 __all__ = [
+    "HardDisabledTensorQuantizer",
     "NVFP4StaticQuantizer",
     "SequentialQuantizer",
     "TensorQuantizer",
@@ -1307,6 +1308,50 @@ def _set_buffer(self, key, value):
             self.register_buffer(key, value)
 
 
+class HardDisabledTensorQuantizer(TensorQuantizer):
+    """TensorQuantizer slot that is structurally never enabled.
+
+    Used where a quantizer attribute must exist for introspection / wildcard-config absorption
+    but the surrounding ``QuantModule`` will never actually invoke it — e.g. ``nn.Embedding``,
+    whose input is integer indices and cannot be fake-quantized.
+
+    Wildcard configs (e.g. the default ``QuantizeConfig`` ``"*"`` rule or
+    ``NVFP4_DEFAULT_CFG``'s ``*input_quantizer``) are accepted silently, then the quantizer
+    is force-disabled — wildcards mean "enable this kind of quantizer in general," not
+    "enable this specific structurally-disabled slot." Direct, explicit attempts (calling
+    ``enable`` / ``enable_quant`` / ``enable_calib``) raise loudly so a misconfigured recipe
+    doesn't silently mislead the user.
+    """
+
+    _DISABLED_ERR = (
+        "This TensorQuantizer is hard-disabled by its parent QuantModule (e.g. the "
+        "input_quantizer slot on a quantized nn.Embedding, whose input is integer indices) "
+        "and cannot be enabled."
+    )
+
+    def enable(self):
+        """Disallowed for hard-disabled quantizers."""
+        raise RuntimeError(self._DISABLED_ERR)
+
+    def enable_quant(self):
+        """Disallowed for hard-disabled quantizers."""
+        raise RuntimeError(self._DISABLED_ERR)
+
+    def enable_calib(self):
+        """Disallowed for hard-disabled quantizers."""
+        raise RuntimeError(self._DISABLED_ERR)
+
+    def set_from_attribute_config(self, attribute_cfg):
+        """Apply the config like any quantizer, then force-disable.
+
+        This absorbs wildcard configs from stock recipes without raising. Other attributes
+        (``num_bits``, ``axis``, etc.) take on the config values for introspection, but
+        ``_disabled`` is forced back to ``True`` so forward is always a no-op.
+        """
+        super().set_from_attribute_config(attribute_cfg)
+        self._disabled = True
+
+
 class NVFP4StaticQuantizer(TensorQuantizer):
     """TensorQuantizer for NVFP4 static block quantization with two-level scaling.
 

modelopt/torch/quantization/utils/core_utils.py

@@ -297,6 +297,11 @@ def is_quantized_linear(module):
     """Check if a module is a quantized linear module."""
     from ..nn import QuantModule, TensorQuantizer
 
+    # Embedding has a 2D weight but is not a GEMM op, so calibration passes that operate
+    # on linear activations (AWQ, SmoothQuant, SVDQuant) must skip it.
+    if isinstance(module, nn.Embedding):
+        return False
+
     return (
         isinstance(module, QuantModule)
         and isinstance(getattr(module, "input_quantizer", None), TensorQuantizer)

modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml

@@ -48,3 +48,6 @@
   - parent_class: 'nn.LeakyReLU'
     quantizer_name: '*'
     enable: false
+  - parent_class: 'nn.Embedding'
+    quantizer_name: '*'
+    enable: false

tests/gpu/torch/export/test_export_embedding.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Export-path tests for QuantEmbedding (CUDA-only).
+
+These tests drive ``_process_quantized_modules`` end-to-end. The export weight-packing
+path bottoms out in ``torch.cuda.empty_cache()``, which raises on systems without an
+NVIDIA driver, so the suite lives under ``tests/gpu/`` rather than ``tests/unit/``.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.export.unified_export_hf import _process_quantized_modules
+from modelopt.torch.quantization.utils import quantizer_attr_names
+
+VOCAB_SIZE = 16
+EMBED_DIM = 32  # multiple of the NVFP4 block size (16) so export packs cleanly
+
+
+def _embedding_nvfp4_cfg() -> dict:
+    """Stock-NVFP4-style cfg that opts the embedding's weight quantizer in."""
+    nvfp4 = {
+        "num_bits": (2, 1),
+        "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+    }
+    return {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "parent_class": "nn.Embedding",
+                "quantizer_name": "*weight_quantizer",
+                "cfg": dict(nvfp4),
+            },
+        ],
+        "algorithm": "max",
+    }
+
+
+class _EmbeddingOnly(nn.Module):
+    """Single-embedding wrapper exposing forward + named_modules iteration."""
+
+    def __init__(self):
+        """Build the lone embedding submodule."""
+        super().__init__()
+        self.embedding = nn.Embedding(VOCAB_SIZE, EMBED_DIM)
+
+    def forward(self, ids):
+        """Look up embeddings for the given token IDs."""
+        return self.embedding(ids)
+
+
+class _TiedEmbeddingLM(nn.Module):
+    """Embedding + Linear lm_head with tied weights (lm_head.weight is embedding.weight)."""
+
+    def __init__(self):
+        """Build embedding + lm_head and tie their weight Parameters."""
+        super().__init__()
+        self.embedding = nn.Embedding(VOCAB_SIZE, EMBED_DIM)
+        self.lm_head = nn.Linear(EMBED_DIM, VOCAB_SIZE, bias=False)
+        self.lm_head.weight = self.embedding.weight  # Python-level tie
+
+    def forward(self, ids):
+        """Embed then project to vocab logits with the tied weight."""
+        return self.lm_head(self.embedding(ids))
+
+
+class TestQuantEmbeddingExport:
+    """Export-path coverage: weight packing and tied-weight guard."""
+
+    def test_quantized_weight_is_packed_and_scales_registered(self):
+        """End-to-end: _process_quantized_modules packs the embedding weight and
+        registers ``weight_scale`` + ``weight_scale_2`` buffers."""
+        model = _EmbeddingOnly()
+        model = mtq.quantize(
+            model, _embedding_nvfp4_cfg(), lambda m: m(torch.randint(0, VOCAB_SIZE, (2, 4)))
+        )
+        _process_quantized_modules(model, dtype=torch.float16)
+
+        attrs = quantizer_attr_names("weight")
+        assert model.embedding.weight.dtype == torch.uint8
+        assert hasattr(model.embedding, attrs.weight_scale)
+        assert hasattr(model.embedding, attrs.weight_scale_2)
+        # input_scale is not registered (input_quantizer is permanently disabled).
+        assert not hasattr(model.embedding, attrs.input_scale)
+
+    def test_tied_embedding_export_skips_packing(self):
+        """When the embedding weight is shared with lm_head, packing is skipped
+        with a warning so the tie survives the export."""
+        model = _TiedEmbeddingLM()
+        assert model.lm_head.weight is model.embedding.weight  # sanity
+
+        model = mtq.quantize(
+            model, _embedding_nvfp4_cfg(), lambda m: m(torch.randint(0, VOCAB_SIZE, (2, 4)))
+        )
+        orig_dtype = model.embedding.weight.dtype
+        with pytest.warns(UserWarning, match="tied"):
+            _process_quantized_modules(model, dtype=torch.float16)
+
+        # Weight Parameter unchanged (not packed to uint8) and still tied.
+        assert model.embedding.weight.dtype == orig_dtype
+        assert model.lm_head.weight is model.embedding.weight