Skip to content

Add a NemotronH aggressive-mse NVFP4 PTQ recipe #2703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
juhi10071998 merged 7 commits into from
Oct 12, 2026
Merged

Add a NemotronH aggressive-mse NVFP4 PTQ recipe #2703

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
46c33a8
fix(hf_ptq): don't require extract_feature during Nemotron VL calibra…
juhi10071998 Oct 8, 2026
5541796
recipe: add a NemotronH aggressive-mse NVFP4 PTQ recipe
juhi10071998 Oct 8, 2026
669db85
fix(hf_ptq): harden the Nemotron VL calibration fallback, and add tests
juhi10071998 Oct 8, 2026
2b177e3
fix(hf_ptq): route omni wrappers through the calibration fallback, an…
juhi10071998 Oct 8, 2026
c306a93
fix(hf_ptq): align both index masks before indexing in the sharded VL…
juhi10071998 Oct 8, 2026
b874d7b
fix(recipe): keep the omni audio branch out of the NemotronH aggressi…
juhi10071998 Oct 9, 2026
5b6fe12
docs: add a CHANGELOG entry for the NemotronH aggressive-MSE NVFP4 re…
juhi10071998 Oct 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
26 changes: 20 additions & 6 deletions examples/hf_ptq/example_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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)
Expand Down
55 changes: 54 additions & 1 deletion examples/hf_ptq/nemotron_vl_calib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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"):

@cjluo-nv cjluo-nv Oct 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

The direct-helper tests reach this branch, but the production caller in create_vlm_calibration_loop() invokes safe_nemotron_vl_forward() only when the wrapper has img_context_token_id. The omni wrappers described here as lacking that attribute instead execute full_model(**call_kwargs), bypassing both this dtype-carrying fallback and its use_cache=False safeguard. Update the caller's dispatch to include the intended omni family without rerouting unrelated encoder-decoder VLMs, and test through the returned calibration loop using a wrapper without img_context_token_id.

@juhi10071998 juhi10071998 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 74615bd.

You're right that the direct-helper tests were exercising a path the production caller never
took. create_vlm_calibration_loop() now computes the dispatch once, before the batch loop:

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"))
)

Two deliberate narrowings:

  • not hasattr(full_model, "extract_feature") keeps the new arm pointed at exactly the branch
    the fallback implements. A Nemotron VL model that has extract_feature but no
    img_context_token_id would otherwise be sent into the InternVL branch and hit an
    AttributeError on full_model.img_context_token_id.
  • not is_enc_dec keeps Nemotron-Parse on the plain forward. It satisfies is_nemotron_vl(),
    but by the time dispatch runs its batch has been renamed to decoder_input_ids;
    safe_nemotron_vl_forward() looks for input_ids, would return without a forward pass, and
    calibration would be silently skipped rather than loudly broken.

tests/examples/hf_ptq/test_nemotron_vl_calib.py now drives the returned calibration loop, not
just the helper: test_calibration_loop_routes_an_omni_wrapper_through_the_fallback asserts
the wrapper's forward received use_cache=False and bf16 pixel_values, and
test_calibration_loop_keeps_the_plain_forward_for_encoder_decoder_vlms asserts the enc-dec
wrapper got decoder_input_ids and fp32 pixel_values — i.e. the fallback was not applied.

# 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)
Expand All @@ -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]

@cjluo-nv cjluo-nv Oct 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

[P2] The earlier cross-device-mask concern is still unresolved. image_flags_s stays on the batch device, so vit_embeds[image_flags_s == 1] can fail before the new embedding transfer when the vision tower is on another GPU. Likewise, selected is built on input_ids.device but indexes flat_embeds on the embedding-table device; the exception handler retries with the same mismatched mask. Move image_flags_s to vit_embeds.device before filtering, and move selected to flat_embeds.device before either assignment. Add a sharded InternVL test exercising both mismatches; the new omni tests bypass this branch entirely.

