[Cherry-pick] PRs #1352 #1351 #1330 #1354 #1355 #1360 #1342 #1324 #1340 #1368 #1373 #1359 #1361 #1325 #1369 #1370 #1371

[kevalmorabia97]

Cherry-picked PRs

• Add release-cherry-pick Claude Code skill #1352
• chore: Move FP8 MHA quantization entry from 0.45 to 0.44 in CHANGELOG #1351
• fix incomplete mapping of safetensors in generated puzzletron checkpoint #1330
• [NVBUG: 6103846] Fix nvfp4_awq export for uncalibrated MoE experts #1354
• Fix regex capture for Megatron KD PP layer renaming #1355
• Add required keys to attention pruning config #1360
• Support EP mcore import for TE Spec and Fix mamba moe config #1342
• [NVBug 6102977] Add _disable_use_cache context manager to fix PTQ AttributeError on custom configs #1324
• [NVBug 6108145] Fix PTQ calibration and export for fused-experts MoE (Qwen3.5-MoE VLM) #1340
• [BUG6108338] Update windows documentation for onnxruntime quantization with Cuda13.x #1368
• [Fix]: Relax Dflash Rregression Test Threshold fo 2GPUs #1373
• Ensure removal of temp files on error in ONNX INT4 quantization #1359
• [6034518] Remove return statement preventing remote auto tuning #1361
• Add Nemotron-Nano-9B-v2 → Pruned 7B e2e tutorial: Prune + Distill + Eval + Quantize + vLLM deployment #1325
• Added fallback to load extra cudnn dlls in the site packages  #1369
• fix: include medusa in data_module assignment in main.py #1370
• fix: guard against None chat_template in _post_process_chat_template #1371

Summary by CodeRabbit

• New Features

  • New end-to-end Minitron pruning→distillation→FP8 quantization→vLLM tutorial and Nemotron-Nano-9B-v2 example
  • FP8 MHA quantization support for vision transformers
  • Windows CUDA‑13.x installation instructions and Megatron dataset pre-tokenization guide
  • Improved grouped MoE export and fused-experts handling; medusa mode enabled in speculative decoding

• Bug Fixes

  • nvfp4_awq export fallback for MoE models
  • Enhanced Windows cuDNN discovery and expert calibration robustness

• Documentation

  • Expanded pruning/distillation/quantization tutorials, ablations, and examples

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds new documentation and examples for pruning/distillation/quantization/deployment, multiple fused-experts quantization/export changes, distributed checkpoint saving for Puzzletron with tests, ONNX quantization cleanup and Windows cuDNN loading, dataset/tokenization and cache-control updates, Megatron exporter/importer fixes, new regression/unit tests, a CI timeout tweak, and a release-cherry-pick workflow doc.

Changes

Documentation & Examples (Nemotron / Pruning / Puzzletron)

