From 46c33a8c74967cdfbf6047969aae14eb21ba466e Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 8 Oct 2026 04:55:41 +0000 Subject: [PATCH 1/7] fix(hf_ptq): don't require extract_feature during Nemotron VL calibration nemotron_vl_calib.py calls full_model.extract_feature(pixel_values) unconditionally. InternVL-style Nemotron VL exposes that helper, but the omni wrappers (NemotronH_Omni_Reasoning_V3 / nemotron_h_omni) do not: they merge vision embeddings with masked_scatter inside their own forward. Calibrating such a model raises AttributeError and aborts before any statistics are collected. Fall back to the model's own forward when the helper is absent. Also align three tensors that can sit on different devices once the model spans GPUs (any sharded device_map, including --use_seq_device_map): vit_embeds against the embedding table before the merge, and attention_mask / position_ids against the LLM's first block. On a single-device model these are no-ops. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/nemotron_vl_calib.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/nemotron_vl_calib.py b/examples/hf_ptq/nemotron_vl_calib.py index 398ef67fe5b..f1fbcdf5b6a 100644 --- a/examples/hf_ptq/nemotron_vl_calib.py +++ b/examples/hf_ptq/nemotron_vl_calib.py @@ -76,9 +76,20 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) flat_ids = input_ids.reshape(b * n) selected = flat_ids == full_model.img_context_token_id - # Vision embeddings + # Vision embeddings. InternVL-style Nemotron VL exposes `extract_feature(pixel_values)`, but + # 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. Calling the helper + # unconditionally raises AttributeError on those, aborting calibration. Fall back to the model's + # own forward when the helper is absent. + if not hasattr(full_model, "extract_feature"): + full_model(**{k: v for k, v in batch.items() if v is not None}) + return vit_embeds = full_model.extract_feature(pixel_values) vit_embeds = vit_embeds[image_flags_s == 1] + # Under a multi-GPU device_map the vision tower and the LLM embedding table can live on + # different devices, so the merge below would raise a cross-device error. Align explicitly. + 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 +99,14 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) inputs_embeds = flat_embeds.reshape(b, n, c) + # Under a sharded device_map the LLM's first block may not sit on the embedding device. + _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, From 55417962ef42f429b76237053dbf24ff54df084d Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 8 Oct 2026 06:59:07 +0000 Subject: [PATCH 2/7] recipe: add a NemotronH aggressive-mse NVFP4 PTQ recipe modelopt_recipes has no architecture-tier PTQ recipe for NemotronH, the hybrid Mamba-MoE family (including the omni VL wrappers). This is the aggressive tier with weight-MSE calibration and an FP8 scale sweep: 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. The MTP rules are inert on checkpoints without an MTP tail. KV cache is FP8 throughout but asymmetric by design: the 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. The MTP rule restates use_constant_amax: false so it cannot inherit the backbone setting whichever way cfg blocks compose. All three numerics imports (nvfp4, nvfp4_static, fp8) and the MTP construction this relies on are already upstream, so the recipe adds no code. ptq.md gains the matching entry, which tests/unit/recipe/test_recipe_docs.py requires for every model_type//ptq/ directory. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- .../nemotron_h/ptq/nvfp4-aggressive-mse.yaml | 138 ++++++++++++++++++ modelopt_recipes/ptq.md | 20 ++- 2 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml 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..5076647ddb2 --- /dev/null +++ b/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml @@ -0,0 +1,138 @@ +# 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 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 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 From 669db851f3bd87a6b94718e30cf5720122862a48 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 8 Oct 2026 07:24:32 +0000 Subject: [PATCH 3/7] fix(hf_ptq): harden the Nemotron VL calibration fallback, and add tests Review follow-ups on the extract_feature fallback. The check now runs before anything InternVL-specific. It previously sat after `full_model.img_context_token_id` and the embedding lookup, so a wrapper lacking `extract_feature` would usually also lack `img_context_token_id` and raise there first -- the fallback was unreachable for most of the models it exists for. The fallback carries the dtype-cast pixel_values rather than the raw batch entry, and sets use_cache=False so calibration does not build a KV cache. Both go through _accepted_kwargs(), which drops keys the wrapper's forward cannot accept unless it takes **kwargs; without it, a narrow forward signature would trade the original AttributeError for a TypeError. A synthesized image_flags is deliberately NOT injected. It is built on pixel_values.device, while a sharded wrapper runs its vision tower on another GPU and indexes image_embeds with it, raising "indices should be either on cpu or on the same device as the indexed tensor". A batch-supplied image_flags is device-consistent and still passes through. This was caught only by an end-to-end 4-GPU run, not by the unit tests, whose mock is single-device; a regression test now pins it. Reworded the device-alignment comment: the code aligns the mask and position ids to the embedding device (where inputs_embeds comes from), not to the first decoder block's as the comment claimed; accelerate's hooks handle later blocks. Tests: tests/examples/hf_ptq/test_nemotron_vl_calib.py, 5 cases. The mock defines neither img_context_token_id nor language_model, so reaching for either regresses the test. Verified end to end: NVFP4 PTQ with --calib_size 32 on 4x184 GiB completed rc=0 in 26m59s, where the pre-fix run died at 6m29s on the device error above. Export structure unchanged: 169,569 tensors, 0 duplicates, 42,179 NVFP4 scales, 4,120 MTP tensors, 8 shards, 457 vision tensors, 69.86 GiB, tensor-count delta 0 against the previous export. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/nemotron_vl_calib.py | 43 +++++-- .../examples/hf_ptq/test_nemotron_vl_calib.py | 119 ++++++++++++++++++ 2 files changed, 153 insertions(+), 9 deletions(-) create mode 100644 tests/examples/hf_ptq/test_nemotron_vl_calib.py diff --git a/examples/hf_ptq/nemotron_vl_calib.py b/examples/hf_ptq/nemotron_vl_calib.py index f1fbcdf5b6a..e580e991644 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) @@ -76,14 +107,6 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) flat_ids = input_ids.reshape(b * n) selected = flat_ids == full_model.img_context_token_id - # Vision embeddings. InternVL-style Nemotron VL exposes `extract_feature(pixel_values)`, but - # 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. Calling the helper - # unconditionally raises AttributeError on those, aborting calibration. Fall back to the model's - # own forward when the helper is absent. - if not hasattr(full_model, "extract_feature"): - full_model(**{k: v for k, v in batch.items() if v is not None}) - return vit_embeds = full_model.extract_feature(pixel_values) vit_embeds = vit_embeds[image_flags_s == 1] # Under a multi-GPU device_map the vision tower and the LLM embedding table can live on @@ -99,7 +122,9 @@ def safe_nemotron_vl_forward(full_model: torch.nn.Module, batch: dict[str, Any]) inputs_embeds = flat_embeds.reshape(b, n, c) - # Under a sharded device_map the LLM's first block may not sit on the embedding device. + # `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: 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..b2f71e2a4fd --- /dev/null +++ b/tests/examples/hf_ptq/test_nemotron_vl_calib.py @@ -0,0 +1,119 @@ +# 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``.""" + +import importlib +from pathlib import Path + +import pytest +import torch + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" + + +@pytest.fixture +def calib(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + return importlib.import_module("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): + super().__init__() + self.calls: list[dict] = [] + + def forward(self, **kwargs): + self.calls.append(kwargs) + return torch.zeros(1) + + +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(calib): + model = _OmniWrapper() + 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(calib): + """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) + 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(calib): + """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() + 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(calib): + def takes_anything(**kwargs): + pass + + payload = {"a": 1, "b": 2} + assert calib._accepted_kwargs(takes_anything, payload) == payload + + +def test_accepted_kwargs_filters_a_fixed_signature(calib): + def takes_only_a(a=None): + pass + + assert calib._accepted_kwargs(takes_only_a, {"a": 1, "b": 2}) == {"a": 1} From 2b177e34b9c2f1643aa24b875535ed785077b946 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 8 Oct 2026 14:50:54 +0000 Subject: [PATCH 4/7] fix(hf_ptq): route omni wrappers through the calibration fallback, and cover the recipe Review follow-ups. `create_vlm_calibration_loop()` reached `safe_nemotron_vl_forward()` only for wrappers carrying `img_context_token_id`. The omni wrappers the fallback exists for have neither that nor `extract_feature`, so in the production calibration loop they still went to `full_model(**call_kwargs)` -- without the dtype-cast `pixel_values` or the `use_cache=False` the fallback supplies. The dispatch now also covers a Nemotron VL model with no `extract_feature`, and explicitly excludes encoder-decoder VL models (Nemotron-Parse): their batch is renamed to `decoder_input_ids`, which the helper does not recognise, so routing them through it would return without a forward pass and skip calibration silently. Tests: - `tests/examples/hf_ptq/test_nemotron_vl_calib.py` now drives the returned calibration loop for both an omni wrapper and an encoder-decoder wrapper, in addition to calling the helper directly, and imports the example modules at collection time through `_test_utils.examples.hf_ptq_example_utils` instead of inside a fixture. - `tests/unit/recipe/test_nemotron_h_recipe.py` pins the NemotronH recipe's effective selection on a synthetic hub-named tree: the KV asymmetry (backbone cast at a fixed range with no `_amax` buffer, MTP calibrated), static-NVFP4 expert weights with dynamic NVFP4 inputs for both the backbone and MTP block 1, FP8 projections and `lm_head`, BF16 vision / router / embeddings, and an exact enabled-quantizer set. Parametrized with and without an MTP tail, so the tail rules are shown to be inert on checkpoints that lack it. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/example_utils.py | 26 +- .../examples/hf_ptq_example_utils.py | 7 +- .../examples/hf_ptq/test_nemotron_vl_calib.py | 95 ++++-- tests/unit/recipe/test_nemotron_h_recipe.py | 289 ++++++++++++++++++ 4 files changed, 385 insertions(+), 32 deletions(-) create mode 100644 tests/unit/recipe/test_nemotron_h_recipe.py 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/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 index b2f71e2a4fd..04c7caa77f2 100644 --- a/tests/examples/hf_ptq/test_nemotron_vl_calib.py +++ b/tests/examples/hf_ptq/test_nemotron_vl_calib.py @@ -13,21 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The Nemotron VL calibration fallback for wrappers without ``extract_feature``.""" +"""The Nemotron VL calibration fallback for wrappers without ``extract_feature``. -import importlib -from pathlib import Path - -import pytest -import torch - -_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" +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 -@pytest.fixture -def calib(monkeypatch): - monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) - return importlib.import_module("nemotron_vl_calib") +import torch +from _test_utils.examples.hf_ptq_example_utils import example_utils, nemotron_vl_calib class _OmniWrapper(torch.nn.Module): @@ -37,8 +34,9 @@ class _OmniWrapper(torch.nn.Module): either before the fallback fires is the regression this guards. """ - def __init__(self): + def __init__(self, config=None): super().__init__() + self.config = config self.calls: list[dict] = [] def forward(self, **kwargs): @@ -46,6 +44,15 @@ def forward(self, **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), @@ -55,9 +62,9 @@ def _batch(): } -def test_fallback_forwards_batch_when_extract_feature_is_absent(calib): +def test_fallback_forwards_batch_when_extract_feature_is_absent(): model = _OmniWrapper() - calib.safe_nemotron_vl_forward(model, _batch()) + 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] @@ -75,17 +82,17 @@ def test_fallback_forwards_batch_when_extract_feature_is_absent(calib): assert "image_flags" not in passed -def test_fallback_forwards_image_flags_the_batch_already_carried(calib): +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) - calib.safe_nemotron_vl_forward(model, batch) + 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(calib): +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): @@ -98,22 +105,64 @@ def forward(self, input_ids=None, pixel_values=None): return torch.zeros(1) model = _Narrow() - calib.safe_nemotron_vl_forward(model, _batch()) # must not raise + 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(calib): +def test_accepted_kwargs_passes_everything_to_a_var_keyword_callable(): def takes_anything(**kwargs): pass payload = {"a": 1, "b": 2} - assert calib._accepted_kwargs(takes_anything, payload) == payload + assert nemotron_vl_calib._accepted_kwargs(takes_anything, payload) == payload -def test_accepted_kwargs_filters_a_fixed_signature(calib): +def test_accepted_kwargs_filters_a_fixed_signature(): def takes_only_a(a=None): pass - assert calib._accepted_kwargs(takes_only_a, {"a": 1, "b": 2}) == {"a": 1} + 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 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..388e4ec106d --- /dev/null +++ b/tests/unit/recipe/test_nemotron_h_recipe.py @@ -0,0 +1,289 @@ +# 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.""" + + def __init__(self): + 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) + + +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 tower, 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 + 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()]) + 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_router_and_embeddings_stay_bf16(with_mtp): + """Everything the blanket `*` disable is meant to leave alone stays unquantized.""" + _, 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", + ] + 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 From c306a93d0b8c64d28f028c476aa7690050b5eca1 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 8 Oct 2026 19:00:06 +0000 Subject: [PATCH 5/7] fix(hf_ptq): align both index masks before indexing in the sharded VL path `safe_nemotron_vl_forward()` built two bool masks on the batch device and then used them to index tensors that a sharded `device_map` can place elsewhere: - `image_flags_s` indexes `vit_embeds`, which comes back on the vision tower's device - `selected` indexes `flat_embeds`, which comes off the embedding table The previous commit moved `vit_embeds` to the embedding device, but did so *after* both indexing operations, so neither mask was ever aligned. The `except` branch retries with the same mismatched `selected`, so it raised a second time rather than recovering. Move each mask to the device of the tensor it indexes, immediately before the indexing. Regression tests simulate the topology rather than allocate it, as this group is CPU-only (the convention `test_example_utils.py` already uses for `torch.cuda.device_count`): a CPU-backed tensor subclass reports a fake device and reproduces CUDA's refusal to index with a mask from another device. Verified against the pre-fix code in a throwaway tree copy: the three sharded cases fail there -- `(got cuda:0 vs cuda:1)` on the image filter and `(got cuda:0 vs cuda:2)` inside the except-branch retry -- while the single-device control passes. On the fixed tree, 11 passed in tests/examples/hf_ptq/test_nemotron_vl_calib.py and 487 in tests/unit/recipe/. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/nemotron_vl_calib.py | 13 +- .../examples/hf_ptq/test_nemotron_vl_calib.py | 178 ++++++++++++++++++ 2 files changed, 189 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/nemotron_vl_calib.py b/examples/hf_ptq/nemotron_vl_calib.py index e580e991644..cffab639401 100644 --- a/examples/hf_ptq/nemotron_vl_calib.py +++ b/examples/hf_ptq/nemotron_vl_calib.py @@ -106,11 +106,20 @@ 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) 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] - # Under a multi-GPU device_map the vision tower and the LLM embedding table can live on - # different devices, so the merge below would raise a cross-device error. Align explicitly. + # 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: diff --git a/tests/examples/hf_ptq/test_nemotron_vl_calib.py b/tests/examples/hf_ptq/test_nemotron_vl_calib.py index 04c7caa77f2..d0e45a3c651 100644 --- a/tests/examples/hf_ptq/test_nemotron_vl_calib.py +++ b/tests/examples/hf_ptq/test_nemotron_vl_calib.py @@ -166,3 +166,181 @@ def test_calibration_loop_keeps_the_plain_forward_for_encoder_decoder_vlms(): 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") From b874d7b9800c481597c3ce16c9170e6206eebd5a Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Fri, 9 Oct 2026 18:10:43 +0000 Subject: [PATCH 6/7] fix(recipe): keep the omni audio branch out of the NemotronH aggressive recipe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FP8 attention rules and the KV cast rule are not language-model scoped. On an omni checkpoint whose config carries a `sound_config`, the wrapper builds a Parakeet encoder, and `*[qkvo]_proj.{weight,input}_quantizer` matches its `self_attn.{q,k,v,o}_proj` — and `relative_k_proj` too, since the character class lands on the `k`. Its attention class ends in `Attention`, so `register_hf_attentions_on_the_fly` registers it for KV-cache quantization exactly like any decoder attention, and `*[kv]_bmm_quantizer` reaches it as well. That leaves the audio branch enabled but never calibrated: the hf_ptq image-calibration path supplies no audio, so those input quantizers collect no amax, and `unified_export_hf` registers `input_scale` only when the amax is present. The weight scale is data-independent and still gets written, so the layer exports as FP8 weights with no activation scale while the image-only run reports success — wrong at inference, silent at export time. Exclude the branch after the enabling rules, matching how the vision tower is already handled and how `diffusion_gemma` handles its never-calibrated self-conditioning network. The test skeleton now carries an attention block in both the vision tower and the audio branch, plus the extra Parakeet projection, so `test_no_quantizer_outside_the_intended_set` covers both; the vision KV quantizers had no coverage before either, and a constant-amax cast exports no scale, so no artefact check would have caught them being left on. Addresses https://github.com/NVIDIA/Model-Optimizer/pull/2703#discussion_r4232413041 Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- .../nemotron_h/ptq/nvfp4-aggressive-mse.yaml | 14 ++++++- tests/unit/recipe/test_nemotron_h_recipe.py | 41 +++++++++++++++---- 2 files changed, 47 insertions(+), 8 deletions(-) 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 index 5076647ddb2..9445dd3b25f 100644 --- a/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml +++ b/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml @@ -23,7 +23,7 @@ # 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 BF16 +# 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 @@ -136,3 +136,15 @@ quantize: 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/tests/unit/recipe/test_nemotron_h_recipe.py b/tests/unit/recipe/test_nemotron_h_recipe.py index 388e4ec106d..b29dd01795f 100644 --- a/tests/unit/recipe/test_nemotron_h_recipe.py +++ b/tests/unit/recipe/test_nemotron_h_recipe.py @@ -37,14 +37,20 @@ class _Attention(nn.Module): - """A NemotronH attention mixer.""" + """A NemotronH attention mixer. - def __init__(self): + ``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): @@ -97,17 +103,26 @@ def __init__(self, mixer): def _nemotron_h(with_mtp: bool) -> nn.Module: - """Hub-named NemotronH skeleton: attention, Mamba and MoE blocks, vision tower, MTP tail.""" + """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 + # 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()]) + 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() @@ -223,8 +238,14 @@ def test_mtp_block_one_experts_match_the_backbone_formats(): @pytest.mark.parametrize("with_mtp", [True, False]) -def test_vision_router_and_embeddings_stay_bf16(with_mtp): - """Everything the blanket `*` disable is meant to leave alone stays unquantized.""" +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 = [ @@ -235,6 +256,12 @@ def test_vision_router_and_embeddings_stay_bf16(with_mtp): "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" From 5b6fe12cf3e86ea89bb9daf72a0ffc9b9051e1f4 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Fri, 9 Oct 2026 19:23:30 +0000 Subject: [PATCH 7/7] docs: add a CHANGELOG entry for the NemotronH aggressive-MSE NVFP4 recipe The recipe is a new user-facing feature, so it gets an entry under *Quantization* alongside the existing Nemotron-H PTQ line. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Juhi Mittal --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) 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