@juhi10071998 juhi10071998 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in db20808 — the concern was still live, and the alignment I added in 74615bd sat after both indexing operations, so it never applied to either mask. Each mask now moves to the device of the tensor it indexes, immediately before the indexing: image_flags_s → vit_embeds.device before the filter, selected → flat_embeds.device before the merge (which also covers the except-branch retry you called out).

On the sharded test: 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. A CPU-backed torch.Tensor subclass reports a fake device and reproduces CUDA's refusal to index with a mask from another device. Four cases: both mismatches together, the retry path (vision rows deliberately outnumber the selected positions so the first assignment raises on shape), a case where only the embedding table moves — which isolates selected, since otherwise the image_flags_s error fires first and execution never reaches it — and a single-device control.

Evidence they bite, run against the pre-fix code in a throwaway copy of the tree:

A) pre-fix:  3 failed, 1 passed   (control passes)
     image filter      RuntimeError: indices should be ... (got cuda:0 vs cuda:1)  :111
     except-retry      RuntimeError: indices should be ... (got cuda:0 vs cuda:2)  :121
B) fixed:   11 passed             (7 existing + 4 new)
C) recipe:  487 passed            (unchanged)

Worth noting for the day-0 omni path specifically: those wrappers have no extract_feature, so they early-return through the fallback and never reach this branch — which is exactly why the omni tests bypassed it, as you said. This fix is for the InternVL-style checkpoints under a sharded device_map.

# The merge target lives on the embedding device; bring the filtered vision embeddings over.
if vit_embeds.device != flat_embeds.device:

@Edwardf0t1 Edwardf0t1 Oct 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Align the index masks before the cross-device transfer. With a sharded device_map, image_flags_s remains on the batch device, so vit_embeds[image_flags_s == 1] can fail before this new transfer if the vision tower runs on another GPU. Likewise, selected remains on input_ids.device while it indexes flat_embeds on the embedding device; the except block retries with the same mask. Please move each mask to the device of the tensor it indexes before indexing.

@juhi10071998 juhi10071998 Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right, and thanks — the transfer I added in 74615bd was in the wrong place. It moved vit_embeds after both indexing operations, so neither mask was ever aligned. Fixed in db20808: each mask now moves to the device of the tensor it indexes, immediately before the indexing.

  • image_flags_s → vit_embeds.device, before vit_embeds[image_flags_s == 1]
  • selected → flat_embeds.device, before the merge
  • vit_embeds → flat_embeds.device after filtering, as before

The point about the except branch was the sharper one: it retried with the same selected, so it raised a second time rather than recovering. Aligning selected once, up front, covers both the fast path and the retry.

Verified by running the new tests against the pre-fix code in a throwaway copy of the tree — 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, 487 in tests/unit/recipe/.

vit_embeds = vit_embeds.to(flat_embeds.device)
try:
flat_embeds[selected] = flat_embeds[selected] * 0.0 + vit_embeds.reshape(-1, c)
except Exception:
Expand All @@ -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,
Expand Down
150 changes: 150 additions & 0 deletions modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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'

@Edwardf0t1 Edwardf0t1 Oct 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Scope these FP8 attention patterns to the language model. They also match q/k/v/o projections in the optional Parakeet audio encoder (audio_tower in the integrated wrapper, sound_encoder in the released remote-code wrapper). The hf_ptq image-calibration path supplies no audio inputs, so those enabled input quantizers never collect an activation amax; HF export skips the input scale when amax is absent. This leaves the audio branch with unintended, uncalibrated quantization even though the image-only run succeeds. Please exclude the audio encoder after the enabling rules (or narrow the patterns), and add an audio-branch case to the recipe selection test. Parakeet projection names: https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/parakeet/modeling_parakeet.py#L288-L301

