diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d3cc67d95bf..78a71f248f6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. +- Add the ``model_type/nemotron_h`` aggressive-MSE NVFP4 PTQ recipe: weight-MSE calibrated NVFP4 routed and shared MoE experts, FP8 attention projections, Mamba ``in_proj``/``out_proj`` and ``lm_head``, and an FP8 KV cache. The latent MoE, router, embeddings and the vision and audio encoders stay in BF16, and the MTP rules apply only to checkpoints that carry an MTP tail. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 52241079d39..bab9c1140cb 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -217,6 +217,25 @@ def create_vlm_calibration_loop(full_model, calib_dataloader): # Import here to avoid circular dependency from nemotron_vl_calib import safe_nemotron_vl_forward + # Check if model is encoder-decoder (needs decoder_input_ids instead of input_ids) + is_enc_dec = getattr(full_model.config, "is_encoder_decoder", False) + + # Which batches go through safe_nemotron_vl_forward: + # - Nano VL and other InternVL-style wrappers, recognised by `img_context_token_id`, need + # its embedding-injection path. + # - The omni wrappers define neither `img_context_token_id` nor `extract_feature`; they + # merge vision inside their own forward. They need the helper's fallback, which carries + # the dtype-cast `pixel_values` and sets `use_cache=False`. Requiring `extract_feature` + # to be absent keeps a Nemotron VL model that has the helper but no `img_context_token_id` + # out of the InternVL branch. + # Encoder-decoder VL models (Nemotron-Parse) keep the plain forward: their batch carries + # `decoder_input_ids`, which the helper does not recognise, so routing them through it would + # return without a forward pass and silently skip calibration. + use_safe_forward = not is_enc_dec and ( + hasattr(full_model, "img_context_token_id") + or (is_nemotron_vl(full_model) and not hasattr(full_model, "extract_feature")) + ) + def calibrate_loop(_model): # Inspect model's forward signature to determine what parameters it accepts forward_params = inspect.signature(full_model.forward).parameters @@ -225,9 +244,6 @@ def calibrate_loop(_model): ) allowed_keys = set(forward_params.keys()) - # Check if model is encoder-decoder (needs decoder_input_ids instead of input_ids) - is_enc_dec = getattr(full_model.config, "is_encoder_decoder", False) - full_model.eval() with torch.no_grad(): for batch in calib_dataloader: @@ -247,9 +263,7 @@ def calibrate_loop(_model): # Remove None values call_kwargs = {k: v for k, v in call_kwargs.items() if v is not None} - # Use safe_nemotron_vl_forward for Nemotron Nano VL (embedding-injection style) - # For other VLMs (like Nemotron-Parse), use standard forward - if hasattr(full_model, "img_context_token_id"): + if use_safe_forward: safe_nemotron_vl_forward(full_model, call_kwargs) else: full_model(**call_kwargs) diff --git a/examples/hf_ptq/nemotron_vl_calib.py b/examples/hf_ptq/nemotron_vl_calib.py index 398ef67fe5b..cffab639401 100644 --- a/examples/hf_ptq/nemotron_vl_calib.py +++ b/examples/hf_ptq/nemotron_vl_calib.py @@ -28,11 +28,25 @@ from __future__ import annotations import contextlib +import inspect from typing import Any import torch +def _accepted_kwargs(fn, kwargs: dict) -> dict: + """Drop keys ``fn`` cannot accept, so a wrapper with a narrow forward signature does not + trade one ``AttributeError`` for a ``TypeError``. Callables taking ``**kwargs`` are left alone. + """ + try: + params = inspect.signature(fn).parameters + except (TypeError, ValueError): + return kwargs + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return kwargs + return {k: v for k, v in kwargs.items() if k in params} + + def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) -> None: """Run a minimal multimodal forward for Nemotron VL that avoids wrapper output packaging.""" pixel_values = batch.get("pixel_values") @@ -67,6 +81,23 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) ): pixel_values = pixel_values.to(dtype=vision_dtype) + # InternVL-style Nemotron VL exposes `extract_feature(pixel_values)`; the newer omni wrappers + # (e.g. NemotronH_Omni_Reasoning_V3 / nemotron_h_omni) do not -- they merge the vision + # embeddings via masked_scatter inside their own forward, and do not necessarily define + # `img_context_token_id` either. Hand the batch to the model's own forward before touching + # anything InternVL-specific, carrying the dtype-cast pixel_values. + if not hasattr(full_model, "extract_feature"): + # Only the batch's own entries, plus the dtype-cast pixel_values. The synthesized + # `image_flags` above is deliberately NOT injected: it is built on pixel_values.device, + # while a sharded wrapper runs its vision tower elsewhere and indexes image_embeds with + # it, so a synthesized tensor raises "indices should be ... on the same device". + # Wrappers that need image_flags either receive the batch's own or build their own. + fallback = {k: v for k, v in batch.items() if v is not None} + fallback["pixel_values"] = pixel_values + fallback["use_cache"] = False + full_model(**_accepted_kwargs(full_model.forward, fallback)) + return + # Token embeddings inputs_embeds = full_model.language_model.get_input_embeddings()(input_ids) image_flags_s = image_flags.squeeze(-1) @@ -75,10 +106,22 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) flat_embeds = inputs_embeds.reshape(b * n, c) flat_ids = input_ids.reshape(b * n) selected = flat_ids == full_model.img_context_token_id + # A bool mask must sit on the device of the tensor it indexes, not the one it was derived + # from. Under a sharded device_map the embedding table can run on a different GPU than the + # batch, so `selected` -- built from `input_ids` -- follows `flat_embeds`. Both the fast + # path and the except-branch retry below index with it. + if selected.device != flat_embeds.device: + selected = selected.to(flat_embeds.device) - # Vision embeddings vit_embeds = full_model.extract_feature(pixel_values) + # Same rule for the image filter: the vision tower returns on its own device, while + # `image_flags` arrived with the batch (or was synthesized on `pixel_values.device`). + if image_flags_s.device != vit_embeds.device: + image_flags_s = image_flags_s.to(vit_embeds.device) vit_embeds = vit_embeds[image_flags_s == 1] + # The merge target lives on the embedding device; bring the filtered vision embeddings over. + if vit_embeds.device != flat_embeds.device: + vit_embeds = vit_embeds.to(flat_embeds.device) try: flat_embeds[selected] = flat_embeds[selected] * 0.0 + vit_embeds.reshape(-1, c) except Exception: @@ -88,6 +131,16 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) inputs_embeds = flat_embeds.reshape(b, n, c) + # `inputs_embeds` comes off the embedding table, so under a sharded device_map it can sit on a + # different device than the mask/ids that came in with the batch. Align them to the embedding + # device so the LLM call starts consistent; accelerate's hooks handle later blocks. + _lm_dev = getattr(full_model.language_model.get_input_embeddings().weight, "device", None) + if _lm_dev is not None: + if attention_mask is not None and attention_mask.device != _lm_dev: + attention_mask = attention_mask.to(_lm_dev) + if position_ids is not None and position_ids.device != _lm_dev: + position_ids = position_ids.to(_lm_dev) + # LLM forward (drives activation stats) full_model.language_model( inputs_embeds=inputs_embeds, diff --git a/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml b/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml new file mode 100644 index 00000000000..9445dd3b25f --- /dev/null +++ b/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml @@ -0,0 +1,150 @@ +# 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. + + +# NemotronH -- aggressive NVFP4 tier, weight-MSE calibration. +# +# For hybrid Mamba-MoE NemotronH checkpoints, including the omni VL wrappers. Broader FP8 weight +# coverage than the conservative tier; the MTP rules apply only when the checkpoint carries an +# MTP tail and are inert otherwise. +# +# MoE routed + shared experts NVFP4 W4A4, static weight scales +# attention q/k/v/o, Mamba in/out_proj, lm_head FP8 +# MTP block-1 routed + shared experts NVFP4 W4A4, static weight scales +# latent MoE, embeddings, router, C-RADIO vision, Parakeet audio BF16 +# KV cache FP8 (see the kv rules below) +# +# quant_cfg is ordered and the LAST match wins, so the blanket disable comes first and the +# vision / MTP rules come last. Only configs/numerics/* is imported, so every rule that +# applies to this model is visible in this file. + +imports: + nvfp4: configs/numerics/nvfp4 + nvfp4_static: configs/numerics/nvfp4_static + fp8: configs/numerics/fp8 + +metadata: + recipe_type: ptq + description: >- + NemotronH aggressive tier, weight-MSE calibrated: NVFP4 routed and shared MoE experts with + static weight scales and dynamic NVFP4 activations, the same for the MTP block-1 experts; + FP8 attention projections, Mamba in/out_proj and lm_head; BF16 latent MoE and vision tower. + FP8 KV cache, with the backbone cast at a fixed range and MTP calibrated. +quantize: + algorithm: + method: mse + # Sweeps the 128 e4m3 scale values; requires static weight quantizers (below). + fp8_scale_sweep: true + # Required for VLMs: decoder layers sit under `model.language_model.layers`, + # where layerwise_calibrate cannot find them. + layerwise: + enable: false + quant_cfg: + - quantizer_name: '*' + enable: false + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*shared_experts*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*shared_experts*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*[qkvo]_proj.weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*[qkvo]_proj.input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*mixer.in_proj.weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*mixer.in_proj.input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*mixer.out_proj.weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*mixer.out_proj.input_quantizer' + cfg: + $import: fp8 + # Backbone KV casts at a fixed FP8 range: no _amax buffer, so no k_scale/v_scale is + # exported, matching configs/ptq/units/kv_fp8_cast. MTP KV below stays calibrated and + # keeps its pair. + - quantizer_name: '*[kv]_bmm_quantizer' + cfg: + $import: fp8 + use_constant_amax: true + - quantizer_name: '*mtp.*' + enable: false + - quantizer_name: '*mtp.layers.1.mixer.experts.*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*mtp.layers.1.mixer.experts.*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mtp.layers.1.mixer.shared_experts.*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*mtp.layers.1.mixer.shared_experts.*input_quantizer' + cfg: + $import: nvfp4 + # Restated so this rule cannot inherit the backbone setting above. + - quantizer_name: '*mtp.layers.0.*[kv]_bmm_quantizer' + cfg: + $import: fp8 + use_constant_amax: false + - quantizer_name: '*embed_vision*' + enable: false + - quantizer_name: '*lm_head*weight_quantizer' + enable: true + cfg: + $import: fp8 + - quantizer_name: '*lm_head*input_quantizer' + enable: true + cfg: + $import: fp8 + - quantizer_name: '*vision_model*' + enable: false + - quantizer_name: '*vision_projector*' + enable: false + - quantizer_name: '*radio_model*' + enable: false + - quantizer_name: '*vision_tower*' + enable: false + - quantizer_name: '*visual*' + enable: false + - quantizer_name: '*multi_modal_projector*' + enable: false + - quantizer_name: 'mlp1.*' + enable: false + - quantizer_name: '*.mlp1.*' + enable: false + # Audio branch of the omni wrapper, present when the config carries a `sound_config`. + # The Parakeet encoder names its projections `self_attn.{q,k,v,o,relative_k}_proj`, so the + # FP8 attention rules above match it, and its attention class is registered for KV-cache + # quantization like any other. Image calibration supplies no audio, so those input + # quantizers collect no amax and the export drops the input scale while still writing an + # FP8 weight scale -- a layer that claims W8A8 with no activation scale. Keep it BF16. + - quantizer_name: '*audio_tower*' + enable: false + - quantizer_name: '*sound_encoder*' + enable: false + - quantizer_name: '*sound_projector*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 4c3751fa394..d7eb46775ea 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -308,7 +308,7 @@ that baseline. The deviations come in four kinds: | Kind | What changes vs. the general recipe | Examples | |------|-------------------------------------|----------| -| **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `minimax_m3_vl`, `qwen3_vl`, `qwen3_5`, `qwen3_5_moe`, `qwen3_6_moe`, `vit`, `nemotron_llama` | +| **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `minimax_m3_vl`, `qwen3_vl`, `qwen3_5`, `qwen3_5_moe`, `qwen3_6_moe`, `vit`, `nemotron_llama`, `nemotron_h` | | **Algorithm override** | Same numerics & scope, but the *calibration algorithm* is tweaked because the default breaks or regresses | `gemma`, `gemma4`, `mpt` | | **Extra exclusions** | Adds disabled-quantizer patterns so non-language branches stay full precision | `nemotron_vl`, `diffusion_gemma` | | **Checkpoint mirror** | A mixed-precision map reproducing one published checkpoint exactly | `models/nvidia/NVIDIA-Nemotron-3-*`, `models/mistralai/Mistral-Medium-3.5-128B` | @@ -326,7 +326,7 @@ folder may carry a `README.md` spelling out that delta. > genuinely deviate and aliases the ones that don't; see > [`models/README.md`](models/README.md) for the alias format and when to write one. -### Architecture-aware `quant_cfg` — `minimax_m3_vl`, `qwen3_vl`, `qwen3_5`, `qwen3_5_moe`, `qwen3_6_moe`, `vit`, `nemotron_llama` +### Architecture-aware `quant_cfg` — `minimax_m3_vl`, `qwen3_vl`, `qwen3_5`, `qwen3_5_moe`, `qwen3_6_moe`, `vit`, `nemotron_llama`, `nemotron_h` **`minimax_m3_vl/ptq/mxfp8_nvfp4_experts`** applies MXFP8 to the language-model linear layers and MSE-calibrated NVFP4 to routed experts, with expert @@ -407,6 +407,22 @@ recipes, for Step-3.7 checkpoints; Step-3.5 has its own recipe above. load-bearing: selectors are Megatron-Core leaf names (`mlp.experts.linear_fc1`, `self_attention.linear_qkv`), so under `hf_ptq.py` nothing matches and `mtq.quantize` raises. +- **`nemotron_h/ptq/nvfp4-aggressive-mse`** is the **aggressive** tier for hybrid Mamba-MoE + NemotronH checkpoints, including the omni VL wrappers: + - MoE routed + shared experts → NVFP4 W4A4, `group_size 16`, **static** weight scales + (`nvfp4_static`) with dynamic NVFP4 activations (`nvfp4`) + - attention q/k/v/out, Mamba `in/out_proj`, `lm_head` → FP8 + - MTP block-1 routed + shared experts → identical treatment to the main experts; the MTP rules + are inert on checkpoints without an MTP tail + - latent MoE, embeddings, router and the vision tower → **BF16** + - KV cache → FP8, **asymmetric**: the language-model backbone casts at a fixed FP8 range + (`use_constant_amax`, as in `configs/ptq/units/kv_fp8_cast`, so no `k_scale`/`v_scale` is + exported) while MTP attention stays data-calibrated and keeps its pair + + Calibration is weight-MSE with an FP8-scale sweep, which is why the weight quantizers are + `static`. Use `--calib_with_images` on VL checkpoints: the vision tower must see image batches + even though it stays BF16. + ### Algorithm overrides — `gemma`, `gemma4`, `mpt` These quantize the **same layers** as the general recipes; only the diff --git a/tests/_test_utils/examples/hf_ptq_example_utils.py b/tests/_test_utils/examples/hf_ptq_example_utils.py index b9a0ebe03b1..0a1b70503ec 100644 --- a/tests/_test_utils/examples/hf_ptq_example_utils.py +++ b/tests/_test_utils/examples/hf_ptq_example_utils.py @@ -12,8 +12,8 @@ # 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. -"""Re-export ``examples/hf_ptq/example_utils`` so tests can import it via -``from _test_utils.examples.hf_ptq_example_utils import example_utils`` +"""Re-export the ``examples/hf_ptq`` helper modules so tests can import them via +``from _test_utils.examples.hf_ptq_example_utils import example_utils, nemotron_vl_calib`` without per-file ``sys.path`` shims. """ @@ -26,5 +26,6 @@ sys.path.insert(0, str(_HF_PTQ_DIR)) import example_utils +import nemotron_vl_calib -__all__ = ["example_utils"] +__all__ = ["example_utils", "nemotron_vl_calib"] diff --git a/tests/examples/hf_ptq/test_nemotron_vl_calib.py b/tests/examples/hf_ptq/test_nemotron_vl_calib.py new file mode 100644 index 00000000000..d0e45a3c651 --- /dev/null +++ b/tests/examples/hf_ptq/test_nemotron_vl_calib.py @@ -0,0 +1,346 @@ +# 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. + +"""The Nemotron VL calibration fallback for wrappers without ``extract_feature``. + +Two layers are covered: the helper itself, and the dispatch in +``create_vlm_calibration_loop()`` that decides which wrappers reach it. The dispatch half is +the one that matters in production -- a wrapper can only benefit from the fallback if the +calibration loop routes it there. +""" + +from types import SimpleNamespace + +import torch +from _test_utils.examples.hf_ptq_example_utils import example_utils, nemotron_vl_calib + + +class _OmniWrapper(torch.nn.Module): + """An omni-style wrapper: merges vision inside forward, so no ``extract_feature``. + + Deliberately also has no ``img_context_token_id`` and no ``language_model`` -- touching + either before the fallback fires is the regression this guards. + """ + + def __init__(self, config=None): + super().__init__() + self.config = config + self.calls: list[dict] = [] + + def forward(self, **kwargs): + self.calls.append(kwargs) + return torch.zeros(1) + + +def _omni_config(): + """A config that ``is_nemotron_vl()`` recognises: multimodal, Nemotron architecture.""" + return SimpleNamespace( + architectures=["NemotronH_Omni_ForConditionalGeneration"], + vision_config=SimpleNamespace(torch_dtype=torch.bfloat16), + is_encoder_decoder=False, + ) + + +def _batch(): + return { + "input_ids": torch.ones(1, 4, dtype=torch.long), + "pixel_values": torch.zeros(2, 3, 8, 8), + "attention_mask": torch.ones(1, 4, dtype=torch.long), + "position_ids": None, + } + + +def test_fallback_forwards_batch_when_extract_feature_is_absent(): + model = _OmniWrapper() + nemotron_vl_calib.safe_nemotron_vl_forward(model, _batch()) + + assert len(model.calls) == 1, "the wrapper's own forward should be called exactly once" + passed = model.calls[0] + # every non-None batch key reaches the model + assert passed["input_ids"].shape == (1, 4) + assert passed["pixel_values"].shape == (2, 3, 8, 8) + assert passed["attention_mask"].shape == (1, 4) + # None-valued keys are dropped, not forwarded + assert "position_ids" not in passed + # calibration must not build a KV cache + assert passed["use_cache"] is False + # a synthesized image_flags must NOT be injected: it is built on pixel_values.device, and a + # sharded wrapper indexes its vision output with it from another device (seen in a real + # 4-GPU run as "indices should be either on cpu or on the same device as the indexed tensor") + assert "image_flags" not in passed + + +def test_fallback_forwards_image_flags_the_batch_already_carried(): + """A batch-supplied image_flags is device-consistent with the rest, so it passes through.""" + model = _OmniWrapper() + batch = _batch() + batch["image_flags"] = torch.ones(2, 1, dtype=torch.long) + nemotron_vl_calib.safe_nemotron_vl_forward(model, batch) + + assert model.calls[0]["image_flags"].shape == (2, 1) + + +def test_fallback_drops_kwargs_a_narrow_forward_cannot_accept(): + """A wrapper without ``**kwargs`` must not get a TypeError instead of the old AttributeError.""" + + class _Narrow(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls: list[dict] = [] + + def forward(self, input_ids=None, pixel_values=None): + self.calls.append({"input_ids": input_ids, "pixel_values": pixel_values}) + return torch.zeros(1) + + model = _Narrow() + nemotron_vl_calib.safe_nemotron_vl_forward(model, _batch()) # must not raise + + assert len(model.calls) == 1 + assert model.calls[0]["pixel_values"].shape == (2, 3, 8, 8) + + +def test_accepted_kwargs_passes_everything_to_a_var_keyword_callable(): + def takes_anything(**kwargs): + pass + + payload = {"a": 1, "b": 2} + assert nemotron_vl_calib._accepted_kwargs(takes_anything, payload) == payload + + +def test_accepted_kwargs_filters_a_fixed_signature(): + def takes_only_a(a=None): + pass + + assert nemotron_vl_calib._accepted_kwargs(takes_only_a, {"a": 1, "b": 2}) == {"a": 1} + + +def test_calibration_loop_routes_an_omni_wrapper_through_the_fallback(): + """The production dispatch, not the helper directly. + + An omni wrapper has no ``img_context_token_id``, so the loop used to call its forward with + the raw batch -- bypassing the dtype cast and the ``use_cache=False`` the fallback supplies. + """ + model = _OmniWrapper(config=_omni_config()) + model.vision_model = torch.nn.Module() + model.vision_model.config = SimpleNamespace(torch_dtype=torch.bfloat16) + + example_utils.create_vlm_calibration_loop(model, [_batch()])(model) + + assert len(model.calls) == 1 + passed = model.calls[0] + assert passed["use_cache"] is False + assert passed["pixel_values"].dtype is torch.bfloat16, "the fallback carries the dtype cast" + assert "position_ids" not in passed + + +def test_calibration_loop_keeps_the_plain_forward_for_encoder_decoder_vlms(): + """Nemotron-Parse is a Nemotron VL model too, but must not be rerouted. + + Its batch is renamed to ``decoder_input_ids``; the helper looks for ``input_ids`` and would + return without a forward pass, silently skipping calibration. + """ + model = _OmniWrapper( + config=SimpleNamespace( + architectures=["NemotronParseForConditionalGeneration"], + is_encoder_decoder=True, + ) + ) + + example_utils.create_vlm_calibration_loop(model, [_batch()])(model) + + assert len(model.calls) == 1 + passed = model.calls[0] + assert "decoder_input_ids" in passed, "the encoder-decoder rename reached the model" + assert "input_ids" not in passed + # the helper's fallback was not applied + assert passed["pixel_values"].dtype is torch.float32 + + +# --- sharded InternVL path ----------------------------------------------------------------- +# +# The omni tests above exit through the fallback and never reach the InternVL branch, where two +# bool masks index tensors that a sharded ``device_map`` can place on other GPUs. This group +# runs CPU-only, so the topology is simulated rather than allocated -- the convention +# ``test_example_utils.py`` already uses for ``torch.cuda.device_count``. + + +def _device_of(tensor): + return getattr(tensor, "_fake_device", torch.device("cpu")) + + +def _shard(tensor, device): + """Tag a CPU tensor as living on ``device``.""" + out = tensor.as_subclass(_ShardedTensor) + out._fake_device = torch.device(device) + return out + + +class _ShardedTensor(torch.Tensor): + """CPU-backed tensor that reports a fake device and enforces CUDA's indexing rule. + + Storage stays on CPU; only the reported device, ``.to(device)`` and the cross-device + indexing error are faked. CUDA raises when a bool mask lives on another device, and that + is the failure this reproduces without needing two GPUs. + """ + + _fake_device = torch.device("cpu") + + @property + def device(self): + return self._fake_device + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if func in (torch.Tensor.__getitem__, torch.Tensor.__setitem__): + target, index = args[0], args[1] + if isinstance(index, torch.Tensor) and _device_of(index) != _device_of(target): + raise RuntimeError( + "indices should be either on cpu or on the same device as the indexed " + f"tensor (got {_device_of(index)} vs {_device_of(target)})" + ) + if func is torch.Tensor.to: + target = kwargs.get("device", args[1] if len(args) > 1 else None) + if isinstance(target, (str, torch.device)): + return _shard(args[0].clone(), target) + out = super().__torch_function__(func, types, args, kwargs) + device = next((a._fake_device for a in args if isinstance(a, _ShardedTensor)), None) + if device is not None: + for tensor in out if isinstance(out, (list, tuple)) else [out]: + if isinstance(tensor, _ShardedTensor): + tensor._fake_device = device + return out + + +class _Embedding: + """Callable embedding table that answers on ``device``, as ``.weight.device`` does.""" + + def __init__(self, device, hidden): + self.weight = _shard(torch.zeros(16, hidden), device) + self._device = device + self._hidden = hidden + + def __call__(self, input_ids): + b, n = input_ids.shape + return _shard(torch.zeros(b, n, self._hidden), self._device) + + +class _LanguageModel(torch.nn.Module): + def __init__(self, device, hidden): + super().__init__() + self.config = SimpleNamespace(torch_dtype=torch.float32) + self.calls: list[dict] = [] + self._embedding = _Embedding(device, hidden) + + def get_input_embeddings(self): + return self._embedding + + def forward(self, **kwargs): + self.calls.append(kwargs) + return (torch.zeros(1),) + + +class _ShardedInternVL(torch.nn.Module): + """InternVL-style wrapper with vision tower, embedding table and batch on three devices.""" + + def __init__(self, vision_device, embed_device, tokens_per_image=1, hidden=4): + super().__init__() + self.config = SimpleNamespace( + architectures=["NemotronVLForConditionalGeneration"], is_encoder_decoder=False + ) + self.vision_model = SimpleNamespace(config=SimpleNamespace(torch_dtype=torch.float32)) + self.language_model = _LanguageModel(embed_device, hidden) + self.img_context_token_id = 7 + self._vision_device = vision_device + self._tokens_per_image = tokens_per_image + self._hidden = hidden + + def extract_feature(self, pixel_values): + return _shard( + torch.ones(pixel_values.shape[0], self._tokens_per_image, self._hidden), + self._vision_device, + ) + + +def _sharded_batch(device): + """Two images and four tokens, of which two are image-context tokens. + + ``image_flags`` is supplied rather than left to the helper's synthesis step: that step + allocates on ``pixel_values.device``, which under this harness is a *fake* device name, so + it would escape the bookkeeping. A processor that emits ``image_flags`` puts it on the + batch device anyway, which is exactly the placement the synthesis step reproduces. + """ + return { + "input_ids": _shard(torch.tensor([[7, 7, 1, 1]], dtype=torch.long), device), + "pixel_values": _shard(torch.zeros(2, 3, 8, 8), device), + "attention_mask": _shard(torch.ones(1, 4, dtype=torch.long), device), + "image_flags": _shard(torch.ones(2, 1, dtype=torch.long), device), + "position_ids": None, + } + + +def test_sharded_internvl_aligns_both_index_masks(): + """``image_flags_s`` must follow the vision output and ``selected`` the embedding table. + + Pre-fix this raised on ``vit_embeds[image_flags_s == 1]``, before any alignment ran. + """ + model = _ShardedInternVL(vision_device="cuda:1", embed_device="cuda:2") + + nemotron_vl_calib.safe_nemotron_vl_forward(model, _sharded_batch("cuda:0")) + + assert len(model.language_model.calls) == 1, "the LLM forward drives the activation stats" + passed = model.language_model.calls[0] + assert passed["inputs_embeds"].device == torch.device("cuda:2") + assert passed["attention_mask"].device == torch.device("cuda:2") + assert passed["use_cache"] is False + + +def test_sharded_internvl_retry_path_also_uses_aligned_masks(): + """The ``except`` branch re-indexes with ``selected``, so it needs the aligned mask too. + + Two images of two tokens each give four vision rows for two selected positions, so the + first assignment raises on shape and the retry runs -- under the same sharding. + """ + model = _ShardedInternVL(vision_device="cuda:1", embed_device="cuda:2", tokens_per_image=2) + + nemotron_vl_calib.safe_nemotron_vl_forward(model, _sharded_batch("cuda:0")) + + assert len(model.language_model.calls) == 1, "the retry still reaches the LLM forward" + + +def test_unsharded_internvl_is_unaffected(): + """Everything on one device: no mask is moved and the merge still happens.""" + model = _ShardedInternVL(vision_device="cuda:0", embed_device="cuda:0") + + nemotron_vl_calib.safe_nemotron_vl_forward(model, _sharded_batch("cuda:0")) + + assert len(model.language_model.calls) == 1 + assert model.language_model.calls[0]["inputs_embeds"].device == torch.device("cuda:0") + + +def test_sharded_internvl_aligns_the_token_mask_when_only_the_embedding_moves(): + """Isolates the second mask, which the first test cannot reach. + + When the vision tower shares the batch device, ``vit_embeds[image_flags_s == 1]`` succeeds + and execution gets as far as ``flat_embeds[selected]``, where ``selected`` -- built from + ``input_ids`` -- meets ``flat_embeds`` on the embedding device. Pre-fix that raises, and the + ``except`` branch retries with the same mask and raises again. + """ + model = _ShardedInternVL(vision_device="cuda:0", embed_device="cuda:2") + + nemotron_vl_calib.safe_nemotron_vl_forward(model, _sharded_batch("cuda:0")) + + assert len(model.language_model.calls) == 1 + assert model.language_model.calls[0]["inputs_embeds"].device == torch.device("cuda:2") diff --git a/tests/unit/recipe/test_nemotron_h_recipe.py b/tests/unit/recipe/test_nemotron_h_recipe.py new file mode 100644 index 00000000000..b29dd01795f --- /dev/null +++ b/tests/unit/recipe/test_nemotron_h_recipe.py @@ -0,0 +1,316 @@ +# 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. + +"""Effective quantizer selection for ``model_type/nemotron_h/ptq/nvfp4-aggressive-mse``. + +The recipe's ``quant_cfg`` is ordered and last-match-wins, so what a module ends up with +depends on rule order as much as on rule content. The KV asymmetry in particular is load +bearing: the backbone casts at a fixed FP8 range and exports no ``k_scale``/``v_scale``, while +the MTP block keeps a calibrated pair. A reordering or a widened glob would silently flip that, +and nothing else in the tree-wide recipe tests would notice. + +The module tree below is synthetic: the real checkpoint is 120B, and what is under test is +name matching, not numerics. +""" + +import pytest +import torch.nn as nn + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry, TensorQuantizer + +_RECIPE = "model_type/nemotron_h/ptq/nvfp4-aggressive-mse" +_WIDTH = 32 + + +class _Attention(nn.Module): + """A NemotronH attention mixer. + + ``relative`` adds the Parakeet audio encoder's extra projection: ``*[qkvo]_proj`` matches + ``relative_k_proj`` too, on the ``k``, so narrowing the attention globs is not a way out. + """ + + def __init__(self, relative: bool = False): + super().__init__() + self.q_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + self.k_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + self.v_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + self.o_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + if relative: + self.relative_k_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + + +class _QuantAttention(QuantModule): + """The KV quantizers ``register_hf_attentions_on_the_fly`` installs on a real checkpoint. + + They have to arrive through the registry rather than be set in ``_Attention.__init__``: + a model that already holds a ``TensorQuantizer`` is treated as quantized, and + ``mtq.quantize`` would then skip the module conversion that creates the weight quantizers. + """ + + def _setup(self): + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + +@pytest.fixture(autouse=True) +def _register_kv_attention(): + mtq.register(original_cls=_Attention, quantized_cls=_QuantAttention) + yield + if QuantModuleRegistry.get(_Attention) is not None: + mtq.unregister(_Attention) + + +class _Mamba(nn.Module): + def __init__(self): + super().__init__() + self.in_proj = nn.Linear(_WIDTH, 2 * _WIDTH, bias=False) + self.out_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + + +class _Expert(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + self.down_proj = nn.Linear(_WIDTH, _WIDTH, bias=False) + + +class _MoE(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.ModuleList([_Expert(), _Expert()]) + self.shared_experts = _Expert() + self.router = nn.Linear(_WIDTH, 2, bias=False) + + +class _Block(nn.Module): + def __init__(self, mixer): + super().__init__() + self.mixer = mixer + + +def _nemotron_h(with_mtp: bool) -> nn.Module: + """Hub-named NemotronH skeleton: attention, Mamba and MoE blocks, vision and audio + encoders, MTP tail. + """ + model = nn.Module() + model.backbone = nn.Module() + model.backbone.embeddings = nn.Linear(_WIDTH, _WIDTH, bias=False) + model.backbone.layers = nn.ModuleList([_Block(_Attention()), _Block(_Mamba()), _Block(_MoE())]) + model.lm_head = nn.Linear(_WIDTH, _WIDTH, bias=False) + # VL wrapper pieces that must stay in BF16. Both encoders carry an attention block, not + # just projections: their attention classes end in `Attention`, so a real run registers + # them for KV-cache quantization and `*[kv]_bmm_quantizer` reaches them. The audio branch + # exists whenever the omni config carries a `sound_config`, and image calibration never + # exercises it, so anything enabled there exports without an activation scale. + model.embed_vision = nn.Linear(_WIDTH, _WIDTH, bias=False) + model.vision_model = nn.Module() + model.vision_model.radio_model = nn.Module() + model.vision_model.radio_model.blocks = nn.ModuleList([_Expert(), _Block(_Attention())]) + model.sound_projector = nn.Module() + model.sound_projector.sound_encoder = nn.Module() + model.sound_projector.sound_encoder.layers = nn.ModuleList([_Block(_Attention(relative=True))]) + if with_mtp: + # block 0 carries the attention (hence the KV pair), block 1 the MoE + model.mtp = nn.Module() + model.mtp.layers = nn.ModuleList([_Block(_Attention()), _Block(_MoE())]) + return model + + +def _quantize(with_mtp: bool): + """Apply the recipe to a fresh skeleton and return it with its named modules.""" + model = _nemotron_h(with_mtp) + config = load_recipe(_RECIPE).quantize.model_dump() + + # Calibration settings the recipe promises, asserted before they are dropped: weight-MSE + # with the FP8 scale sweep, and layerwise calibration off (VL decoder layers sit under + # `model.language_model.layers`, where layerwise_calibrate cannot find them). + assert config["algorithm"]["method"] == "mse" + assert config["algorithm"]["fp8_scale_sweep"] is True + assert config["algorithm"]["layerwise"]["enable"] is False + + # No forward loop here; selection and precedence are what is under test. + config["algorithm"] = None + mtq.quantize(model, config) + return model, dict(model.named_modules()) + + +def _assert_fp8(quantizer, name): + """Assert ``quantizer`` is enabled per-tensor FP8 E4M3.""" + assert quantizer.is_enabled, f"{name} should be enabled" + assert quantizer.num_bits == (4, 3), f"{name} should be FP8 E4M3" + assert quantizer.block_sizes is None, f"{name} should be per-tensor" + + +def _assert_nvfp4(quantizer, name, *, block_type): + """Assert ``quantizer`` is enabled NVFP4 E2M1 over 16-element blocks with FP8 scales.""" + assert quantizer.is_enabled, f"{name} should be enabled" + assert quantizer.num_bits == (2, 1), f"{name} should be NVFP4 E2M1" + assert quantizer.block_sizes[-1] == 16, f"{name} should use 16-element blocks" + assert quantizer.block_sizes["type"] == block_type, f"{name} should be {block_type}" + assert quantizer.block_sizes["scale_bits"] == (4, 3), f"{name} should carry FP8 scales" + + +@pytest.mark.parametrize("with_mtp", [True, False]) +def test_backbone_kv_casts_at_a_fixed_range_and_exports_no_scale(with_mtp): + """`use_constant_amax` pins amax to the FP8 E4M3 max and registers no `_amax` buffer, so the + export carries no backbone `k_scale`/`v_scale`. The shipped checkpoint has exactly one of + each, both from MTP; adding backbone scales here would break that match. + """ + _, modules = _quantize(with_mtp) + + for which in ("k", "v"): + name = f"backbone.layers.0.mixer.{which}_bmm_quantizer" + quantizer = modules[name] + _assert_fp8(quantizer, name) + assert quantizer._use_constant_amax is True, f"{name} should cast at a fixed range" + assert not hasattr(quantizer, "_amax"), f"{name} must register no exportable scale" + + +def test_mtp_kv_stays_calibrated(): + """The MTP rule restates `use_constant_amax: false` so it cannot inherit the backbone rule.""" + _, modules = _quantize(with_mtp=True) + + for which in ("k", "v"): + name = f"mtp.layers.0.mixer.{which}_bmm_quantizer" + quantizer = modules[name] + _assert_fp8(quantizer, name) + assert quantizer._use_constant_amax is False, f"{name} should stay calibrated" + + +@pytest.mark.parametrize("with_mtp", [True, False]) +def test_backbone_expert_and_projection_formats(with_mtp): + """Experts are NVFP4 (static weights, dynamic inputs); projections and lm_head are FP8.""" + _, modules = _quantize(with_mtp) + + for expert in ("experts.0", "experts.1", "shared_experts"): + for proj in ("up_proj", "down_proj"): + stem = f"backbone.layers.2.mixer.{expert}.{proj}" + _assert_nvfp4(modules[f"{stem}.weight_quantizer"], stem, block_type="static") + _assert_nvfp4(modules[f"{stem}.input_quantizer"], stem, block_type="dynamic") + + for proj in ("q_proj", "k_proj", "v_proj", "o_proj"): + stem = f"backbone.layers.0.mixer.{proj}" + _assert_fp8(modules[f"{stem}.weight_quantizer"], stem) + _assert_fp8(modules[f"{stem}.input_quantizer"], stem) + + for proj in ("in_proj", "out_proj"): + stem = f"backbone.layers.1.mixer.{proj}" + _assert_fp8(modules[f"{stem}.weight_quantizer"], stem) + _assert_fp8(modules[f"{stem}.input_quantizer"], stem) + + _assert_fp8(modules["lm_head.weight_quantizer"], "lm_head") + _assert_fp8(modules["lm_head.input_quantizer"], "lm_head") + + +def test_mtp_block_one_experts_match_the_backbone_formats(): + """The MTP experts are NVFP4 with static weight scales and dynamic inputs, like block 2. + + `*mtp.*` disables the whole tail first, so only the rules that follow it bring anything + back: the block-1 experts and the block-0 KV pair. The MTP attention projections stay BF16. + """ + _, modules = _quantize(with_mtp=True) + + for expert in ("experts.0", "experts.1", "shared_experts"): + for proj in ("up_proj", "down_proj"): + stem = f"mtp.layers.1.mixer.{expert}.{proj}" + _assert_nvfp4(modules[f"{stem}.weight_quantizer"], stem, block_type="static") + _assert_nvfp4(modules[f"{stem}.input_quantizer"], stem, block_type="dynamic") + + for proj in ("q_proj", "o_proj"): + for kind in ("weight_quantizer", "input_quantizer"): + name = f"mtp.layers.0.mixer.{proj}.{kind}" + assert not modules[name].is_enabled, f"{name} should stay BF16" + assert not modules["mtp.layers.1.mixer.router.weight_quantizer"].is_enabled + + +@pytest.mark.parametrize("with_mtp", [True, False]) +def test_vision_audio_router_and_embeddings_stay_bf16(with_mtp): + """Everything the blanket `*` disable is meant to leave alone stays unquantized. + + The encoder attention entries are the load-bearing ones: they are named by the FP8 + attention and KV rules earlier in the file and only the trailing disables turn them back + off. A constant-amax KV cast exports no scale either way, so no artefact check would + catch it being left on. + """ + _, modules = _quantize(with_mtp) + + disabled = [ + "backbone.embeddings.weight_quantizer", + "backbone.layers.2.mixer.router.weight_quantizer", + "embed_vision.weight_quantizer", + "embed_vision.input_quantizer", + "vision_model.radio_model.blocks.0.up_proj.weight_quantizer", + "vision_model.radio_model.blocks.0.up_proj.input_quantizer", + "vision_model.radio_model.blocks.0.down_proj.weight_quantizer", + "vision_model.radio_model.blocks.1.mixer.q_proj.weight_quantizer", + "vision_model.radio_model.blocks.1.mixer.k_bmm_quantizer", + "sound_projector.sound_encoder.layers.0.mixer.q_proj.weight_quantizer", + "sound_projector.sound_encoder.layers.0.mixer.o_proj.input_quantizer", + "sound_projector.sound_encoder.layers.0.mixer.relative_k_proj.weight_quantizer", + "sound_projector.sound_encoder.layers.0.mixer.v_bmm_quantizer", + ] + for name in disabled: + assert not modules[name].is_enabled, f"{name} should stay BF16" + + +def test_mtp_rules_are_inert_without_an_mtp_tail(): + """The recipe is shipped for checkpoints with and without MTP; the tail rules must not + change what the backbone gets. + """ + _, with_mtp = _quantize(with_mtp=True) + _, without_mtp = _quantize(with_mtp=False) + + def summary(modules): + return { + name: (module.is_enabled, module.num_bits, module.block_sizes) + for name, module in modules.items() + if isinstance(module, TensorQuantizer) and not name.startswith("mtp.") + } + + assert summary(with_mtp) == summary(without_mtp) + assert not any(name.startswith("mtp.") for name in without_mtp) + + +def test_no_quantizer_outside_the_intended_set_is_enabled(): + """A widened glob would show up here first: the blanket `*` disable comes first, so every + enabled quantizer has to be named by a later rule. + """ + _, modules = _quantize(with_mtp=True) + + enabled = { + name + for name, module in modules.items() + if isinstance(module, TensorQuantizer) and module.is_enabled + } + expected = set() + for which in ("k", "v"): + expected |= { + f"backbone.layers.0.mixer.{which}_bmm_quantizer", + f"mtp.layers.0.mixer.{which}_bmm_quantizer", + } + for kind in ("weight_quantizer", "input_quantizer"): + for proj in ("q_proj", "k_proj", "v_proj", "o_proj"): + expected.add(f"backbone.layers.0.mixer.{proj}.{kind}") + for proj in ("in_proj", "out_proj"): + expected.add(f"backbone.layers.1.mixer.{proj}.{kind}") + for prefix in ("backbone.layers.2", "mtp.layers.1"): + for expert in ("experts.0", "experts.1", "shared_experts"): + for proj in ("up_proj", "down_proj"): + expected.add(f"{prefix}.mixer.{expert}.{proj}.{kind}") + expected.add(f"lm_head.{kind}") + + assert enabled == expected