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, 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 33c43e42359..6ec36f218d7 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` | @@ -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