Skip to content

Commit 022767c

Browse files
authored
Add Muse Glimmer AutoQuant recipe (#2269)
### What does this PR do? Type of change: ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> <!-- Details about the change. --> ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing <!-- Mention how have you tested your change if applicable. --> ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A <!--- Mandatory --> - Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A <!--- Very short summary of changes only for new features, backward breaking changes, deprecations, or fixes for critical bugs present in previous releases. --> - Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run `/claude review`. NVIDIA org members can self-trigger for complex changes; orthogonal to CodeRabbit. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added a Muse Glimmer mixed-precision AutoQuantize recipe. * Supports W4A16 NVFP4 Four-Over-Six and FP8 quantization with BF16 fallback at 5.5 effective bits. * Keeps vision modules and the key-value cache unquantized. * Added an Alpamayo QAD example with FSDP2 support and optional full-checkpoint export. * **Documentation** * Added calibration, export, target-module, and fallback-precision guidance for the recipe. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Jenny Chen <jennifchen@nvidia.com>
1 parent 72e48d5 commit 022767c

3 files changed

Lines changed: 101 additions & 0 deletions

File tree

CHANGELOG.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ Changelog
88

99
*Quantization*
1010

11+
- Add a Muse Glimmer AutoQuantize recipe that searches language-model MLP projections, self-attention projections, and ``lm_head`` over W4A16 NVFP4 Four-Over-Six, FP8, and BF16 fallback at 5.5 effective bits while leaving the vision tower unquantized.
1112
- Add ``examples/alpamayo/qad.py``, which runs quantization-aware distillation on the quantized Alpamayo checkpoint produced by ``examples/alpamayo/quantize.py``. It distills the quantized VLM against the original FP16 VLM with ``QADTrainer``, supports FSDP2 for multi-GPU runs, and ``--export`` reassembles the trained VLM into a full AlpamayoR1 checkpoint that ``AlpamayoR1.from_pretrained`` can reload.
1213
- Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path.
1314
- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported.
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Muse Glimmer AutoQuantize recipes
2+
3+
`w4a16_nvfp4_4o6_mixed` searches the Muse Glimmer
4+
language-model MLP projections, self-attention projections, and `lm_head` over
5+
W4A16 NVFP4 Four-Over-Six, FP8, and BF16 fallback at 5.5 effective bits. The
6+
vision tower and unmatched modules remain BF16.
7+
8+
Use unquantized KV cache and representative text calibration:
9+
10+
```bash
11+
python examples/hf_ptq/hf_ptq.py \
12+
--pyt_ckpt_path <muse-glimmer-checkpoint> \
13+
--recipe models/meta-models/Muse-Glimmer-30B/auto_quantize/w4a16_nvfp4_4o6_mixed \
14+
--kv_cache_qformat none \
15+
--dataset nemotron-post-training-v3 \
16+
--calib_size 512 \
17+
--calib_seq 2048 \
18+
--export_path <output-path>
19+
```
Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
# Muse Glimmer language-model search over W4A16 NVFP4 Four-Over-Six, FP8,
5+
# and BF16 fallback. Vision and unmatched module families remain BF16.
6+
7+
# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
8+
imports:
9+
base_disable_all: configs/ptq/units/base_disable_all
10+
base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
11+
base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers
12+
default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
13+
fp8_quantizers: configs/ptq/units/w8a8_fp8_fp8
14+
nvfp4_four_over_six: configs/numerics/nvfp4_four_over_six
15+
16+
metadata:
17+
recipe_type: auto_quantize
18+
description: >-
19+
Muse Glimmer language-model W4A16 NVFP4 Four-Over-Six, FP8, and BF16
20+
fallback search at 5.5 weight-effective bits. Use --kv_cache_qformat none.
21+
22+
# Unmatched modules retain this BF16-by-default baseline. The searched language-model
23+
# families start from W4A16 NVFP4 and are reassigned by AutoQuantize below.
24+
quantize:
25+
effective_bits: 4.5
26+
algorithm:
27+
method: mse
28+
fp8_scale_sweep: false
29+
start_multiplier: 1.0
30+
stop_multiplier: 1.5
31+
step_size: 0.5
32+
quant_cfg:
33+
- $import: base_disable_all
34+
- quantizer_name: "*language_model.layers.*.mlp*.weight_quantizer"
35+
cfg:
36+
$import: nvfp4_four_over_six
37+
- quantizer_name: "*language_model.layers.*.self_attn.*_proj.weight_quantizer"
38+
cfg:
39+
$import: nvfp4_four_over_six
40+
- quantizer_name: "*lm_head*.weight_quantizer"
41+
cfg:
42+
$import: nvfp4_four_over_six
43+
- $import: default_disabled_quantizers
44+
45+
auto_quantize:
46+
constraints:
47+
effective_bits: 5.5
48+
cost_model: weight
49+
50+
module_search_spaces:
51+
- module_name_patterns:
52+
- "*language_model.layers.*.mlp*"
53+
- "*language_model.layers.*.self_attn.*_proj"
54+
- "*lm_head*"
55+
candidate_formats:
56+
- effective_bits: 4.5
57+
algorithm:
58+
method: mse
59+
fp8_scale_sweep: false
60+
start_multiplier: 1.0
61+
stop_multiplier: 1.5
62+
step_size: 0.5
63+
quant_cfg:
64+
- $import: base_disable_all
65+
- quantizer_name: "*weight_quantizer"
66+
cfg:
67+
$import: nvfp4_four_over_six
68+
- algorithm: max
69+
quant_cfg:
70+
- $import: base_disable_all
71+
- $import: fp8_quantizers
72+
allow_no_quant: true
73+
74+
auto_quantize_method: gradient
75+
score_size: 128
76+
77+
disabled_layers:
78+
- $import: base_disabled_layers
79+
80+
cost_excluded_layers:
81+
- $import: base_cost_excluded_layers

0 commit comments

Comments
 (0)