[Cherry-pick] PRs #1975 #2076 #2071 #2093 #2084 #2115 #2133 #2146 #2064 #2159 #2112 - #2179
Merged
Merged
Conversation
### What does this PR do?
Type of change: new feature
Adds online DFlash training support for Qwen3-VL–style vision-language
models.
Changes include:
- Load VLMs through the Transformers 5 `AutoModelForImageTextToText`
API, while retaining compatibility with the legacy VLM auto-model API.
- Run the base model through its top-level multimodal forward when
image/video inputs are present, ensuring vision embeddings are injected
before collecting DFlash target hidden states.
- Extend `VisionLanguageDataCollator` to:
- propagate `answer_only_loss`, chat-template, and DFlash
label-alignment settings;
- apply `VLM_MIN_PIXELS` / `VLM_MAX_PIXELS` processor limits;
- derive assistant-only masks from ChatML/Llama chat boundaries when
processor generation masks are unavailable;
- enforce the fixed `training_seq_len` required by DFlash block
training.
- Preserve the existing text-only DFlash path.
### Usage
```bash
python -m torch.distributed.run \
--nproc_per_node 4 \
examples/speculative_decoding/main.py \
--config modelopt_recipes/general/speculative_decoding/dflash.yaml \
model.model_name_or_path=/path/to/qwen3-vl-model \
model.trust_remote_code=true \
data.data_path=/path/to/train.jsonl \
data.vlm_processor=/path/to/qwen3-vl-model \
data.vlm_img_dir=/path/to/image/root \
training.training_seq_len=4096 \
training.answer_only_loss=true \
dflash.dflash_block_size=8 \
dflash.dflash_mask_token_id=151669
### Testing
- git diff --check
- Parsed all modified Python modules successfully.
- Ran iterative multi-node Slurm smoke tests with a Qwen3-VL-family model and mixed multimodal data:
- validated VLM model loading with Transformers 5;
- validated distributed initialization, DFlash conversion, and VLM collation paths;
- identified and addressed processor padding/truncation behavior required by fixed-size DFlash blocks.
This PR remains draft pending a completed end-to-end training smoke test and automated regression coverage.
### 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?: ✅
- If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
- Did you write any new necessary tests?: ❌ — automated Qwen3-VL/DFlash regression coverage still needs to be added before review.
- Did you update Changelog (https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — evaluate and add an entry before marking ready for review if this is considered user-facing speculative-decoding support.
- Did you get Claude approval on this PR?: N/A
### Additional Information
The PR intentionally excludes local Slurm launch scripts, logs, model paths, datasets, and environment-specific configuration.
<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit
* **New Features**
* Expanded VLM data-collation controls, including `shift_labels` and more robust `answer_only_loss` masking.
* Improved Qwen3-VL speculative decoding for Transformers 5.3+ with correct video frame grouping and safer position-id handling.
* Improved DFlash RoPE export to reliably read `rope_theta` from newer config formats.
* **Bug Fixes**
* Hardened multimodal preprocessing and training loss masking to keep label/attention alignment consistent.
* Improved behavior when anchor sampling yields no valid blocks.
* More resilient VLM model loading when certain Transformers auto classes are unavailable.
* **Tests**
* Added coverage for RoPE export, Qwen3-VL position-id logic across Transformers versions, and VLM label-mode/collator options.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Co-authored-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…me (NVBug 6525511) (#2076) ## What does this PR do? **Type of change:** Bug fix Fixes [NVBug 6525511](https://nvbugspro.nvidia.com/bug/6525511) / [OMNIML-5599](https://jirasw.nvidia.com/browse/OMNIML-5599): FP8 PTQ of `llava-1.5-13b` on `transformers>=5.12` produces a checkpoint vLLM refuses to load: ``` ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration ``` A cross-architecture sweep run for this PR shows **`google/gemma-3-4b-it` is broken the same way** (884/884 keys mangled) and is fixed by the same change, though no bug was filed for it. ### Root cause transformers collects conversion mappings **recursively** and tags each sub-model's transforms with the sub-module path they belong to, then matches only keys under that prefix (`WeightTransform._scoped_match`: strip `scope_prefix` → match → re-attach): ```python transform.scope_prefix = scope_prefix transform.base_model_prefix = model.base_model_prefix ``` `LlavaForConditionalGeneration` therefore carries the **vision tower's own** `PrefixChange` — "add a `vision_model.` prefix" — scoped to `model.vision_tower`: ``` [5] rev=PrefixChange scope_prefix = 'model.vision_tower' source_patterns = ['^(?:(?!vision_model\.))(.+)$'] target_patterns = ['vision_model.\1'] ``` `_build_reverse_rules` read `rev.source_patterns` / `rev.target_patterns` raw and **discarded `scope_prefix`**. Read as an unscoped regex, that pattern means "any key not already starting with `vision_model.`" — i.e. everything. All 758 llava-1.5-13b tensors were moved under a bogus top-level `vision_model.`: ``` model.language_model.layers.0.self_attn.q_proj.weight -> vision_model.model.language_model.layers.0.self_attn.q_proj.weight lm_head.weight -> vision_model.lm_head.weight ``` This is the same class of defect as #2032 (NVBug 6525534), but that fix's shadowing heuristic cannot catch it: it drops a rule whose target is an **existing** namespace, whereas this rule **invents** one (`vision_model` exists nowhere in the module tree). ### The fix `RenameRule` carries `scope_prefixes`; `_sub_scoped` applies a scoped rule only to keys under one of them — stripping the prefix before the match and re-attaching after, mirroring transformers' own semantics (trying `base_model_prefix + scope_prefix` before `scope_prefix`). The same scoping flows through `build_reverse_name_mapper`, so `exclude_modules` (which lists the BF16 vision tower) stays aligned with the weights. `_drop_shadowed_prefix_renames` skips scoped rules, which are already confined to their subtree. Converter-derived rules (`_expert_leaf_renames`, `_dense_split_rule`) match by **module suffix** rather than an anchored pattern, so they cannot be confined to a subtree the same way. A survey of 9 architectures (LLaVA, LLaVA-Next, Gemma-3, Gemma-4, Qwen2-VL, Qwen3-VL-MoE, Llama-4-Scout, Mixtral, DeepSeek-V2-Lite) found **every `WeightConverter` has `scope_prefix=None`** — transformers only scopes `WeightRenaming`/`PrefixChange` — so the case is unreachable today. Rather than emit rules that could silently reach a sibling namespace if that ever changes, a scoped converter now raises `QuantConversionUnsupportedError` and the caller falls back to in-memory names with a warning. ## Testing **Name-level oracle against the real hub checkpoint.** Exported names must equal the original checkpoint's keys. All 758 `llava-hf/llava-1.5-13b-hf` state-dict keys round-trip exactly: ``` MISSING (hub key not produced): 0 SPURIOUS (name not in hub) : 0 exported: {language_model: 363, vision_tower: 391, multi_modal_projector: 4} hub : {language_model: 363, vision_tower: 391, multi_modal_projector: 4} ``` Pre-fix, 758/758 keys were mangled. This check involves no vLLM. **Cross-architecture regression sweep** — old (unscoped) vs new (scoped) mapping over each model's real state dict: | Result | Models | |---|---| | Identical (no behavior change) | Llama-3.2, Qwen2.5, Qwen3, Mistral, Phi-3, SmolLM2, gpt-oss-20b, DeepSeek-V2-Lite, Mixtral-8x7B, gemma-2-2b, gemma-4-31B, Qwen2-VL-2B | | Differs (fixed) | llava-1.5-7b (686/686), gemma-3-4b-it (884/884) | Note Qwen2-VL carries a scoped rule yet is unchanged — the fix only bites where a scoped rule would have wrongly matched. **End-to-end**, on the exact image from the bug report (`vllm/vllm-openai:v0.26.0`, verified `vllm.__version__ == 0.26.0`, transformers 5.14.1): real FP8 PTQ (`general/ptq/fp8_default-kv_fp8_cast`, `--calib_size 512`) on llava-1.5-13b, then served with the bug's exact `api_server` command. ``` PTQ_EXIT=0 top-level namespaces: {'language_model': 923, 'multi_modal_projector': 4, 'vision_tower': 391} KEYS UNDER BOGUS vision_model.* : 0 occurrences of "no module or parameter named 'vision_model'": 0 Loading weights took 5.58 seconds Model loading took 13.2 GiB and 7.45 seconds ``` The reported failure is gone and weight loading completes. - `pytest tests/unit/torch/export` → **116 passed, 1 skipped** (3 new regression tests) - `pre-commit run --files ...` → all hooks pass ## Additional Information **Out of scope, for whoever picks up the QA ticket.** After this fix the bug's exact repro hits a *different* error: llava-1.5-13b ships `"dtype": "float16"` and vLLM's FP8 kernel requires BF16 output (`RuntimeError: For FP8 input, output must have dtype BF16`). This is not an export defect — it also occurs with `kv_cache_dtype=auto`, i.e. with the FP8 KV cache entirely out of the picture, and the same checkpoint loads and generates correctly under `--dtype bfloat16` (`" Paris. with a population of about 2,249,03"`). QA will need `--dtype bfloat16`. **[NVBug 6525597](https://nvbugspro.nvidia.com/bug/6525597)** (`gemma-4-31B-it`, `assert layer.k_scale > 0.0`) is unrelated to this PR: Gemma4 exposes zero transformers conversions, so this change provably does not touch it (confirmed identical in the sweep above). Re-tested separately on vLLM 0.26.0 with that bug's environment (transformers 5.5.0, TP=1, batch_size 8) it did not reproduce — the server reached `Application startup complete` with zero asserts — but that is tracked outside this PR. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved Hugging Face exports for multimodal models with nested prefix conversions. - Ensured vision-tower prefixes apply only within the correct model scope. - Prevented unrelated language-model and head parameters from being incorrectly renamed. - Preserved scoped behavior for module mappings and wildcard exclusions. - Added safer handling for unsupported scoped weight conversions. - **Tests** - Added regression coverage for scoped prefix handling across weights, module names, exclusions, and unsupported conversion scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: James Shen <yueshen@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…VBug 6518665) (#2071) ### What does this PR do? Type of change: Bug fix Fixes [NVBug 6518665](https://nvbugspro.nvidia.com/bug/6518665) / [OMNIML-5583](https://jirasw.nvidia.com/browse/OMNIML-5583): PTQ of `Step-3.7-Flash` (`--recipe general/ptq/nvfp4_mlp_only-kv_fp8`, transformers 5.5.4) calibrates fine and then dies while writing the checkpoint: ``` File "modelopt/torch/export/unified_export_hf.py", line 1510, in export_hf_checkpoint model.save_pretrained( File "modelopt/torch/opt/plugins/transformers.py", line 155, in _save_pretrained_with_checks File "transformers/modeling_utils.py", line 3352, in save_pretrained state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save) File "transformers/modeling_utils.py", line 338, in _get_tied_weight_keys tied_weight_keys.extend([f"{name}.{k}" if name else k for k in tied.keys()]) AttributeError: 'list' object has no attribute 'keys' ``` **Root cause is version drift in `transformers`, not in quantization.** transformers 5.0 changed `_tied_weights_keys` from `list[str]` to a `{target: source}` dict, and `save_pretrained` → `_get_tied_weight_keys` calls `.keys()` on every submodule's declaration without a type check. `stepfun-ai/Step-3.7-Flash`'s remote code still uses the 4.x list format — `_tied_weights_keys = ["lm_head.weight"]` on `Step3p7TextModel`, `Step3p7Model` and `Step3p7ForConditionalGeneration`. The load-time consumer (`get_expanded_tied_weights_keys`) returns early when `config.tie_word_embeddings` is false, which it is for this checkpoint, so the model loads and calibrates normally and only fails at the end of the run, after the expensive part. The same crash reproduces with a plain `AutoModel.from_pretrained(..., trust_remote_code=True).save_pretrained(...)`, and transformers 5.12 tolerates `None` there but still not a list, so it is not fixed upstream either. ### Fix `_save_pretrained_with_checks` — the entry point every ModelOpt HF save routes through (unified export, `ModelOptHFTrainer`, user-called `save_pretrained`) — now normalizes a list-style declaration to the equivalent dict for the duration of the save and restores the original attribute afterwards. Mapping each entry to itself preserves the legacy semantics: those list entries were exactly the dedup patterns `_get_tied_weight_keys` is expected to return. No-op on `transformers<5`, where the list format is native, and no-op for models that already declare a dict. The restore tracks whether the instance owned the attribute, so for the usual class-level declaration nothing is left shadowing it. ### Usage ```bash python hf_ptq.py --model /local/Step-3.7-Flash --recipe general/ptq/nvfp4_mlp_only-kv_fp8 \ --dataset /local/cnn_dailymail --calib_size 32 --export_path /local/Step-3.7-Flash-nvfp4 --trust_remote_code ``` ### Testing Two new tests in `tests/unit/torch/opt/plugins/test_transformers_save_load.py`: - `test_save_pretrained_with_legacy_tied_weights_keys` (parametrized over `tie_word_embeddings`) — a quantized tiny-Llama declaring list-style keys saves and round-trips; without the fix it fails with the exact reported `AttributeError` at `modeling_utils.py:338`. It asserts the saved `model.safetensors` keys directly, so no weight can be silently dropped by the save-time dedup: `model.embed_tokens.weight` is always present, and `lm_head.weight` is present iff the weights are not actually tied (when they are, transformers drops the alias and re-ties it on load — the transformers 4.x behavior). - `test_legacy_tied_weights_keys_as_dict_restores_class_attribute` — the shim leaves no instance attribute shadowing a class-level declaration (skipped on `transformers<5`). Ran locally against both transformers majors: - transformers 5.5.4 (the version in the bug report) + torch 2.11: `tests/unit/torch/opt/plugins/` 27 passed, `tests/unit/torch/export/` 124 passed. - transformers 4.57.6: `tests/unit/torch/opt/plugins/` 26 passed (the shim no-ops, new dict-restore test skips). Not yet re-run end-to-end on the real Step-3.7-Flash checkpoint (1.4 TB / 8×B200); QA can re-run the reported command against this branch. ### Known limitation The shim covers the save side only. A legacy-list model that *also* sets `tie_word_embeddings=True` still crashes on the **load** side, in `get_expanded_tied_weights_keys` during `PreTrainedModel.post_init` — before ModelOpt has a model object to patch. Such a checkpoint cannot be loaded by transformers 5 at all, with or without ModelOpt, so it is out of scope here. `Step-3.7-Flash` is unaffected (`tie_word_embeddings` is unset, i.e. false). ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information The upstream-correct fix is for StepFun to migrate their modeling code to the transformers 5 dict format; this shim unblocks every not-yet-migrated `trust_remote_code` checkpoint in the meantime. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility with Transformers 5 when exporting models with legacy tied-weight declarations. * Preserved tied-weight settings after saving, including class-level declarations. * Ensured saved models reload with equivalent state and outputs. * **Tests** * Added coverage for quantized models, tied and untied embeddings, and legacy tied-weight formats. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
### What does this PR do? Type of change: Bug fix Fixes calibration failures when running fake-quantization with vLLM 0.26.0. Two root causes are addressed: 1. **`finish_requests` must be called explicitly after `add_requests`.** In vLLM 0.26.0 the scheduler calls `finish_requests` *before* `add_requests` inside `execute_model`, so request IDs are not registered yet and cleanup never runs. The calibration loop now calls `finish_requests` directly after each batch using `dataclasses.replace`. Wrapped in `try/finally` with an inner `try/except` so it always runs and never masks the original exception. A warning is emitted when `finish_requests` is absent so the regression is self-diagnosing on future vLLM API changes. 2. **`NewRequestData` gained a `prefill_token_ids` field.** vLLM 0.26.0 added this required argument; the calibration helper now passes it. Additional: - Dockerfile updated to vLLM 0.26.0 with `USER vllm` (non-root). - README updated to include vLLM 0.26.0 in tested versions. ### Usage ```bash cd examples/vllm_serve QUANT_CFG=FP8_DEFAULT_CFG QUANT_CALIB_SIZE=8 CALIB_BATCH_SIZE=1 \ python3 vllm_serve_fakequant.py Qwen/Qwen1.5-MoE-A2.7B-Chat -tp 1 \ --host 0.0.0.0 --port 8000 ``` ### Testing Tested end-to-end FQ calibration with vLLM 0.26.0 using the Docker image built from `examples/vllm_serve/Dockerfile`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A (examples change only) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information Changes are confined to `examples/vllm_serve/` and do not affect the core library. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Added compatibility with vLLM 0.26.0 in the serving example. * Improved calibration request handling and cleanup after model execution. * Enabled the serving container to run with a non-root user. * **Bug Fixes** * Calibration cleanup failures no longer obscure the original model execution error. * Added warnings when calibration cleanup cannot be completed. * **Documentation** * Updated the serving example documentation to list vLLM 0.26.0 among tested versions. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
) ### What does this PR do? Type of change: Bug fix The ModelOptSparseAttentionImpl.forward (FlashAttention path) assumes the paged KV cache has shape [2, num_blocks, page_size, num_kv_heads, head_dim], where dimension 0 is the K/V split. unbind(0) on this shape returns exactly 2 tensors. Newer vLLM versions changed the FlashAttention kv_cache layout to [num_blocks, 2, page_size, num_kv_heads, head_dim] (same as FlashInfer). unbind(0) now returns num_blocks tensors, causing the unpack error. Now, we detect the layout first and then unbind. ### Usage save as before. ### Testing unittest ### 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?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - 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 * **Bug Fixes** * Improved FlashAttention compatibility with vLLM key/value cache layouts. * Automatically detects supported cache dimension ordering for accurate sparse-attention processing across legacy and newer layouts. * **Tests** * Expanded coverage for supported cache layouts, cache validation, and paged-cache scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
) Type of change: Deprecation Resolves [NVBug 6563509](https://nvbugspro.nvidia.com/bug/6563509), where `hf_ptq.py` on Phi-4-multimodal-instruct died with `RuntimeError: Tensor.item() cannot be called on meta tensors`. The crash is real but not fixable on our side, and it is not the reason the model is unusable. Phi-4-multimodal's bundled remote code predates Transformers v5 and does not load on **any** version in our supported range (`transformers>=4.57,<5.15`): | Blocker | Where | |---|---| | `peft.get_peft_model` reads `prepare_inputs_for_generation`, gone since transformers 4.52 dropped `GenerationMixin` from `PreTrainedModel` | `modeling_phi4mm.py:1959` | | `_tied_weights_keys` declared as a list; Transformers 5.x calls `.keys()` on it in `post_init` | `modeling_phi4mm.py:1937` | | `int(torch.tensor(...))` in `__init__`, which cannot run on a meta device — the reported crash | `speech_conformer_encoder.py:1435` | The model card pins `transformers==4.48.2` / `peft==0.13.2`, so there is no overlap with our floor and nothing on our side can bridge it. The model is therefore dropped rather than worked around. **Phi-3-vision is dropped alongside it because it is the older, superseded model in the same family** — with its successor unsupportable there is no reason to keep carrying the predecessor. This is a product-scope call, not a separate compatibility finding: Phi-3-vision shares the list-valued `_tied_weights_keys` defect (`modeling_phi3_v.py:1214`) and so is likewise broken on Transformers 5.x, but it does **not** hit the `peft` blocker, and it was not re-verified on 4.57. Per the 0.46 changelog we have already bumped the floor to 4.57 and noted that "Transformers 4.x support will be dropped in a future release", so any remaining window closes on its own. Same reasoning already applied to VILA / NVILA in this release. **Removed** - the support-matrix row in `examples/hf_ptq/README.md` - `"Phi4MMForCausalLM": "phi4mm"` from `MODEL_NAME_TO_TYPE` - the multimodal-detection heuristics that only ever matched these two — `vision_lora`, `audio_processor`, `embd_layer.image_embd_layer`, and the `phi4mm` model-type check — in both `is_multimodal_model` and `_is_multimodal_config` - the `Phi3Image` / `PhiImage` exclusions in `is_embedding` - the phi4mm input-mode warning in `hf_ptq.py` - `modelopt_recipes/huggingface/phi4mm/` and its references in `modelopt_recipes/ptq.md` **Not changed:** the device-map sizing path (meta-device skeleton, `infer_auto_device_map`, and the `--gpu_max_mem_percentage` cap) keeps its original behavior. That cap is wanted exactly where it already fires — when the model is already offloading to CPU, where it costs little and the headroom is required. With the affected checkpoints removed, there is no supported model that trips the meta-device build, so there is nothing to work around here. Text-only **Phi-3/Phi-4** and **Phi-3.5-MoE** are natively supported by transformers and are untouched. On H200, `nvcr.io/nvidia/tensorrt-llm/release` (torch 2.12, transformers 5.5.4), against the real checkpoint: - **Version matrix** (vanilla transformers, no modelopt) — Phi-4-MM loads at 4.48.2 / 4.49.0 / 4.50.0 / 4.51.3 and fails at 4.53.3 / 4.56.2 / 4.57.1 (`AttributeError: 'Phi4MMModel' object has no attribute 'prepare_inputs_for_generation'`) and at 5.5.4 (meta-init, then tied-keys). This is what establishes that no supported version works. - `tests/examples/hf_ptq/test_example_utils.py` — 28 passed. - **Sweep**: `tests/examples/hf_ptq` + `tests/unit/torch/export` — failure set identical to the pre-change tree (GPU/model-dependent `test_vlm_ptq`, plus `test_quant_aware_conversion` scoped-mapping tests), so none are introduced here. - `pre-commit` clean on all changed files, including recipe validation. - Is this change backward compatible?: ❌ — PTQ for Phi-3-vision and Phi-4-multimodal is removed, along with the `huggingface/phi4mm/ptq/*` recipes. Phi-4-multimodal is already unloadable on every supported transformers version, so no working workflow regresses; Phi-3-vision is a deliberate scope removal as its superseded predecessor. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A — this is a deletion; the existing `test_get_model_*` / `test_resolve_init_config_*` tests are unchanged and still pass. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ — not yet run. Two related references were left in place deliberately; say the word and I'll fold them in: - `tests/examples/hf_ptq/test_deploy.py` still deploys the already-published `nvidia/Phi-4-multimodal-instruct-{NVFP4,FP8}` checkpoints. Those artifacts exist and serve fine; this PR only removes the ability to *produce* them. - `examples/torch_onnx/README.md` still lists Phi-4-multimodal-instruct. That is a separate ONNX pipeline that does not go through `get_model()` and was not tested here. Earlier revisions of this branch also reworked the device-map sizing so the meta-tensor crash could not occur. That was reverted in 701180e: the guard is correct as written, and every alternative either changed behavior for models that fit today or moved the guard somewhere it does not belong, for a crash that only ever affected the checkpoints this PR removes. --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
### What does this PR do? Type of change: Bug fix Fixes NVBug 6571812 by enabling NVFP4 quantization for NemotronH dense MLP projections registered as `backbone.layers.*.mixer.up_proj` and `backbone.layers.*.mixer.down_proj`. The existing MLP-only recipes enable `*mlp*` quantizers after disabling all quantizers. NemotronH's dense projection names do not contain `mlp`, so no dense MLP weights were selected and HF export emitted `quant_algo: null`. This change: - adds a reusable `mixer_mlp_nvfp4` recipe unit for dynamic NVFP4 - applies it to the MLP-only and OMLP-only recipes and presets - adds equivalent static selectors to the MSE recipe - verifies that dense mixer projections match while Mamba `in_proj`/`out_proj` and `shared_experts` remain excluded - documents the fix in the 0.46 changelog ### Usage ```bash python examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 \ --recipe general/ptq/nvfp4_mlp_only-kv_fp8_cast \ --export_fmt hf \ --export_path /path/to/output ``` ### Testing - `python -m pytest tests/unit/recipe`: 231 passed - `pre-commit run`: all applicable hooks passed - Remote GB300 PTQ/export smoke test: all 34 expected dense MLP projections used NVFP4, with 0 missed or unexpected projections - TensorRT-LLM `1.3.0rc6` TP1 canary: loaded all 445 tensors, `/health` and `/v1/models` returned HTTP 200, and chat completion returned Paris - Slurm job `2900754`: `COMPLETED`, exit `0:0` ### 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?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: N/A ### Additional Information - NVBug 6571812 - Commit `e9b8f9ffc55e66cdbbacdbbb287e0dedf96ab082` is signed and includes the DCO sign-off. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added NVFP4 quantization support for dense mixer MLP projections. - Expanded NVFP4 MLP and OMLP preset recipes, including dynamic input and static-weight configurations. - Improved compatibility across supported quantization recipe combinations. - **Documentation** - Documented the new NVFP4 mixer MLP quantization option. - Added release notes covering calibration, grouped-MoE quantization, logging, integrations, and compatibility updates. - **Tests** - Added coverage validating NVFP4 recipe behavior and projection selection. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Type of change: New example Add Nemotron Lightning 3.5 NVFP4 recipe and QAD example Also exclude MTP in default disabled quantizers ```python ``` <!-- Mention how have you tested your change if applicable. --> 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 <!--- 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. --> <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> - **New Features** - Added an end-to-end NVFP4 quantization, distillation, and export workflow for NVIDIA Nemotron 3.5 Lightning 30B-A3B. - Added a PTQ configuration supporting NVFP4 W4A16 quantization with optimized scaling and FP8 support for selected components. - **Bug Fixes** - Preserved custom model output locations when provided, while retaining the existing default path. - **Configuration** - Disabled quantization for MTP modules by default. - Updated an existing Nemotron workflow to use the aggressive NVFP4 quantization profile. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Type of change: Bug fix Fixes `AssertionError: Model already has modelopt state!` when exporting a QLoRA checkpoint (NVBug 6542481). The QLoRA output is adapter-only, so `from_pretrained` resolves the quantized base model and already restores the ModelOpt state; `export.py` then restored a second time. Fixing that exposed two more breakages on the same path, also fixed here: - `_restore_qtensor_wrappers` missed every module — PEFT renames the compressed linears to `<name>.base_layer`, so no weight got re-wrapped and the packed NVFP4 weight hit a shape error. - `postprocess_state_dict` dropped `weight_scale_2` (missing from the QLoRA rename map), leaving the exported checkpoint impossible to dequantize. No API change — `examples/llm_qat/export.py --pyt_ckpt_path <qlora_ckpt> --export_path <out>` now completes on the documented quantize → train → export flow. Reproduced in the reported environment (TRT-LLM 1.3.0rc22, transformers 5.5.4, NVFP4). - Added the missing export step to `test_qwen3_qlora_nvfp4` and a unit test for the QLoRA `base_layer` rename; both fail without the fix. - Exported base model is byte-identical to a plain PTQ export; dequantized NVFP4 weights match the bf16 original (worst rel. error 0.10). - No regressions: `tests/gpu/torch/export/test_export.py` (49 passed), save/load plugin tests. - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ❌ — can add if wanted - Did you get Claude approval on this PR?: ❌ — not run yet Fixes NVBug 6542481. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> * **Bug Fixes** * Improved QLoRA checkpoint export and restoration across supported model configurations. * Preserved secondary weight-scale information and other deployment tensors in exported checkpoints. * Corrected handling of quantized base-layer weights after adapter reparenting. * Prevented duplicate state restoration when checkpoints already include the required model state. * Removed internal adapter prefixes and quantizer details from exported state data. * **Tests** * Added validation for packed weights, quantization metadata, required scales, and removal of embedded adapter layers. * Added regression coverage for QLoRA state processing and quantized weight restoration. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
#2159) ### What does this PR do? Type of change: Bug fix + new feature Two model families that could not be pruned end-to-end now can: - **Nemotron-3.5-Lightning** (`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`) — a native `NemotronHForCausalLM` that ships without remote code and carries MTP heads. Fixes a calibration crash and HF-export failures on the modern Megatron-Bridge / transformers stack. - **DeepSeek-V3** — fixes an MLA Q-LoRA crash during calibration, and adds a `candidate_filter` search option to `mcore_minitron` so its MoE-FFN dimensions stay prunable while remaining representable in HF. Also makes a rank-local failure under pipeline parallelism fail fast instead of stalling. #### 1. Nemotron Lightning: prune + HF export (`examples/megatron_bridge/prune_minitron.py`) 1. **MTP calibration crash.** On newer Megatron-LM, `mtp_process` is derived from the hybrid *pattern*, not from `mtp_num_layers`. Setting only `mtp_num_layers=0` in the calibration provider overrides was insufficient: the provider's `finalize()` re-appended the MTP suffix to `hybrid_layer_pattern` (because `mtp_hybrid_override_pattern` was still set and `mtp_use_repeated_layer=True`), so `mtp_process=True` while `mtp_num_layers=0` and the calibration forward hit `assert self.config.mtp_num_layers > 0`. Fix: also clear `mtp_hybrid_override_pattern` in the calibration overrides so MTP is fully disabled (MTP heads are dropped from the pruned model, as before). 2. **HF export via a config-only bridge (hybrid models only).** The old export built a *dummy* HF model to obtain the bridge, then streamed weights. This breaks on native NemotronH because (a) native `NemotronHConfig` makes `hybrid_override_pattern` a read-only property, and (b) transformers 5.12 saves the input embedding under a different key than the bridge mapping expects (`backbone.embedding` vs `backbone.embeddings`); the mismatch made `build_conversion_tasks` drop the embedding task on its owning rank, leaving an owner-less PP placeholder that crashed `save_hf_weights` with `Object must exist on at least one PP rank`. Fix: stream weights through a **config-only** bridge (`AutoBridge.from_hf_config(hf_cfg).save_hf_pretrained(...)`), available since Megatron-Bridge 0.5.0 (nemo:26.06). A config-only bridge has `hf_keys=None`, so the embedding task is never dropped, no dummy model is built, and the output uses the canonical HF key names. This is **restricted to hybrid providers**, which are the only models that need it; non-hybrids keep the dummy-model path that CI has always exercised. Writing the source artifacts is now rank-0-only. Every rank used to write the source `config.json`, which races with the pruned `config.json` that `save_hf_pretrained` writes from rank 0 alone: a late write from another rank leaves a checkpoint whose config does not match its weights. This reproduced intermittently on both Qwen3 and NemotronH before the fix, and 3/3 clean runs after. `save_hf_pretrained` takes no `trust_remote_code` argument — it reads the flag **off the bridge** to fetch the source checkpoint's artifacts, and `from_hf_config` cannot infer it because `AutoConfig.from_pretrained` consumes the kwarg rather than storing it on the config. So the flag is set explicitly on the bridge instance; otherwise remote-code models would silently lose it. 3. **Config write-back correctness:** - `hybrid_override_pattern` is only written for older remote-code configs that lack `layer_types`; native configs carry the cadence in `layer_types` (read-only `hybrid_override_pattern` is skipped). - `n_shared_experts` is preserved (a fixed count) instead of being re-derived by `moe_shared_expert_intermediate_size // moe_ffn_hidden_size`, which is DeepSeek-style logic that would corrupt NemotronH's count. Non-hybrids, VLMs, and Megatron-Bridge builds without config-only export keep the dummy-model path, with a `warn_rank_0` when a hybrid has to fall back. The README's `transformers<5` workaround is **removed**: it existed because the dummy-model path broke on transformers 5, and the config-only path handles NemotronH on every supported container. #### 2. `candidate_filter` for `mcore_minitron` (`modelopt/torch/prune/plugins/mcore_minitron.py`) DeepSeek-style MoE configs have no explicit shared-expert-size field: they size the shared expert as `n_shared_experts * moe_intermediate_size`, where `moe_intermediate_size` is the (also prunable) **routed** expert size. So only candidates with `moe_shared_expert_intermediate_size % moe_ffn_hidden_size == 0` can be written back to HF at all. Candidates come from a Cartesian `product()` of independent per-hparam choice lists, so no per-hparam restriction can express a constraint *between* two hparams. New optional `candidate_filter` search-config key (default `None`, so existing behaviour is unchanged): a callable that rejects candidate configs before the metric computation, making the search cheaper rather than more expensive. It receives **every** supported hparam, with non-searched ones filled in from the model config, so a filter still works when one of its hparams was skipped or had a single choice. Rejected candidates are not cached, so — like `score_func`, whose cached scores are reused without re-validation — the filter is assumed unchanged when resuming from a `checkpoint`. `prune_minitron.py` wires this up for DeepSeek-style configs, so **both** `moe_ffn_hidden_size` and `moe_shared_expert_intermediate_size` stay prunable (the search then only picks shared sizes that are a multiple of the routed one). A `--prune_export_config` that violates the constraint never reaches the filter, so the export path now raises `ValueError` instead of writing a checkpoint whose config disagrees with its weights. #### 3. MLA Q-LoRA pruning (`modelopt/torch/prune/plugins/mcore_minitron.py`) Pruning any MLA model with `q_lora_rank` set died during calibration with `AttributeError: 'tuple' object has no attribute 'view'`. `hidden_size` importance estimation blanket-patches every `TELayerNormColumnParallelLinear` with `return_layernorm_output=True` to capture post-layernorm activations. When `q_lora_rank` is set, MCore builds `linear_q_up_proj` as a `TELayerNormColumnParallelLinear` — the Q-LoRA layernorm is fused into it, which is why `q_layernorm` is `IdentityOp` — so it was patched too, even though its layernorm is over the **latent rank**, not `hidden_size`. TE then returns `((out, ln_out), bias)` and MCore's `q, _ = self.linear_q_up_proj(...)` leaves `q` a tuple. Isolated by probing the module before and after dynamic conversion: | Setup | `linear_q_up_proj` returns | Forward | | --- | --- | --- | | Before conversion | `tuple(Tensor, NoneType)` | — | | After conversion, no hooks | `tuple(Tensor, NoneType)` | OK | | After conversion **+ importance hooks** | `tuple(tuple(Tensor, Tensor), NoneType)` | AttributeError | So conversion is innocent; registering the importance hooks is the trigger. Fix: exclude MLA's Q/KV up-projections from both the patch and unpatch loops. `test_mcore_mla_pruning` did not catch this because it builds MLA without `q_lora_rank`, where MCore uses a plain `linear_q_proj` and nothing is patched. #### 4. Fail fast instead of stalling on a rank-local error under PP (`modelopt/torch/utils/distributed.py`) A rank raising inside a distributed entrypoint left the whole job stalled until the process group timed out, with **no diagnostic output at all**: the failing rank blocked in `cleanup()`'s barrier while its peers blocked in `recv_from_prev_pipeline_rank_`, and Python only prints a traceback once the enclosing `finally` returns. A crash on one rank was indistinguishable from a slow job. - `dist.cleanup()` skips the barrier when unwinding from an exception. - New `dist.abort()` prints the traceback, flushes and exits immediately. Skipping the barrier alone is **not** enough — a stack dump showed the failing rank then blocking in `destroy_process_group` for the same reason — so the error path must not tear the process group down at all. `SystemExit` is re-raised rather than aborted, so an intentional exit (e.g. the `--score_lower_bound` gate) keeps its exit code and prints no traceback. Kept out of `cleanup()` so no library caller gets a surprise process exit. - Called from the entrypoints that wrap `main()` in `try/finally`: the five `examples/megatron_bridge` scripts. Measured on a 2-GPU PP run whose rank 0 raises during calibration: **10 min timeout kill with no visible error → 31s, exit 1, real traceback.** This is a latent, pre-existing issue (the `try/finally` predates this PR); it only surfaces on a failing PP run, which is why CI never hit it. ### Usage ```bash torchrun --nproc_per_node 4 examples/megatron_bridge/prune_minitron.py \ --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 \ --pp_size 4 \ --prune_target_active_params 3e9 \ --output_hf_path /path/to/Nemotron-3.5-Lightning-30B-A3B-Pruned-A3.0B ``` ### Testing - **End-to-end on nemo:26.08.rc6** (4× GB300, transformers 5.12.1, Megatron-Bridge with config-only export): pruning + export complete (`EXIT=0`, "Saved pruned model … Done!"). The exported checkpoint has canonical **plural** `backbone.embeddings.weight` keys, **0 MTP tensors**, and a config that reloads correctly (`num_hidden_layers=52` from `layers_block_type`, `n_shared_experts=1`, `num_nextn_predict_layers=0`, pruned `hidden_size`/`mamba_*`/MoE dims, reconstructed `hybrid_override_pattern`). <details> <summary>Pruning search log (<code>--prune_target_active_params 3e9</code>)</summary> ```text Top 10 Candidates with Scores ┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━┓ ┃ # ┃ export_config ┃ active_params ┃ params ┃ score ┃ ┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━┩ │ 1 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.49B │ 0.5406 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3584} │ │ │ │ │ 2 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 56, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2427 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 3 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 21.61B │ 0.2643 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 4 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 19.28B │ 0.4552 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3712} │ │ │ │ │ 5 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 104, │ 3.00B │ 22.28B │ 0.5860 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 6 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 21.99B │ 0.2294 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3328} │ │ │ │ │ 7 │ {'num_layers': 48, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.5231 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 8 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 96, │ 3.00B │ 21.81B │ 0.5042 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3584} │ │ │ │ │ 9 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2462 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 10 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 20.70B │ 0.5685 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ └────┴───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┴───────────────┴────────┴────────┘ ╭──────────────────────────────────────────────────────────────────────── Best Subnet ─────────────────────────────────────────────────────────────────────────╮ │ export_config {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 104, 'moe_ffn_hidden_size': 1856, │ │ 'moe_shared_expert_intermediate_size': 3072} │ │ active_params 3.00B │ │ params 22.28B │ │ score 0.5860 │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭────────────────────────────────────────────────────── Pruned Model Stats ───────────────────────────────────────────────────────╮ │ Total Parameters 22.28B │ │ Active Parameters 3.00B │ │ Memory (BF16, seq_length=8192, batch_size=8) weights: 42489.7 MB, kv_cache: 384.0 MB, mamba_state: 190.5 MB, Total: 43064.2 MB │ ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ``` </details> - **`tests/examples/megatron_bridge/test_prune_minitron.py`** — `nemotron_h` now exports to HF and reloads (previously it stopped at a Megatron checkpoint, since the dummy-model path needed `transformers<5`), plus an `n_shared_experts` config assertion; the dead `megatron_format` branch is gone. It runs on the CI container: **verified on nemo:26.06.01 (transformers 5.8.1) and nemo:26.08.rc6.** - **`tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py`** — the `nas_memory_mb` search test now passes a `candidate_filter` and asserts the exact number of rejected candidates (256 of the 512-combo grid) plus the surviving candidates' validity; its `expected_top_k` goldens are regenerated accordingly. Because `moe_shared_expert_intermediate_size` is in that test's skip list, this also covers the model-config fallback for hparams that are not in the search space. Verified on 2 GPUs, on both the CI container (nemo:26.06.01) and nemo:26.08.rc6: | Test | Result | | --- | --- | | `test_prune_minitron[qwen3]` | PASSED on 26.06.01 and 26.08.rc6 | | `test_prune_minitron[deepseek_v3]` | PASSED (52s) — MLA Q-LoRA + `candidate_filter` end-to-end | | `test_prune_minitron[nemotron_h]` | PASSED on 26.06.01 (58s) and 26.08.rc6 (61s) | | `test_mcore_mamba_hybrid_pruning_nas_memory_mb` | PASSED | | `test_mcore_mamba_hybrid_pruning_nas_params` | PASSED (unchanged sibling, run to check the regenerated goldens did not disturb it) | | 2-GPU PP run failing on rank 0 | fails in 31s with a real traceback (was a 10 min stall) | ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — `candidate_filter` defaults to `None` (existing searches unchanged), and the config-only export is limited to hybrid providers on nemo:26.08+, so dense / MoE / VLM exports keep the path they use today. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: ✅ ### Additional Information Enables the Prune + Distill workflow for `NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` (native, no-remote-code `NemotronHForCausalLM` with MTP heads). Pruning-time MTP support was scoped and intentionally deferred — MTP heads are dropped and can be re-derived via a short SFT with `mtp_num_layers=1` on the pruned+distilled model. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…n-Bridge (#2112) **Type of change:** Bug fix **Overview:** A quantized `output_layer` (lm_head) is **silently exported as BF16** when the model is built through Megatron-Bridge — no error, no warning, just an unquantized output layer in the exported checkpoint. This restores it. **1. The extra-state callbacks are never registered.** ```python if name.endswith("output_layer") and not getattr( getattr(module, "weight_quantizer", None), "is_enabled", False): continue ``` This hook also runs *before* `QuantModule` replacement (e.g. on restore), when `weight_quantizer` does not exist yet — so the check always skipped, and `output_layer` never received the ModelOpt extra-state callbacks. Its quantizer state (promotion to `StaticBlockScaleQuantizer`, `_amax`, `_global_amax`) was therefore never saved or restored. **2. Tiedness is read from a Megatron-LM global that Megatron-Bridge does not have.** ```python from megatron.training import get_args as _mlm_get_args _untied = bool(getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False)) except Exception as e: warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") ``` Megatron-Bridge has no `megatron.training` global args store, so the import raises, the `except` path treats the layer as tied, and quantization is dropped. The two combine into a silent failure: the save path writes nothing for the output layer, and the load path has nothing to restore. Resolve tiedness from the model instead of from a framework global: * `_resolve_output_layer_untied()` walks `named_modules()` for the framework-agnostic `share_embeddings_and_output_weights` flag and records it on the model config as `modelopt_output_layer_untied`. * `sharded_state_dict` prefers that flag and falls back to `get_args()` only when it is absent — **Megatron-LM behaviour is unchanged**. * Callback registration falls back to the tiedness flag when `weight_quantizer` does not exist yet, so an untied `output_layer` is registered. * When the load plan is built, the per-block NVFP4 scale buffers are materialized so the loader has somewhere to write. A Megatron dist-checkpoint load **silently skips any key the model does not advertise**, which is why the missing buffers produced no error. Divisibility is checked against `weight.shape[-1]` (the axis `_process_quantizer_amax` later views over), and a warning is emitted rather than leaving the buffers unallocated. `modelopt/torch/distill/plugins/megatron.py` gains a matching one-shot promotion on the first forward, for the single module the restore path does not promote. Measured on Nemotron-Nano-3: `already_sbsq=460 converted=1 disabled=0` — exactly `output_layer`. Verified end-to-end on Nemotron-Nano-3, W4A16 NVFP4 `four_over_six`, 4x GB200: * exported checkpoint carries `lm_head.weight` (U8 `[131072, 1344]`) plus `weight_scale` and `weight_scale_2` * matches a known-good reference export exactly: 52 shards, 18487 keys, 72 exclusions * converts to compressed-tensors and serves correctly on stock vLLM 0.26.0 at TP=2 * during QAD, `disabled=0` (was `disabled=1` before the fix) independently confirms the output layer stays quantized through distillation Without the fix the same pipeline produces a BF16 `lm_head` with no diagnostic of any kind. - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes. Megatron-LM keeps its existing `get_args()` path; the new model-derived flag is only consulted first, and only affects cases that are currently broken. - **Did you write any new necessary tests?**: Yes — unit coverage for both tiedness resolvers (`_resolve_output_layer_untied`, `_output_layer_untied`). Also verified end-to-end on Nemotron-Nano-3.5: byte-identical export, and QAD matching the reference iteration-10 KD loss at identical LR. - **Did you add or update any necessary documentation?**: No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No `quant_module_get_extra_state` now returns `{}` for a module with nothing quantized, so registering the extra-state callbacks unconditionally does not give unquantized layers non-empty extra state. An earlier revision of this PR gated registration on output-layer tiedness and did change the state-dict schema for untied-but-disabled `output_layer`; that revision has been replaced. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> * **Bug Fixes** * Improved quantized Megatron model state handling when quantization is disabled. * Corrected output-layer weight-tying detection across model configurations, including vision and teacher components. * Improved compatibility with Megatron-LM settings when determining output-layer behavior. * Ensured sharded model state is generated consistently using the resolved weight-tying configuration. * Added reliable handling and caching of weight-tying information for consistent model state output. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: James Shen <yueshen@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
requested review from
AAnoosheh,
h-guo18 and
sugunav14
and removed request for
a team
August 12, 2026 17:56
kevalmorabia97
requested review from
kaix-nv,
kinjalpatel27 and
realAsma
and removed request for
a team
August 12, 2026 17:56
Contributor
|
Note Currently processing new changes in this PR. This may take a few minutes, please wait... ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (66)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
|
Contributor
|
Caution CodeRabbit couldn't update its existing comment. The review summary may be out of date. Error details |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/utils/plugins/transformers_dataset.py (1)
433-448: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPrefer the template's
{% generation %}mask over the marker heuristic.
derive_masks_from_markersis true wheneveranswer_only_lossis set and the template contains<|im_start|>and<|im_end|>. A ChatML template can contain both those markers and{% generation %}tags. For such a template this code stops requestingreturn_assistant_tokens_maskfrom the processor and replaces the authoritative tag-based mask with the token-subsequence heuristic. That changes the loss mask for templates that already worked.Derive masks from markers only when the template has no
{% generation %}tags.🐛 Proposed fix to gate the heuristic on the absence of generation tags
def _process_multimodal_sample(self, examples): - derive_masks_from_markers = self.answer_only_loss and bool(self._assistant_marker_specs()) + template = self.tokenizer.chat_template or "" + # The template's own {% generation %} tags are authoritative; the ChatML + # marker heuristic is only a fallback for templates that lack them. + derive_masks_from_markers = ( + self.answer_only_loss + and "{% generation %}" not in template + and "{%- generation %}" not in template + and bool(self._assistant_marker_specs()) + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 433 - 448, Update _process_multimodal_sample so derive_masks_from_markers is enabled only when answer_only_loss is true, assistant markers are present, and the chat template contains no {% generation %} tags. Preserve the processor-provided return_assistant_tokens_mask path for templates that define generation tags.
🧹 Nitpick comments (4)
tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py (1)
95-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the produced labels, not just the stored flag.
The docstring states that the real VLM collator must support DFlash's unshifted labels. The test only asserts
collator.shift_labels is False, which checks constructor plumbing. It does not verify that_process_multimodal_samplewriteslabels[i] == input_ids[i]instead of the shifted variant. A regression in the label-construction branch atmodelopt/torch/utils/plugins/transformers_dataset.pylines 453-458 would not fail this test.Add a fake
processor.apply_chat_templatethat returnsinput_idsandattention_masktensors, then assert the returned labels equalinput_ids.As per path instructions, "Tests must exercise the behavior they claim to validate."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py` around lines 95 - 110, Update test_vlm_data_collator_accepts_unshifted_labels to mock processor.apply_chat_template with input_ids and attention_mask tensors, invoke the collator’s multimodal sample-processing path, and assert the produced labels exactly equal input_ids. Retain the shift_labels assertion only as supplemental constructor coverage.Source: Path instructions
modelopt/torch/utils/plugins/transformers_dataset.py (1)
460-471: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLog when answer-only masking discards every label.
When
assistant_maskcontains no non-zero entry, the code sets the whole label tensor toIGNORE_TOKEN_ID. For the marker-derived path an all-zero mask usually means the marker heuristic did not match the rendered template, not that the assistant content was truncated. Training then proceeds with zero supervised tokens and no diagnostic output.Emit a rank-aware warning in that branch so the misconfiguration is visible.
♻️ Proposed warning for the fully masked batch
assistant_mask = tokenized_messages["assistant_masks"] if not isinstance(assistant_mask, torch.Tensor) or not assistant_mask.any(): + print_rank_0( + "answer_only_loss produced an empty assistant mask for this batch; " + "all labels are ignored. Check the chat template's assistant markers." + ) labels[:] = IGNORE_TOKEN_ID🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 460 - 471, Add a rank-aware warning in the answer_only_loss branch of the tokenized_messages masking logic when assistant_mask is non-tensor or contains no non-zero entries, immediately before setting labels[:] to IGNORE_TOKEN_ID. Keep the existing masking behavior unchanged and include enough context to identify that the entire batch has been masked.tests/unit/torch/speculative/plugins/test_hf_dflash.py (1)
239-259: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the remaining validation branches.
The new tests cover the video-group mismatch, the missing
mm_token_type_ids, and the shape mismatch. Three error paths inmodelopt/torch/speculative/plugins/hf_dflash.pystay uncovered:
_expand_qwen3_video_grid_thwrejecting a grid whose shape is not[num_videos, 3]or whose temporal length is not positive (lines 140-146).- The
compute_3d_position_idsfallback whenget_rope_indexis absent (lines 337-342).- The final
RuntimeErrorfor an unexpected mRoPE output shape (lines 344-352).These branches guard against Transformers API drift, which is the exact failure mode this workaround exists for.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/speculative/plugins/test_hf_dflash.py` around lines 239 - 259, Add tests covering the remaining validation paths in HFDFlashModel: verify _expand_qwen3_video_grid_thw rejects grids with invalid rank/shape and non-positive temporal lengths, verify _qwen3_vl_position_ids uses compute_3d_position_ids when get_rope_index is absent, and verify it raises RuntimeError for an unexpected mRoPE output shape. Keep the existing Transformers-version and Qwen3-VL test setup, asserting the relevant exception or fallback behavior for each branch.modelopt/torch/speculative/utils.py (1)
613-619: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winThe VLM auto-class preference order and its test share one root cause. The loader prefers the pre-5
AutoModelForVision2Seqand treatsAutoModelForImageTextToTextas the fallback, while the test acknowledges thattransformersre-creates the legacy name lazily. If Transformers 5 still resolves the legacy name, the fallback never runs outside the test, and the stated Transformers 5 support is unverified.
modelopt/torch/speculative/utils.py#L613-L619: preferAutoModelForImageTextToTextand fall back toAutoModelForVision2Seqso Transformers 5 selects the current class and Transformers 4 still works.tests/unit/torch/speculative/plugins/test_fakebase.py#L155-L164: after the preference order is inverted, drop the forcedNonefor the legacy attribute and assert the selected class directly, so the test reflects a real Transformers 5 environment.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/speculative/utils.py` around lines 613 - 619, Invert the VLM auto-class preference in modelopt/torch/speculative/utils.py lines 613-619: select transformers.AutoModelForImageTextToText first and fall back to AutoModelForVision2Seq for Transformers 4 compatibility. Update tests/unit/torch/speculative/plugins/test_fakebase.py lines 155-164 by removing the forced legacy-attribute None setup and asserting the selected class directly, so the test exercises the real Transformers 5 resolution behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/vllm_serve/Dockerfile`:
- Around line 33-34: Replace the recursive world-writable permission change in
the Dockerfile with ownership of /workspace by vllm:vllm, and grant write access
only to /workspace/torch_extensions if needed. Keep the rest of the workspace
non-world-writable so editable source and cached extensions cannot be modified
by other container users.
In `@examples/vllm_serve/vllm_ptq_utils.py`:
- Around line 105-132: The cleanup block after execute_model currently
suppresses finish_requests failures; capture the cleanup exception and propagate
it when execution completed successfully, while preserving any original
execute_model exception. Update the surrounding execute_model/finally control
flow and finish_requests handling so warnings do not replace a cleanup failure
after a successful batch.
In
`@modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml`:
- Around line 4-7: Update the quantization configuration comments to document
the FP8 activation exception for the Mamba projections governed by the input
quantizers around the Mamba projection settings, while retaining BF16 as the
general activation precision. Do not remove those input quantizers unless the
configuration is changed to enforce BF16 activations for those projections.
In `@modelopt/torch/utils/plugins/transformers_dataset.py`:
- Around line 416-429: Update the assistant-mask range in the turn-processing
logic around _find_subsequence so a matched end marker, such as Qwen’s
<|im_end|>, is included through its full token span in assistant_masks. Adjust
content_end and subsequent search_from handling consistently, while preserving
sequence_end behavior when no terminator is found.
---
Outside diff comments:
In `@modelopt/torch/utils/plugins/transformers_dataset.py`:
- Around line 433-448: Update _process_multimodal_sample so
derive_masks_from_markers is enabled only when answer_only_loss is true,
assistant markers are present, and the chat template contains no {% generation
%} tags. Preserve the processor-provided return_assistant_tokens_mask path for
templates that define generation tags.
---
Nitpick comments:
In `@modelopt/torch/speculative/utils.py`:
- Around line 613-619: Invert the VLM auto-class preference in
modelopt/torch/speculative/utils.py lines 613-619: select
transformers.AutoModelForImageTextToText first and fall back to
AutoModelForVision2Seq for Transformers 4 compatibility. Update
tests/unit/torch/speculative/plugins/test_fakebase.py lines 155-164 by removing
the forced legacy-attribute None setup and asserting the selected class
directly, so the test exercises the real Transformers 5 resolution behavior.
In `@modelopt/torch/utils/plugins/transformers_dataset.py`:
- Around line 460-471: Add a rank-aware warning in the answer_only_loss branch
of the tokenized_messages masking logic when assistant_mask is non-tensor or
contains no non-zero entries, immediately before setting labels[:] to
IGNORE_TOKEN_ID. Keep the existing masking behavior unchanged and include enough
context to identify that the entire batch has been masked.
In `@tests/unit/torch/speculative/plugins/test_hf_dflash.py`:
- Around line 239-259: Add tests covering the remaining validation paths in
HFDFlashModel: verify _expand_qwen3_video_grid_thw rejects grids with invalid
rank/shape and non-positive temporal lengths, verify _qwen3_vl_position_ids uses
compute_3d_position_ids when get_rope_index is absent, and verify it raises
RuntimeError for an unexpected mRoPE output shape. Keep the existing
Transformers-version and Qwen3-VL test setup, asserting the relevant exception
or fallback behavior for each branch.
In `@tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py`:
- Around line 95-110: Update test_vlm_data_collator_accepts_unshifted_labels to
mock processor.apply_chat_template with input_ids and attention_mask tensors,
invoke the collator’s multimodal sample-processing path, and assert the produced
labels exactly equal input_ids. Retain the shift_labels assertion only as
supplemental constructor coverage.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: aafd9785-1d46-4fb2-ae79-780ebbce83cd
📒 Files selected for processing (66)
CHANGELOG.rstexamples/hf_ptq/README.mdexamples/hf_ptq/example_utils.pyexamples/hf_ptq/hf_ptq.pyexamples/llm_qat/export.pyexamples/megatron_bridge/README.mdexamples/megatron_bridge/distill.pyexamples/megatron_bridge/export_distilled_megatron_to_hf.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/prune_minitron.pyexamples/megatron_bridge/quantize.pyexamples/speculative_decoding/eagle_utils.pyexamples/vllm_serve/Dockerfileexamples/vllm_serve/README.mdexamples/vllm_serve/vllm_ptq_utils.pymodelopt/torch/export/layer_utils.pymodelopt/torch/export/model_utils.pymodelopt/torch/export/plugins/hf_spec_export.pymodelopt/torch/export/quant_aware_conversion.pymodelopt/torch/export/quant_utils.pymodelopt/torch/opt/plugins/transformers.pymodelopt/torch/prune/plugins/mcore_minitron.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/sparsity/attention_sparsity/plugins/vllm.pymodelopt/torch/speculative/plugins/hf_dflash.pymodelopt/torch/speculative/utils.pymodelopt/torch/utils/distributed.pymodelopt/torch/utils/plugins/transformers_dataset.pymodelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yamlmodelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yamlmodelopt_recipes/configs/ptq/units/README.mdmodelopt_recipes/configs/ptq/units/default_disabled_quantizers.yamlmodelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yamlmodelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yamlmodelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yamlmodelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yamlmodelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yamlmodelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yamlmodelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yamlmodelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16.yamlmodelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-max-calib.yamlmodelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yamlmodelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6.yamlmodelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yamlmodelopt_recipes/huggingface/phi4mm/ptq/README.mdmodelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yamlmodelopt_recipes/ptq.mdtests/examples/llm_qat/test_llm_qat.pytests/examples/megatron_bridge/test_prune_minitron.pytests/gpu/torch/export/test_export.pytests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.pytests/unit/recipe/test_loader.pytests/unit/torch/export/test_hf_spec_rope_export.pytests/unit/torch/export/test_quant_aware_conversion.pytests/unit/torch/opt/plugins/test_hf_patching.pytests/unit/torch/opt/plugins/test_transformers_save_load.pytests/unit/torch/speculative/plugins/test_fakebase.pytests/unit/torch/speculative/plugins/test_hf_dflash.pytests/unit/torch/speculative/plugins/test_hf_speculative_offline.pytools/launcher/common/megatron_lm/quantize/quantize.shtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml
💤 Files with no reviewable changes (7)
- modelopt_recipes/huggingface/phi4mm/ptq/README.md
- examples/megatron_bridge/README.md
- examples/hf_ptq/hf_ptq.py
- examples/hf_ptq/example_utils.py
- modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml
- examples/hf_ptq/README.md
- modelopt/torch/export/model_utils.py
Comment on lines
+33
to
34
| # Allow the non-root vllm user to access the workspace | ||
| RUN chmod -R 777 /workspace |
Contributor
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Do not make the complete workspace world-writable.
Line 34 makes Model-Optimizer writable by every container user. A second process can modify editable Python source or cached extensions before the vllm process imports them.
Use chown -R vllm:vllm /workspace and grant write access only to /workspace/torch_extensions if required.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vllm_serve/Dockerfile` around lines 33 - 34, Replace the recursive
world-writable permission change in the Dockerfile with ownership of /workspace
by vllm:vllm, and grant write access only to /workspace/torch_extensions if
needed. Keep the rest of the workspace non-world-writable so editable source and
cached extensions cannot be modified by other container users.
Comment on lines
+105
to
+132
| finally: | ||
| # finish_requests runs before add_requests inside execute_model, so | ||
| # req IDs aren't registered yet at that point — call it directly after. | ||
| # Wrap in try/except so a cleanup error never masks the original exception. | ||
| try: | ||
| if hasattr(self.model_runner, "finish_requests"): | ||
| cleanup_output = _create_new_data_cls( | ||
| type(scheduler_output), | ||
| scheduled_new_reqs=[], | ||
| scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs, | ||
| num_scheduled_tokens={}, | ||
| total_num_scheduled_tokens=0, | ||
| scheduled_spec_decode_tokens={}, | ||
| scheduled_encoder_inputs={}, | ||
| num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks, | ||
| finished_req_ids=set(num_scheduled_tokens.keys()), | ||
| free_encoder_mm_hashes=[], | ||
| kv_connector_metadata=None, | ||
| structured_output_request_ids={}, | ||
| grammar_bitmask=None, | ||
| ) | ||
| self.model_runner.finish_requests(cleanup_output) | ||
| else: | ||
| warnings.warn( | ||
| "model_runner.finish_requests not found; request state may leak during calibration." | ||
| ) | ||
| except Exception: | ||
| warnings.warn("Failed to clean up request state after calibration batch.") |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Fail calibration when request cleanup fails after a successful batch.
If finish_requests fails, the model runner can retain request state and KV-cache allocations. The next batches use new request IDs, so the leak can accumulate until calibration produces incorrect state or runs out of memory.
Preserve an exception from execute_model. If execution succeeded, propagate the cleanup exception instead of only warning.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vllm_serve/vllm_ptq_utils.py` around lines 105 - 132, The cleanup
block after execute_model currently suppresses finish_requests failures; capture
the cleanup exception and propagate it when execution completed successfully,
while preserving any original execute_model exception. Update the surrounding
execute_model/finally control flow and finish_requests handling so warnings do
not replace a cleanup failure after a successful batch.
Comment on lines
+4
to
+7
| # Quantization config for nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4: | ||
| # NVFP4 weights use the four_over_six (4/6) per-block scale selection. | ||
| # Weight-only; activations BF16. MSE selects per-block M=6 (mult 1.0) vs | ||
| # M=4 (mult 1.5); the choice is baked into the amax by MSE. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the activation-precision description.
Lines 6-7 state that activations are BF16. Lines 52-60 enable FP8 input quantizers for both Mamba projections. Document the FP8 Mamba activation exception, or remove those input quantizers if BF16 activations are intended.
Also applies to: 19-23
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml`
around lines 4 - 7, Update the quantization configuration comments to document
the FP8 activation exception for the Mamba projections governed by the input
quantizers around the Mamba projection settings, while retaining BF16 as the
general activation precision. Do not remove those input quantizers unless the
configuration is changed to enforce BF16 activations for those projections.
Comment on lines
+416
to
+429
| end_positions = [ | ||
| position | ||
| for marker in end_markers | ||
| if ( | ||
| position := self._find_subsequence( | ||
| tokens, marker, content_start, sequence_end | ||
| ) | ||
| ) | ||
| != -1 | ||
| ] | ||
| content_end = min(end_positions) if end_positions else sequence_end | ||
| if content_start < content_end: | ||
| assistant_masks[row_index, content_start:content_end] = 1 | ||
| search_from = max(content_start + 1, content_end + 1) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Hugging Face chat template generation tag include im_end token assistant tokens mask
💡 Result:
In Hugging Face chat templates, the {% generation %} and {% endgeneration %} tags are used to explicitly mark the boundaries of assistant-generated content [1][2][3]. When return_assistant_tokens_mask=True is set in the apply_chat_template method, these tags allow the tokenizer to generate a mask that identifies which tokens belong to the assistant, facilitating assistant-only loss masking during Supervised Fine-Tuning (SFT) [1][2]. Tokens such as <|im_start|> and <|im_end|> (or similar variants like <|eot_id|>) act as control tokens that define the structure of the chat, indicating the start and end of messages [4][5]. It is critical that these markers remain outside the {% generation %} blocks to ensure correct masking: the assistant's content and the trailing stop token (like <|im_end|>) should be wrapped within the {% generation %}/{% endgeneration %} tags, while the assistant header (e.g., <|im_start|>assistant) should remain outside to correctly represent the prompt cue that the model does not generate [1][2]. For example, a template might be structured as follows: {% if message['role'] == 'assistant' %} {% generation %} {{ message['content'] + '<|im_end|>\n' }} {% endgeneration %} {% endif %} Using these markers ensures that the loss function can effectively ignore non-assistant tokens (such as system or user prompts) by masking them, typically with a value like -100 [1][6]. Libraries such as TRL provide utility functions like get_training_chat_template to automatically inject these markers and ensure templates are prefix-preserving for complex tasks like tool calling [7].
Citations:
- 1: https://huggingface.co/docs/trl/main/chat_templates
- 2: https://huggingface.co/docs/trl/chat_templates
- 3: https://huggingface.co/docs/trl/en/chat_templates
- 4: https://huggingface.co/docs/transformers/chat_templating
- 5: https://huggingface.co/docs/transformers/main/en/chat_template_advanced
- 6: return mask of user messages when calling
tokenizer.apply_chat_template(c,tokenize=True)huggingface/transformers#28950 - 7: https://huggingface.co/docs/trl/main/en/chat_template_utils
🏁 Script executed:
#!/bin/bash
set -e
file="modelopt/torch/utils/plugins/transformers_dataset.py"
printf '%s\n' '--- target implementation ---'
sed -n '330,455p' "$file"
printf '%s\n' '--- generation-tag and marker references ---'
rg -n -C 3 "_verify_generation_tags|end_markers|im_end|generation|assistant_masks|content_end" modelopt tests examples 2>/dev/null | head -n 300Repository: NVIDIA/Model-Optimizer
Length of output: 33274
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from pathlib import Path
p = Path("modelopt/torch/utils/plugins/transformers_dataset.py")
text = p.read_text()
for needle in ("def _verify_generation_tags", "def _find_subsequence", "end_markers", "assistant_masks"):
print(f"\n--- {needle} ---")
starts = [i for i, line in enumerate(text.splitlines(), 1) if needle in line]
for start in starts:
lines = text.splitlines()
lo = max(1, start - 15)
hi = min(len(lines), start + 90)
for n in range(lo, hi + 1):
print(f"{n:4}: {lines[n-1]}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 50378
🌐 Web query:
site:github.com/huggingface/transformers "{% generation %}" "im_end"
💡 Result:
As of August 12, 2026, the term "quantum supremacy" is widely understood in the scientific and technical community as a milestone that was achieved years ago, beginning with Google's Sycamore processor in 2019. Since that time, the field has transitioned its focus from demonstrating quantum supremacy—the capability of a quantum computer to solve a task that is practically impossible for a classical supercomputer—to the more practical objectives of achieving "quantum advantage" and "fault-tolerant quantum computing." Current research and development in 2026 are focused on scaling quantum systems, improving qubit coherence times, and developing error-correction techniques to enable reliable computations. While quantum supremacy demonstrated a proof-of-concept, the industry's current efforts are directed toward delivering useful, noise-resilient quantum applications in cryptography, materials science, and complex optimization. Technical developments in the broader computing ecosystem, including advancements in AI architectures (such as multimodal models and specialized quantization techniques), are occurring in parallel, but quantum computing remains a distinct discipline progressing toward scalable, fault-tolerant hardware.
Citations:
- 1: [Research] Fundamental Equation of Consciousness: Ψ = argmax H(p) s.t. Φ > Φ_min huggingface/transformers#45102
- 2: huggingface/transformers@b6d5084
- 3: [Energy] N6 Arithmetic: 50-70% AI Training/Inference Energy Reduction — 17 Techniques with Code huggingface/transformers#45145
- 4: https://github.com/huggingface/transformers/blob/1f8daee0/docs/source/en/model_doc/qwen2_5_vl.md
- 5: https://github.com/huggingface/transformers/blob/main/docs/source/en/model_doc/vibevoice_asr.md
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from pathlib import Path
p = Path("modelopt/torch/utils/plugins/transformers_dataset.py")
lines = p.read_text().splitlines()
for lo, hi in ((1, 180), (433, 475)):
print(f"--- lines {lo}-{hi} ---")
for n in range(lo, min(hi, len(lines)) + 1):
print(f"{n:4}: {lines[n-1]}")
PY
printf '%s\n' '--- repository templates and focused tests ---'
rg -n -l --glob '*.py' --glob '*.json' --glob '*.jinja' --glob '*.yaml' --glob '*.yml' \
'im_start|im_end|endgeneration|return_assistant_tokens_mask|_build_assistant_masks' . \
| head -n 120Repository: NVIDIA/Model-Optimizer
Length of output: 11516
🏁 Script executed:
#!/bin/bash
set -e
tmp=$(mktemp)
curl -L --fail --silent --show-error \
https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/processing_utils.py \
-o "$tmp"
printf '%s\n' '--- generation-mask implementation references ---'
rg -n -C 8 'assistant_masks|generation_indices|return_assistant_tokens_mask|generation' "$tmp" | head -n 240
rm -f "$tmp"Repository: NVIDIA/Model-Optimizer
Length of output: 12795
🏁 Script executed:
#!/bin/bash
set -e
for url in \
https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/utils/chat_template_utils.py \
https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/tokenization_utils_base.py
do
tmp=$(mktemp)
curl -L --fail --silent --show-error "$url" -o "$tmp"
printf '\n--- %s ---\n' "$url"
rg -n -C 12 'generation_indices|GenerationMarker|generation' "$tmp" | head -n 320
rm -f "$tmp"
done
printf '\n--- focused tests ---\n'
sed -n '1,280p' tests/examples/llm_qat/test_assistant_mask.py
printf '\n--- ChatML-like templates ---\n'
for f in tools/launcher/examples/Qwen/Qwen3-0.6B/chat_template_train.jinja \
tools/launcher/examples/Qwen/Qwen3-8B/chat_template_train.jinja \
tools/launcher/examples/MiniMax/MiniMax-M2.7-DFlash/chat_template_train.jinja; do
echo "--- $f ---"
rg -n -C 8 'im_start|im_end|generation|assistant' "$f"
doneRepository: NVIDIA/Model-Optimizer
Length of output: 38518
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- VisionLanguageDataCollator call sites and VLM configuration ---'
rg -n -C 5 'VisionLanguageDataCollator|answer_only_loss|chat_template_train|Cosmos|Qwen.*VL|Qwen2.*VL' \
modelopt tests tools examples --glob '*.py' --glob '*.yaml' --glob '*.yml' --glob '*.json' --glob '*.jinja' \
| head -n 360
printf '%s\n' '--- all tracked generation-tag templates ---'
rg -n -C 5 '<\|im_end\|>|endgeneration|generation' --glob '*.jinja' --glob '*.json' --glob '*.py' . \
| head -n 360Repository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
# Read-only verifier for the boundary claimed in the review:
# a ChatML sequence with one assistant content token followed by <|im_end|>.
tokens = ["<|im_start|>", "assistant", "\\n", "answer", "<|im_end|>", "\\n"]
start_marker = ["<|im_start|>", "assistant", "\\n"]
end_marker = ["<|im_end|>", "\\n"]
start = 0
content_start = start + len(start_marker)
end = next(i for i in range(content_start, len(tokens)) if tokens[i:i+len(end_marker)] == end_marker)
exclusive_mask = list(range(content_start, end))
inclusive_mask = list(range(content_start, end + len(end_marker)))
print({"content_start": content_start, "end": end,
"exclusive_mask": [tokens[i] for i in exclusive_mask],
"inclusive_mask": [tokens[i] for i in inclusive_mask]})
PYRepository: NVIDIA/Model-Optimizer
Length of output: 272
🏁 Script executed:
#!/bin/bash
set -e
for ref in \
Qwen/Qwen2.5-VL-7B-Instruct \
Qwen/Qwen3-VL-8B-Instruct \
Qwen/Qwen2-VL-7B-Instruct
do
echo "--- $ref tokenizer_config.json ---"
encoded=${ref//\//%2F}
url="https://huggingface.co/${ref}/raw/main/tokenizer_config.json"
if curl -L --fail --silent --show-error "$url" | \
python3 -c 'import json,sys; d=json.load(sys.stdin); print(d.get("chat_template",""))' | \
rg -n -C 5 'im_start|im_end|generation|assistant'; then
:
else
echo "template unavailable"
fi
doneRepository: NVIDIA/Model-Optimizer
Length of output: 5709
🏁 Script executed:
#!/bin/bash
set -e
for ref in \
Qwen/Qwen2.5-VL-7B-Instruct \
Qwen/Qwen3-VL-8B-Instruct \
Qwen/Qwen2-VL-7B-Instruct
do
echo "--- $ref tokenizer_config.json ---"
url="https://huggingface.co/${ref}/raw/main/tokenizer_config.json"
if curl -L --fail --silent --show-error "$url" | \
python3 -c 'import json,sys; d=json.load(sys.stdin); print(d.get("chat_template",""))' | \
rg -n -C 5 'im_start|im_end|generation|assistant'; then
:
else
echo "template unavailable"
fi
doneRepository: NVIDIA/Model-Optimizer
Length of output: 5709
Include the turn terminator in the assistant mask.
Qwen VLM ChatML templates emit <|im_end|> after assistant content. The current slice excludes it, so shifted labels ignore the turn-terminator target. Extend content_end past the matched end marker, or document and test the intentional exclusion.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/utils/plugins/transformers_dataset.py` around lines 416 - 429,
Update the assistant-mask range in the turn-processing logic around
_find_subsequence so a matched end marker, such as Qwen’s <|im_end|>, is
included through its full token span in assistant_masks. Adjust content_end and
subsequent search_from handling consistently, while preserving sequence_end
behavior when no terminator is found.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## release/0.46.0 #2179 +/- ##
===================================================
+ Coverage 66.86% 77.98% +11.11%
===================================================
Files 519 519
Lines 59323 59576 +253
===================================================
+ Hits 39665 46458 +6793
+ Misses 19658 13118 -6540
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This was referenced Aug 12, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
10 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-picked PRs
Summary by CodeRabbit
New Features
Bug Fixes
Documentation