Layer / File(s)	Summary
High-level docs CHANGELOG.rst, examples/*/README.md	Added changelog entries and multiple README edits linking to new documents and updating references.
Megatron data prep guide examples/dataset/MEGATRON_DATA_PREP.md, examples/dataset/README.md	New standalone Megatron pre-tokenization guide; README updated to link to it.
Minitron pruning tutorial & artefacts examples/pruning/*, examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/*	Added end-to-end Minitron pruning+distillation+quantization+vLLM deployment README, ablations, evaluator YAML, and supporting README/link updates.
Puzzletron docs examples/pruning/puzzletron/README.md, examples/puzzletron/README.md	Added Puzzletron pruning/distillation results README and cross-reference updates.

Fused-Experts Quantization & Export

Layer / File(s)	Summary
Core utility addition modelopt/torch/quantization/utils/core_utils.py, modelopt/torch/quantization/utils/__init__.py	Added representative_weight_quantizer() and exported it; updated weight_attr_names() to handle plural/fused-experts quantizer layouts.
Name normalization modelopt/torch/quantization/conversion.py	Added regex normalization to collapse plural per-expert quantizer naming for wildcard matching.
Calibration robustness modelopt/torch/quantization/model_calib.py, modelopt/torch/quantization/config.py	Defensive handling for AWQ-lite experts (disable/mask unexercised experts) and extended disabled-quantizer patterns for Mamba-MoE/self-attention names.
Plugin changes modelopt/torch/quantization/plugins/huggingface.py	Added _QuantFusedExperts.fold_weight() to fold per-expert fake-quant and force_eager_experts_impl_on_the_fly() registered in plugins to force eager expert impl when needed.
Export & fake-quant modelopt/torch/export/plugins/vllm_fakequant_hf.py, modelopt/torch/export/quant_utils.py, modelopt/torch/export/unified_export_hf.py	Broadened weight-quantizer state-key regex; added _fakequant_fused_experts_weights() and ensured fused-experts are exported/handled before generic quantized-module logic; export detection now uses representative_weight_quantizer().
Tests tests/unit/torch/quantization/plugins/test_fused_experts.py	New unit tests covering eager-impl forcing, per-expert calibration amax checks, mixed-precision export enumeration, and quantizer-name normalization.

Puzzletron Distributed Checkpoint Saving & Tests

Layer / File(s)	Summary
Implementation modelopt/torch/puzzletron/tools/checkpoint_utils_hf.py	Added save_checkpoint_from_shards() to gather per-rank state dicts to rank 0, merge, save once, broadcast errors, and synchronize. Exported it via __all__.
Integration modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py	Switched checkpoint save calls to save_checkpoint_from_shards.
Tests tests/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py	New tests validating single-process and distributed sharded checkpoint saving, index/shard/config outputs, tied-embedding exclusion, and numeric tensor equality after reload.
Test import behavior tests/gpu/torch/puzzletron/test_puzzletron.py	Removed module-level pytest import-or-skip guard for mip so module can load when optional dependency missing.

ONNX Quantization Cleanup & cuDNN DLL Loading

Layer / File(s)	Summary
Cleanup helper modelopt/onnx/quantization/int4.py	Added _remove_augmented_onnx() and refactored _quantize_awq_clip / _quantize_awq_lite to use try/finally with session cleanup, gc.collect(), and unified ONNX removal.
Windows cuDNN loading modelopt/onnx/quantization/ort_utils.py	Added helpers to find nvidia-cudnn-cu12/cu13 site-packages DLL dir and explicitly load missing cudnn*.dll on Windows; adjusted preload output handling and core-DLL checks.

Dataset Utilities, Tokenization, and Speculative Decoding

Layer / File(s)	Summary
Cache-control context modelopt/torch/utils/dataset_utils.py, tests/unit/torch/utils/test_dataset_utils.py	Added _disable_use_cache() context manager; updated get_max_batch_size and _forward_loop to use it; added unit tests for behavior and restoration.
Tokenizer & streaming changes modelopt/torch/utils/plugins/megatron_preprocess_data.py, examples/dataset/MEGATRON_DATA_PREP.md	Tokenizer apply_chat_template now wrapped with try/except to log input on failure; removed fallback to non-streaming when --hf_streaming used; updated docs and examples.
Chat template guard modelopt/torch/utils/plugins/transformers_dataset.py	Added early-return when tokenizer chat_template is None to avoid .replace on None.
Speculative decoding examples/speculative_decoding/main.py	Construct data module for medusa mode by adding "medusa" to the mode condition.

Megatron Importer & Exporter / Model Support

Layer / File(s)	Summary
Exporter validation modelopt/torch/export/unified_export_megatron.py	Added optional guarded import of HybridModel and allowed it in exporter input-type validation.
Grouped MoE import modelopt/torch/export/plugins/megatron_importer.py	Compute initial global expert index via get_expert_model_parallel_rank() instead of hardcoded 0; assert divisibility and map local expert weights into weight{local_id} slots.
Descriptor mapping modelopt/torch/puzzletron/anymodel/.../model_descriptor_factory.py	Adjusted auto-detection mapping from gpt_oss_20b → gpt_oss.
Replacement library robustness modelopt/torch/puzzletron/replacement_library/build_replacement_library.py	Treat pandas NaN as missing checkpoint_dir when building block replacement paths.

Quantization Recipes & Misc

Layer / File(s)	Summary
Recipe tweak modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml	Set quantize.algorithm.layerwise: false with comment about VLM decoder nesting requirement.
TensorRT autotune arg modelopt/onnx/quantization/autotune/benchmark.py	Ensure --skipInference present in trtexec args for remote autotuning and log TensorRT API version info.

Tests, Regression & CI

Layer / File(s)	Summary
DFlash offline regression tests/regression/torch/speculative/test_dflash_offline.py	New regression test that dumps hidden states for synthetic conversations, runs offline DFlash training, and asserts loss decreases and final loss < 5.0.
Dataset utils unit tests tests/unit/torch/utils/test_dataset_utils.py	Added tests for _disable_use_cache and _forward_loop cache behavior and restoration.
CI workflow .github/workflows/gpu_tests.yml	Increased GPU job timeout from 60 to 75 minutes.

Release Cherry-pick Skill doc

Layer / File(s)	Summary
Skill doc .claude/skills/release-cherry-pick/SKILL.md	New skill documenting an interactive workflow to cherry-pick merged PRs labeled cherry-pick-<version> onto release/<version>, create a cherry-picks/release-<version> branch and PR, and label cherry-picked PRs with cherry-pick-done.

Sequence Diagram(s)

(omitted — conditions for sequence diagram generation not met)

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• NVIDIA/Model-Optimizer#1352: Adds the same .claude/skills/release-cherry-pick/SKILL.md release workflow file.

Suggested reviewers

• jenchen13
• realAsma
• cjluo-nv

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No security anti-patterns found. All potentially dangerous functions have proper justifications or use secure patterns.
Title check	✅ Passed	The title clearly identifies this as a cherry-pick PR combining multiple upstream PRs (#1352 #1351 #1330 #1354 #1355 #1360 #1342 #1324 #1340 #1368 #1373 #1359 #1361 #1325 #1369 #1370 #1371) into release/0.44.0, which matches the changeset content containing numerous bug fixes, features, and documentation updates across the codebase.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch cherry-picks/release-0.44.0

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1385/
Built to branch gh-pages at 2026-05-04 20:34 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[coderabbitai]

Actionable comments posted: 13

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py (1)

175-192: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Potential NameError when realizable_as_symlinks=True and skip_validation=True.

If both conditions are true, the model variable is never assigned (line 175-176 condition is False), but line 192 still attempts to use it, causing a NameError.

This appears to be pre-existing behavior, but since the line is modified, worth noting. Consider adding model to the condition or guarding line 192.

🛡️ One possible fix

         if args.save_models:
             checkpoint_dir = (
                 args.solutions_path.with_name(f"{args.solutions_path.stem}--checkpoints")
                 / f"solution_{i_solution}"
             )

             model_config.dtype = resolve_torch_dtype(getattr(args, "model_dtype", "torch.bfloat16"))
             Converter.copy_checkpoint_files(args.teacher_dir, checkpoint_dir)
             if realizable_as_symlinks:
                 if dist.is_master():
                     # TODO: Loo into internal Puzzleron code to see how to save as symlinks
                     # save_checkpoint_as_symlinks is currently not supported
                     pass
-            save_checkpoint_from_shards(model, checkpoint_dir, descriptor)
+            if not realizable_as_symlinks:
+                save_checkpoint_from_shards(model, checkpoint_dir, descriptor)

Or ensure model is always loaded when args.save_models is True:

-        if (args.save_models and not realizable_as_symlinks) or (not args.skip_validation):
+        if args.save_models or (not args.skip_validation):
             model = replacement_library.load_model(layer_replacements)

🤖 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/puzzletron/tools/validate_puzzle_with_multi_replacements.py`
around lines 175 - 192, The code can raise NameError because model is only set
when (args.save_models and not realizable_as_symlinks) or (not
args.skip_validation) but save_checkpoint_from_shards(model, ...) always runs;
ensure model exists before use by either: move/recompute model assignment so
replacement_library.load_model(layer_replacements) runs whenever
args.save_models is True (e.g., include realizable_as_symlinks branch), or guard
the call to save_checkpoint_from_shards behind the same condition that sets
model (use the same boolean logic involving args.save_models,
realizable_as_symlinks, and args.skip_validation) so save_checkpoint_from_shards
only receives a valid model; update related model_config assignment
(model_config.dtype) to match the chosen approach.

modelopt/torch/export/plugins/vllm_fakequant_hf.py (1)

632-649: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Disable and restore the plural fused-expert weight quantizers too.

The new fused-expert path strips *_weight_quantizers.<idx> from saved state, but Lines 632-649 still only disable singular *_weight_quantizer children. nn.ModuleList-backed fused-expert quantizers therefore remain enabled when modelopt_state() is captured, so the reload metadata can disagree with the fakequantized weights, and _check_all_weight_quantizers_disabled() will miss it for the same reason.

🧩 Suggested fix

         for _, module in model.named_modules():
             if isinstance(module, QuantModule):
                 for attr_name, quantizer in module.named_children():
-                    if not (attr_name.endswith("weight_quantizer") and quantizer.is_enabled):
-                        continue
-                    if isinstance(quantizer, SequentialQuantizer):
-                        quantizer.disable()
-                        for sub in quantizer:
-                            orig_rotate = sub._rotate
-                            if sub.rotate_is_enabled:
-                                sub._rotate = disable_rotate(sub)
-                            wqs_to_restore.append((sub, orig_rotate))
-                    elif isinstance(quantizer, TensorQuantizer):
-                        quantizer.disable()
-                        orig_rotate = quantizer._rotate
-                        if quantizer.rotate_is_enabled:
-                            quantizer._rotate = disable_rotate(quantizer)
-                        wqs_to_restore.append((quantizer, orig_rotate))
+                    quantizers: list[TensorQuantizer | SequentialQuantizer] = []
+                    if attr_name.endswith("weight_quantizers") and isinstance(quantizer, nn.ModuleList):
+                        quantizers = [
+                            q
+                            for q in quantizer
+                            if isinstance(q, (TensorQuantizer, SequentialQuantizer)) and q.is_enabled
+                        ]
+                    elif (
+                        attr_name.endswith("weight_quantizer")
+                        and isinstance(quantizer, (TensorQuantizer, SequentialQuantizer))
+                        and quantizer.is_enabled
+                    ):
+                        quantizers = [quantizer]
+
+                    for q in quantizers:
+                        if isinstance(q, SequentialQuantizer):
+                            q.disable()
+                            for sub in q:
+                                orig_rotate = sub._rotate
+                                if sub.rotate_is_enabled:
+                                    sub._rotate = disable_rotate(sub)
+                                wqs_to_restore.append((sub, orig_rotate))
+                        else:
+                            q.disable()
+                            orig_rotate = q._rotate
+                            if q.rotate_is_enabled:
+                                q._rotate = disable_rotate(q)
+                            wqs_to_restore.append((q, orig_rotate))

Please mirror the same plural/singular test in _check_all_weight_quantizers_disabled().

🤖 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/export/plugins/vllm_fakequant_hf.py` around lines 632 - 649,
The loop that disables per-module weight quantizers only checks for singular
child names ending with "weight_quantizer" and misses plural fused-expert
ModuleList children (e.g., "weight_quantizers"), so update the disable logic in
the block that iterates model.named_modules() to also detect attr_name ending
with "weight_quantizers" (or the ModuleList-typed plural container), iterate its
elements, disable each element (handling SequentialQuantizer and TensorQuantizer
cases the same as the singular branch), call disable_rotate when
rotate_is_enabled, and append the (quantizer, orig_rotate) tuples to
wqs_to_restore; then mirror the same plural-vs-singular checks in
_check_all_weight_quantizers_disabled() so it verifies both single
"weight_quantizer" children and ModuleList "weight_quantizers" entries when
asserting all weight quantizers are disabled.

🧹 Nitpick comments (3)

modelopt/onnx/quantization/ort_utils.py (1)

146-167: 💤 Low value

Consider adding a platform guard for defensive programming.

The function uses ctypes.windll (line 174) which only exists on Windows. While currently this function is only called within a Windows check (line 245-246), adding an early guard would prevent potential AttributeError if the function is ever called directly from elsewhere.

🛡️ Optional defensive guard

 def _load_extra_cudnn_dlls():
     """Load any cuDNN DLLs from site-packages that ORT's preload_dlls() missed.
     ...
     """
+    if platform.system() != "Windows":
+        return
+
     import ctypes
     import ctypes.wintypes

🤖 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/onnx/quantization/ort_utils.py` around lines 146 - 167, The
_load_extra_cudnn_dlls function should defensively return early on non-Windows
platforms to avoid AttributeError from using ctypes.windll; add a platform guard
at the top of _load_extra_cudnn_dlls (e.g., check sys.platform or os.name for
Windows) so the function exits immediately when not on Windows, while keeping
the existing _find_cudnn_bin_dir check and logging behavior unchanged.

tests/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py (1)

63-63: 💤 Low value

Minor: Redundant model instantiation.

get_tiny_llama() is called again just to access config.num_hidden_layers. Reuse the model from line 41 instead.

♻️ Suggested fix

-        assert cfg["num_hidden_layers"] == get_tiny_llama().config.num_hidden_layers
+        assert cfg["num_hidden_layers"] == model.config.num_hidden_layers

🤖 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/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py` at line 63,
The assertion redundantly calls get_tiny_llama() again to read
config.num_hidden_layers; instead reuse the already-instantiated model (variable
named model) created earlier (from get_tiny_llama()) and change the check to
compare cfg["num_hidden_layers"] with model.config.num_hidden_layers so you
remove the extra get_tiny_llama() call.

examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml (1)

18-23: ⚡ Quick win

Avoid requiring unrelated API secrets in this template.

Line 18–23 and Line 110–113 ask for API/OpenAI credentials even though this workflow’s listed tasks rely on HF-backed evals. Prefer scoping extra secrets only to tasks that need them to reduce unnecessary secret exposure.

Proposed tightening

-#   # Set additional unused but required environment variables:
-#   export API_KEY=xxxxxx
-#   export INFERENCE_API_KEY=xxxxxx
-#   export OPENAI_CLIENT_ID=xxxxxx
-#   export OPENAI_CLIENT_SECRET=xxxxxx
+#   # If a specific task/provider requires extra credentials, add them per-task.

   env_vars:
     HF_TOKEN: HF_TOKEN
     HF_HOME: HF_HOME
     VLLM_CACHE_ROOT: VLLM_CACHE_ROOT
-    API_KEY: API_KEY
-    INFERENCE_API_KEY: INFERENCE_API_KEY
-    OPENAI_CLIENT_ID: OPENAI_CLIENT_ID
-    OPENAI_CLIENT_SECRET: OPENAI_CLIENT_SECRET

Also applies to: 106-113

🤖 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/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml`
around lines 18 - 23, Remove the globally listed unrelated secret environment
variables (API_KEY, INFERENCE_API_KEY, OPENAI_CLIENT_ID, OPENAI_CLIENT_SECRET)
from the top-level template and instead scope them only to the specific tasks
that require them (or mark them as optional) so HF-backed evals don’t prompt for
unnecessary secrets; update the commented blocks around those env vars (the
comments containing API_KEY / INFERENCE_API_KEY / OPENAI_CLIENT_ID /
OPENAI_CLIENT_SECRET) and any duplicated occurrences later in the file to either
remove them or relocate them into the relevant task-level env sections.

🤖 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 @.claude/skills/release-cherry-pick/SKILL.md:
- Around line 20-24: The gh api command in the SKILL.md snippet uses per_page=50
which will truncate results when more than 50 matching PRs exist; update the
command that constructs the search (the gh api "search/issues?q=...+per_page=50"
invocation) to handle pagination—either increase per_page to the GitHub max
(100) and/or add --paginate to the gh api call so all pages are returned, and
ensure the pipeline (the jq and sort usage) still consumes streamed results
correctly.
- Around line 28-35: The current text says "creating it from the remote if it
doesn't exist locally" but the shown git fetch/checkout will fail if
release/<VERSION> doesn't exist on the remote; update the Step 3 wording to
reflect that, or replace the commands with logic that first git fetch origin,
then test for origin/release/<VERSION> (e.g., using git rev-parse --verify
origin/release/<VERSION>) and if it exists run git checkout release/<VERSION>,
otherwise create the branch locally from main (git checkout -b release/<VERSION>
origin/main) and push it upstream (git push -u origin release/<VERSION>).

In `@examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md`:
- Line 146: Fix the user-facing typos in the README: change "atmost" to "at
most" and replace the shorthand "hparams" with "hyperparameters" (e.g., the
sentence "Only considering atmost 40% for width and 20% for depth pruning
hparams" should read "Only considering at most 40% for width and 20% for depth
pruning hyperparameters"); apply the same wording corrections to the equivalent
lines around Line 184 and Line 228 so all tutorial sentences use "at most" and
"hyperparameters" consistently.

In `@examples/pruning/minitron/README.md`:
- Line 3: The intro sentence in README.md has a missing space after the comma
("evaluation,and"); update the sentence in the README (the Minitron tutorial
intro line) to read "evaluation, and" by inserting a space after the comma so it
becomes "evaluation, and vLLM deployment."

In `@examples/pruning/README.md`:
- Around line 295-300: The "Data Composition" guidance mixes the spellings
"pre-training" and "pretraining"; pick one consistent spelling (e.g.,
"pretraining") and update every occurrence in this section including the entries
under "Data Composition" and the related lines noted (previously 295-300 and
307-310) so all user-facing text uses the same term; search for both
"pre-training" and "pretraining" in the README.md and replace them consistently
(references: the "Data Composition" row and the adjacent bullet lines).

In `@modelopt/onnx/quantization/int4.py`:
- Around line 1322-1324: The bug is that new_tensor.get() is called
unconditionally which fails when cupy is unavailable; update the conversion
before calling numpy_helper.from_array to use np.asnumpy(new_tensor) when
has_cupy is True and otherwise pass the numpy array directly (i.e., use the same
pattern as other conversions: conditionally call np.asnumpy on new_tensor before
numpy_helper.from_array), referencing the variables/new calls new_tensor,
has_cupy, np.asnumpy and numpy_helper.from_array so the fix is applied where
new_tensor is created and passed into numpy_helper.from_array.

In `@modelopt/torch/distill/plugins/megatron.py`:
- Line 175: The current line uses submodule_name.replace(f".{match.group(0)}",
f".{new_layer_idx}") which replaces every identical ".<idx>" token; instead,
only replace the specific match instance. Locate the code referencing
submodule_name, match, new_layer_idx and change it to perform a single-instance
replacement using the match span (e.g., compute start,end = match.span(0) and
build new_submodule_name = submodule_name[:start] + f".{new_layer_idx}" +
submodule_name[end:]) or use re.sub with count=1 anchored to the exact match;
ensure only the matched token is changed.

In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 653-659: The current check sends any module with
gate_up_proj_weight_quantizers to _export_fused_experts regardless of whether
quantization is actually enabled; update the condition to only call
_export_fused_experts when the module both exposes
gate_up_proj_weight_quantizers and is actually quantized (e.g., change the if to
require hasattr(sub_module, "gate_up_proj_weight_quantizers") and
get_quantization_format(sub_module) != QUANTIZATION_NONE), or alternatively
inspect the items in sub_module.gate_up_proj_weight_quantizers for an
enabled/active flag before entering the
fsdp2_aware_weight_update/_export_fused_experts path so non-quantized
_QuantFusedExperts are not rewritten. Ensure you adjust/remove the following
branches accordingly: the if that references gate_up_proj_weight_quantizers, the
subsequent elif using get_quantization_format, and keep calls to
_export_fused_experts and fsdp2_aware_weight_update only when quantization is
active.

In `@modelopt/torch/puzzletron/tools/checkpoint_utils_hf.py`:
- Around line 231-235: The merge loop using full_sd.update(shard_sd) can
silently overwrite real tensors with placeholder zeros from other ranks; change
the merge to iterate each shard_sd's items and for each key k: if k not in
full_sd set full_sd[k]=v, else verify torch.equal(full_sd[k], v) (and raise a
ValueError on mismatch) or skip if you prefer keeping the first-seen owner—apply
this logic where full_sd and gathered are used to build the combined state_dict
(the shard_sd aggregation loop).

In `@modelopt/torch/utils/dataset_utils.py`:
- Around line 522-537: The probe tensor target_input is only built once before
the retry loop, so after an OOM and reducing target_data_batch you must rebuild
target_input using sample_input_single_batch.expand(...) with the new
target_data_batch before calling infer_method; move the expand logic (or
recreate target_input) inside the while loop (referencing
sample_input_single_batch, target_data_batch, and infer_method) so each retry
tests the correctly sized probe batch under the
torch.set_grad_enabled(enable_grad) context.

In `@modelopt/torch/utils/plugins/megatron_preprocess_data.py`:
- Around line 195-202: The exception handler in the try/except around
_Encoder.tokenizer.apply_chat_template currently prints the entire data payload
(variable data), which risks leaking sensitive content; replace the print of
json.dumps(data, ...) with a safe, minimal identifier and summary: log a stable
id if present (e.g., data.get("id") or data.get("row_id")), or compute and log a
short hash/fingerprint (e.g., SHA256 hex prefix) of the serialized row plus a
small sanitized field summary (e.g., length and first 100 chars of a sanitized
"text" field or list of keys), and include the exception message and context
(function name/tokenizer call) instead of the full payload in the raise/LOG call
in the except block that follows the call to
_Encoder.tokenizer.apply_chat_template.

In `@tests/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py`:
- Around line 120-132: The test_distributed_save_creates_valid_checkpoint
currently only verifies keys; update it to also load the saved tensors
referenced by SAFE_WEIGHTS_INDEX_NAME/weight_map (and shards in
SAFETENSORS_SUBBLOCKS_DIR_NAME) and compare each tensor's data to the original
model from get_tiny_llama() (use the same key mapping) using an element-wise
comparison (e.g., torch.equal or torch.allclose with a tolerance) to ensure
values match—mirror the approach used in test_saved_weights_match_original and
perform comparisons for every entry in index["weight_map"] to detect silent
corruption from _distributed_save_worker.

In `@tests/unit/torch/quantization/plugins/test_fused_experts.py`:
- Around line 386-441: The test is flaky because routing depends on random model
init; seed the model and make the forward_loop drive tokens to every expert
deterministically: call torch.manual_seed(...) before creating _TinyMoEModel()
(so weights are deterministic) and replace the current forward_loop with one
that, for idx in range(NUM_EXPERTS), constructs inputs that target expert idx
(e.g., unique per-expert input vectors or one-hot-like patterns) and runs m(x)
so each expert in model.moe.experts is exercised at least once before
calibration assertions.

---

Outside diff comments:
In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py`:
- Around line 632-649: The loop that disables per-module weight quantizers only
checks for singular child names ending with "weight_quantizer" and misses plural
fused-expert ModuleList children (e.g., "weight_quantizers"), so update the
disable logic in the block that iterates model.named_modules() to also detect
attr_name ending with "weight_quantizers" (or the ModuleList-typed plural
container), iterate its elements, disable each element (handling
SequentialQuantizer and TensorQuantizer cases the same as the singular branch),
call disable_rotate when rotate_is_enabled, and append the (quantizer,
orig_rotate) tuples to wqs_to_restore; then mirror the same plural-vs-singular
checks in _check_all_weight_quantizers_disabled() so it verifies both single
"weight_quantizer" children and ModuleList "weight_quantizers" entries when
asserting all weight quantizers are disabled.

In `@modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py`:
- Around line 175-192: The code can raise NameError because model is only set
when (args.save_models and not realizable_as_symlinks) or (not
args.skip_validation) but save_checkpoint_from_shards(model, ...) always runs;
ensure model exists before use by either: move/recompute model assignment so
replacement_library.load_model(layer_replacements) runs whenever
args.save_models is True (e.g., include realizable_as_symlinks branch), or guard
the call to save_checkpoint_from_shards behind the same condition that sets
model (use the same boolean logic involving args.save_models,
realizable_as_symlinks, and args.skip_validation) so save_checkpoint_from_shards
only receives a valid model; update related model_config assignment
(model_config.dtype) to match the chosen approach.

---

Nitpick comments:
In `@examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml`:
- Around line 18-23: Remove the globally listed unrelated secret environment
variables (API_KEY, INFERENCE_API_KEY, OPENAI_CLIENT_ID, OPENAI_CLIENT_SECRET)
from the top-level template and instead scope them only to the specific tasks
that require them (or mark them as optional) so HF-backed evals don’t prompt for
unnecessary secrets; update the commented blocks around those env vars (the
comments containing API_KEY / INFERENCE_API_KEY / OPENAI_CLIENT_ID /
OPENAI_CLIENT_SECRET) and any duplicated occurrences later in the file to either
remove them or relocate them into the relevant task-level env sections.

In `@modelopt/onnx/quantization/ort_utils.py`:
- Around line 146-167: The _load_extra_cudnn_dlls function should defensively
return early on non-Windows platforms to avoid AttributeError from using
ctypes.windll; add a platform guard at the top of _load_extra_cudnn_dlls (e.g.,
check sys.platform or os.name for Windows) so the function exits immediately
when not on Windows, while keeping the existing _find_cudnn_bin_dir check and
logging behavior unchanged.

In `@tests/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py`:
- Line 63: The assertion redundantly calls get_tiny_llama() again to read
config.num_hidden_layers; instead reuse the already-instantiated model (variable
named model) created earlier (from get_tiny_llama()) and change the check to
compare cfg["num_hidden_layers"] with model.config.num_hidden_layers so you
remove the extra get_tiny_llama() call.

🪄 Autofix (Beta)

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: c29df816-dd02-41f0-81cf-28e64d80103b

📥 Commits

Reviewing files that changed from the base of the PR and between b1ec471 and 1b2f029.

⛔ Files ignored due to path filters (1)

• examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/figures/learning_curves.png is excluded by !**/*.png