@juhi10071998 juhi10071998 Oct 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks Zhiyu, the bf16 ckpt only contains the vision and the language model
https://huggingface.co/nvidia/Nemotron-3.5-Super-VL-120B-A12B-GA-candidate-MTP-boosted-20261001/blob/main/model.safetensors.index.json

I see @meenchen has an exclusion for the audio, speech and audio_projector here https://github.com/NVIDIA/Model-Optimizer/pull/2582/changes#r4232236633:~:text=%2B-,disabled_layers,-%3A

Would adding a similar disabling suffice to ensure that the q/k/v/o projections are scoped in the language model only?

@shengliangxu shengliangxu Oct 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is a model_type specific recipe, this nemotron_h model type does not seem to have audio_tower and sound_encoder, so no need to add these.

@juhi10071998 juhi10071998 Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it seems like the model had the support but the config declared null in this case

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

@cjluo-nv cjluo-nv Oct 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

The existing tree-wide recipe tests validate loading and algorithm construction, not this recipe's effective quantizer selection/precedence. Please add a focused NemotronH regression test (following test_minimax_m3_recipe.py) that applies the resolved config to representative backbone and MTP names. Assert backbone K/V use constant amax with no _amax buffer, MTP K/V remain enabled and calibrated, block-1 routed/shared experts use static NVFP4 weights and dynamic NVFP4 inputs, and vision/router/latent modules stay disabled. Include a no-MTP case. This protects the intentional shipped-checkpoint KV asymmetry from silent selector or ordering regressions.

@juhi10071998 juhi10071998 Oct 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in 74615bd as tests/unit/recipe/test_nemotron_h_recipe.py, following
test_minimax_m3_recipe.py: a synthetic hub-named tree (backbone.layers.N.mixer.*,
mtp.layers.{0,1}.mixer.*, lm_head, vision_model.radio_model.*, embed_vision), the
resolved config applied with algorithm = None, then assertions on the effective quantizers.

Ten cases:

  • backbone K/V — enabled FP8 E4M3, _use_constant_amax is True, and not hasattr(q, "_amax"),
    so no k_scale/v_scale can be exported. Parametrized with and without MTP.
  • MTP block-0 K/V — enabled FP8 and _use_constant_amax is False, which is what the restated
    rule after *mtp.* buys.
  • expert formats — routed and shared experts NVFP4 E2M1, 16-element blocks, FP8 scales,
    type: static on weights and type: dynamic on inputs, for both backbone.layers.2 and
    mtp.layers.1. FP8 on q/k/v/o, Mamba in/out_proj and lm_head.
  • BF16 exclusions — embeddings, both routers, embed_vision, and the C-RADIO tower.
  • MTP attention projections stay BF16 (*mtp.* disables them and nothing after re-enables).
  • test_mtp_rules_are_inert_without_an_mtp_tail compares the full non-MTP quantizer state
    between a tree with and without the tail and requires it identical.
  • test_no_quantizer_outside_the_intended_set_is_enabled asserts the exact set of enabled
    quantizers, so a widened glob fails here before it reaches a checkpoint.

One thing worth recording, since it bit me: the KV quantizers cannot be set in the mock's
__init__. A model that already holds a TensorQuantizer makes is_quantized(model) true,
mtq.quantize then skips apply_mode entirely, no nn.Linear is converted, and the run dies
in _check_weight_quantization_took_effect. They now arrive through QuantModuleRegistry via a
registered _QuantAttention stub, the way register_hf_attentions_on_the_fly installs them on a
real checkpoint, with an autouse fixture that unregisters afterwards.

Verified: 7 passed in tests/examples/hf_ptq/test_nemotron_vl_calib.py, 487 passed in
tests/unit/recipe/ (477 pre-existing + 10 new), zero failures.

- 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
20 changes: 18 additions & 2 deletions modelopt_recipes/ptq.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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` |
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions tests/_test_utils/examples/hf_ptq_example_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""

Expand All @@ -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"]
Loading
Loading