📒 Files selected for processing (44)

• .claude/skills/release-cherry-pick/SKILL.md
• CHANGELOG.rst
• docs/source/getting_started/windows/_installation_standalone.rst
• examples/dataset/MEGATRON_DATA_PREP.md
• examples/dataset/README.md
• examples/megatron_bridge/README.md
• examples/pruning/README.md
• examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/ABLATIONS.md
• examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md
• examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml
• examples/pruning/minitron/README.md
• examples/pruning/puzzletron/Llama-3.1-8B-Instruct.md
• examples/pruning/puzzletron/README.md
• examples/puzzletron/README.md
• examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml
• examples/speculative_decoding/main.py
• modelopt/onnx/quantization/autotune/benchmark.py
• modelopt/onnx/quantization/int4.py
• modelopt/onnx/quantization/ort_utils.py
• modelopt/torch/distill/plugins/megatron.py
• modelopt/torch/export/plugins/megatron_importer.py
• modelopt/torch/export/plugins/vllm_fakequant_hf.py
• modelopt/torch/export/quant_utils.py
• modelopt/torch/export/unified_export_hf.py
• modelopt/torch/export/unified_export_megatron.py
• modelopt/torch/puzzletron/anymodel/model_descriptor/model_descriptor_factory.py
• modelopt/torch/puzzletron/replacement_library/build_replacement_library.py
• modelopt/torch/puzzletron/tools/checkpoint_utils_hf.py
• modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py
• modelopt/torch/quantization/config.py
• modelopt/torch/quantization/conversion.py
• modelopt/torch/quantization/model_calib.py
• modelopt/torch/quantization/plugins/huggingface.py
• modelopt/torch/quantization/utils/__init__.py
• modelopt/torch/quantization/utils/core_utils.py
• modelopt/torch/utils/dataset_utils.py
• modelopt/torch/utils/plugins/megatron_preprocess_data.py
• modelopt/torch/utils/plugins/transformers_dataset.py
• modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml
• tests/gpu/torch/puzzletron/test_puzzletron.py
• tests/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py
• tests/regression/torch/speculative/test_dflash_offline.py
• tests/unit/torch/quantization/plugins/test_fused_experts.py
• tests/unit/torch/utils/test_dataset_utils.py

💤 Files with no reviewable changes (2)

• tests/gpu/torch/puzzletron/test_puzzletron.py
• modelopt/onnx/quantization/autotune/benchmark.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Handle pagination for repositories with many pending cherry-picks.

The GitHub API query uses per_page=50, which limits results to 50 PRs. If more than 50 PRs are labeled cherry-pick-<VERSION> without cherry-pick-done, later PRs will be silently omitted from the cherry-pick batch.

📄 Suggested approaches

Option 1: Increase the page size to a safer limit:

-gh api "search/issues?q=repo:NVIDIA/Model-Optimizer+is:pr+is:merged+base:main+label:cherry-pick-<VERSION>+-label:cherry-pick-done&sort=updated&order=asc&per_page=50" \
+gh api "search/issues?q=repo:NVIDIA/Model-Optimizer+is:pr+is:merged+base:main+label:cherry-pick-<VERSION>+-label:cherry-pick-done&sort=updated&order=asc&per_page=100" \

Option 2: Add pagination using --paginate:

-gh api "search/issues?q=repo:NVIDIA/Model-Optimizer+is:pr+is:merged+base:main+label:cherry-pick-<VERSION>+-label:cherry-pick-done&sort=updated&order=asc&per_page=50" \
+gh api --paginate "search/issues?q=repo:NVIDIA/Model-Optimizer+is:pr+is:merged+base:main+label:cherry-pick-<VERSION>+-label:cherry-pick-done&sort=updated&order=asc&per_page=100" \
   --jq '.items[] | [.number, .title, .pull_request.merged_at] | `@tsv`' \

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	```bash
	gh api "search/issues?q=repo:NVIDIA/Model-Optimizer+is:pr+is:merged+base:main+label:cherry-pick-<VERSION>+-label:cherry-pick-done&sort=updated&order=asc&per_page=50" \
	--jq '.items[] | [.number, .title, .pull_request.merged_at] | @tsv' \
	| sort -t$'\t' -k3
	```

🤖 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 @.claude/skills/release-cherry-pick/SKILL.md around lines 20 - 24, The gh api
command in the SKILL.md snippet uses per_page=50 which will truncate results
when more than 50 matching PRs exist; update the command that constructs the
search (the gh api "search/issues?q=...+per_page=50" invocation) to handle
pagination—either increase per_page to the GitHub max (100) and/or add
--paginate to the gh api call so all pages are returned, and ensure the pipeline
(the jq and sort usage) still consumes streamed results correctly.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Clarify release branch creation behavior.

The description states "creating it from the remote if it doesn't exist locally," but the provided commands will fail if the release/<VERSION> branch doesn't exist on the remote. The commands only check out an existing remote branch—they don't create it.

📝 Suggested clarification

Either update the description to reflect the actual behavior:

-Check out `release/<VERSION>`, creating it from the remote if it doesn't exist locally:
+Check out the existing `release/<VERSION>` branch from the remote:

Or provide commands that handle both cases:

+If the release branch doesn't exist yet, create it from main:
+
+```bash
+git fetch origin
+if git rev-parse --verify origin/release/<VERSION> >/dev/null 2>&1; then
+  git checkout release/<VERSION>
+else
+  git checkout -b release/<VERSION> origin/main
+  git push -u origin release/<VERSION>
+fi
+```
+
+Otherwise, check out the existing branch:
+
 ```bash
 git fetch origin release/<VERSION>
 git checkout release/<VERSION>

</details>

<details>
<summary>🤖 Prompt for AI Agents</summary>

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In @.claude/skills/release-cherry-pick/SKILL.md around lines 28 - 35, The
current text says "creating it from the remote if it doesn't exist locally" but
the shown git fetch/checkout will fail if release/ doesn't exist on the
remote; update the Step 3 wording to reflect that, or replace the commands with
logic that first git fetch origin, then test for origin/release/ (e.g.,
using git rev-parse --verify origin/release/) and if it exists run git
checkout release/, otherwise create the branch locally from main (git
checkout -b release/ origin/main) and push it upstream (git push -u
origin release/).

</details>

<!-- fingerprinting:phantom:triton:puma -->

<!-- 4e71b3a2 -->

<!-- This is an auto-generated comment by CodeRabbit -->

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Fix small typos in user-facing instructions.

There are a few wording typos in key tutorial lines (Line 146, Line 184, Line 228) that are quick to clean up.

Suggested edits

-Only considering atmost 40% for width and 20% for depth pruning hparams
+Only considering at most 40% for width and 20% for depth pruning hparams

-> While we use 96 nodes here for faster training, you can also run with 1 node. If you dont want to do full distillation run, you can stop earlier and take intermediate checkpoints as well.
+> While we use 96 nodes here for faster training, you can also run with 1 node. If you don't want to do a full distillation run, you can stop earlier and take intermediate checkpoints as well.

-The eval config xin [nemo_evaluator.yaml](nemo_evaluator.yaml) is for Slurm-based evaluation — it submits a vLLM serving job and runs evals against it.
+The eval config in [nemo_evaluator.yaml](nemo_evaluator.yaml) is for Slurm-based evaluation — it submits a vLLM serving job and runs evals against it.

Also applies to: 184-184, 228-228

🤖 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/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md` at line 146,
Fix the user-facing typos in the README: change "atmost" to "at most" and
replace the shorthand "hparams" with "hyperparameters" (e.g., the sentence "Only
considering atmost 40% for width and 20% for depth pruning hparams" should read
"Only considering at most 40% for width and 20% for depth pruning
hyperparameters"); apply the same wording corrections to the equivalent lines
around Line 184 and Line 228 so all tutorial sentences use "at most" and
"hyperparameters" consistently.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Fix the missing space in the intro sentence.

evaluation,and should read evaluation, and.

🤖 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/pruning/minitron/README.md` at line 3, The intro sentence in
README.md has a missing space after the comma ("evaluation,and"); update the
sentence in the README (the Minitron tutorial intro line) to read "evaluation,
and" by inserting a space after the comma so it becomes "evaluation, and vLLM
deployment."

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Use one spelling of “pretraining” in this section.

The edited guidance now mixes pre-training and pretraining, which reads like an edit artifact in user-facing docs.

Also applies to: 307-310

🧰 Tools

🪛 LanguageTool

[uncategorized] ~300-~300: Do not mix variants of the same word (‘pre-train’ and ‘pretrain’) within a single text.
Context: ...Composition** | • Standard models: 100% pre-training data
• Reasoning models: 70% reasoni...

(EN_WORD_COHERENCY)

---

[uncategorized] ~300-~300: Do not mix variants of the same word (‘pre-train’ and ‘pretrain’) within a single text.
Context: ...soning models: 70% reasoning data + 30% pre-training data | > [!TIP] > If you know the maxi...

(EN_WORD_COHERENCY)

🤖 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/pruning/README.md` around lines 295 - 300, The "Data Composition"
guidance mixes the spellings "pre-training" and "pretraining"; pick one
consistent spelling (e.g., "pretraining") and update every occurrence in this
section including the entries under "Data Composition" and the related lines
noted (previously 295-300 and 307-310) so all user-facing text uses the same
term; search for both "pre-training" and "pretraining" in the README.md and
replace them consistently (references: the "Data Composition" row and the
adjacent bullet lines).

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Merge logic is fragile with overlapping keys across ranks.

The full_sd.update(shard_sd) approach silently overwrites values when multiple ranks have the same key. If a rank provides a "placeholder" value (e.g., zeros) for keys it doesn't own, later ranks will overwrite earlier ranks' real values.

This is exactly what the distributed test does—each rank zeroes keys it doesn't own, so the final merged state will have zeros from the last rank overwriting real values from earlier ranks for those keys.

Consider either:

1. Only updating keys that are missing from full_sd, or
2. Documenting that ranks must provide disjoint key sets (no overlapping keys).

🛡️ Proposed fix to avoid overwriting existing values

             full_sd: dict[str, torch.Tensor] = {}
             for shard_sd in gathered:
                 if shard_sd is None:
                     continue
-                full_sd.update(shard_sd)
+                for k, v in shard_sd.items():
+                    if k not in full_sd:
+                        full_sd[k] = v

Alternatively, if overlapping keys with identical values are expected, you could add a check:

for k, v in shard_sd.items():
    if k in full_sd:
        if not torch.equal(full_sd[k], v):
            raise ValueError(f"Key {k} has conflicting values across ranks")
    else:
        full_sd[k] = v

🤖 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/puzzletron/tools/checkpoint_utils_hf.py` around lines 231 -
235, The merge loop using full_sd.update(shard_sd) can silently overwrite real
tensors with placeholder zeros from other ranks; change the merge to iterate
each shard_sd's items and for each key k: if k not in full_sd set full_sd[k]=v,
else verify torch.equal(full_sd[k], v) (and raise a ValueError on mismatch) or
skip if you prefer keeping the first-seen owner—apply this logic where full_sd
and gathered are used to build the combined state_dict (the shard_sd aggregation
loop).

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Rebuild the probe tensor after each OOM retry.

target_input is expanded once before the retry loop, so every retry still uses the original oversized batch. If the first estimate OOMs, this path just counts target_data_batch down to 1 without ever validating a smaller candidate, which can collapse auto-batching to 1 even when a larger batch would fit.

Suggested fix

-        target_input = sample_input_single_batch.expand(
-            [
-                target_data_batch if index == 0 else dim
-                for index, dim in enumerate(sample_input_single_batch.shape)
-            ]
-        )
-
         # For some models on multi GPU, we observe the memory per batch is not a constant.
         # So we just test the target batch size and make sure we do not go OOM.
         while target_data_batch > 1:
+            target_input = sample_input_single_batch.expand(
+                [
+                    target_data_batch if index == 0 else dim
+                    for index, dim in enumerate(sample_input_single_batch.shape)
+                ]
+            )
             with torch.set_grad_enabled(enable_grad):
                 try:
                     infer_method(target_input)
                     break
                 except torch.cuda.OutOfMemoryError:  # pragma: no cover - GPU OOM retry path
+                    torch.cuda.empty_cache()  # pragma: no cover
                     target_data_batch = target_data_batch // 2  # pragma: no cover

🤖 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/dataset_utils.py` around lines 522 - 537, The probe
tensor target_input is only built once before the retry loop, so after an OOM
and reducing target_data_batch you must rebuild target_input using
sample_input_single_batch.expand(...) with the new target_data_batch before
calling infer_method; move the expand logic (or recreate target_input) inside
the while loop (referencing sample_input_single_batch, target_data_batch, and
infer_method) so each retry tests the correctly sized probe batch under the
torch.set_grad_enabled(enable_grad) context.

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Avoid dumping the full failing sample into stdout.

Line 199 prints the entire JSON row when apply_chat_template() fails. On real SFT/pretraining data that can leak prompts, tool payloads, PII, or secrets into shared logs, and very large rows will also blow up log volume. Log a stable identifier / field summary instead of the full payload.

🔒 Suggested fix

                 try:
                     text = _Encoder.tokenizer.apply_chat_template(value, tokenize=False, **kwargs)
                 except Exception as e:
+                    sample_id = data.get("id") or data.get("uuid") or data.get("conversation_id")
                     print(
-                        f"apply_chat_template failed: {e}\nData:\n{json.dumps(data, indent=2, default=str)}",
+                        "apply_chat_template failed "
+                        f"for sample={sample_id!r}: {e}. "
+                        f"Available keys: {sorted(data.keys())}",
                         flush=True,
                     )
                     raise

As per coding guidelines, SECURITY.md says to treat all model/weight/config artifacts and data as untrusted.

🤖 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/megatron_preprocess_data.py` around lines 195 -
202, The exception handler in the try/except around
_Encoder.tokenizer.apply_chat_template currently prints the entire data payload
(variable data), which risks leaking sensitive content; replace the print of
json.dumps(data, ...) with a safe, minimal identifier and summary: log a stable
id if present (e.g., data.get("id") or data.get("row_id")), or compute and log a
short hash/fingerprint (e.g., SHA256 hex prefix) of the serialized row plus a
small sanitized field summary (e.g., length and first 100 chars of a sanitized
"text" field or list of keys), and include the exception message and context
(function name/tokenizer call) instead of the full payload in the raise/LOG call
in the except block that follows the call to
_Encoder.tokenizer.apply_chat_template.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Multi-process test should verify saved tensor values, not just keys.

The test only checks that weight_map.keys() match expected keys, but doesn't verify the actual tensor values like test_saved_weights_match_original does. Given the merge logic concern in the main implementation, verifying tensor correctness would catch silent data corruption.

💚 Proposed enhancement to verify tensor values

         shard_files = list((tmp_path / SAFETENSORS_SUBBLOCKS_DIR_NAME).glob("*.safetensors"))
         assert len(shard_files) > 0
+
+        # Verify saved values match original model
+        reloaded_sd = {}
+        for shard in shard_files:
+            reloaded_sd.update(safe_load_file(str(shard)))
+
+        original_sd = {k: v.cpu() for k, v in model.state_dict().items()}
+        for key in expected_keys:
+            torch.testing.assert_close(reloaded_sd[key], original_sd[key])

🤖 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/gpu/torch/puzzletron/tools/test_save_ckpt_from_shards.py` around lines
120 - 132, The test_distributed_save_creates_valid_checkpoint currently only
verifies keys; update it to also load the saved tensors referenced by
SAFE_WEIGHTS_INDEX_NAME/weight_map (and shards in
SAFETENSORS_SUBBLOCKS_DIR_NAME) and compare each tensor's data to the original
model from get_tiny_llama() (use the same key mapping) using an element-wise
comparison (e.g., torch.equal or torch.allclose with a tolerance) to ensure
values match—mirror the approach used in test_saved_weights_match_original and
perform comparisons for every entry in index["weight_map"] to detect silent
corruption from _distributed_save_worker.

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Make expert coverage deterministic in this calibration test.

This assertion requires every expert to receive tokens, but _TinyMoEModel() is created from unseeded random weights and the router path is data-dependent. That makes the test flaky: a different initialization can leave one or more experts untouched and keep their amax at None even when the fused-experts path is working correctly. Please force a deterministic routing pattern that exercises all experts before asserting on every per-expert quantizer.

🤖 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/quantization/plugins/test_fused_experts.py` around lines 386
- 441, The test is flaky because routing depends on random model init; seed the
model and make the forward_loop drive tokens to every expert deterministically:
call torch.manual_seed(...) before creating _TinyMoEModel() (so weights are
deterministic) and replace the current forward_loop with one that, for idx in
range(NUM_EXPERTS), constructs inputs that target expert idx (e.g., unique
per-expert input vectors or one-hot-like patterns) and runs m(x) so each expert
in model.moe.experts is exercised at least once before calibration assertions.

[codecov]

Codecov Report

❌ Patch coverage is 64.52906% with 177 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.88%. Comparing base (b1ec471) to head (3720a7a).

Files with missing lines	Patch %	Lines
modelopt/onnx/quantization/int4.py	67.42%	86 Missing ⚠️
modelopt/onnx/quantization/ort_utils.py	15.09%	45 Missing ⚠️
...lopt/torch/puzzletron/tools/checkpoint_utils_hf.py	25.00%	18 Missing ⚠️
modelopt/torch/export/plugins/vllm_fakequant_hf.py	67.85%	9 Missing ⚠️
modelopt/torch/export/plugins/megatron_importer.py	0.00%	6 Missing ⚠️
...pt/torch/utils/plugins/megatron_preprocess_data.py	0.00%	5 Missing ⚠️
modelopt/torch/quantization/plugins/huggingface.py	90.00%	3 Missing ⚠️
modelopt/torch/distill/plugins/megatron.py	0.00%	2 Missing ⚠️
modelopt/torch/quantization/model_calib.py	60.00%	2 Missing ⚠️
...delopt/torch/utils/plugins/transformers_dataset.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@                Coverage Diff                 @@
##           release/0.44.0    #1385      +/-   ##
==================================================
+ Coverage           75.41%   76.88%   +1.46%     
==================================================
  Files                 463      463              
  Lines               50208    50387     +179     
==================================================
+ Hits                37865    38739     +874     
+ Misses              12343    11648     -695     

Flag	Coverage Δ
examples	41.58% <16.63%> (+0.87%)	⬆️
gpu	59.65% <59.31%> (+0.52%)	⬆️
unit	52.84% <51.30%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.