feat(megatron-bridge): adopt TEGroupedMLP per-expert quantizers (#1550) for W4A16 NVFP4 four_over_six

[yueshen2016]

What does this PR do ?

Type of change: Bug fix + integration validation

Overview: Adopts the per-expert TEGroupedMLP quantization from #1550 for the
Nemotron-Nano-3 W4A16 NVFP4 four_over_six recipe, and fixes one bug that blocks it.

Until now this recipe had to force SequentialMLP (non-grouped experts), because
TEGroupedLinear._process_quantizer_amax asserted v.numel() == 1 — i.e. per-tensor only —
while four_over_six is static per-block NVFP4 whose _amax has one entry per block.
With #1550 that restriction is gone, so the grouped path now works and SequentialMLP is no
longer required for this recipe.

The fix

_QuantTEGroupedLinear._setup() builds the new per-expert GroupedQuantizer from
self.default_quant_desc_weight, but that attribute does not resolve: _ParallelLinear derives
from QuantModule, not _QuantLinear, and _ParallelLinear._setup() itself references the class
attribute directly (plugins/custom.py). Building any grouped-experts model fails with:

AttributeError: QuantTEColumnParallelGroupedLinear object has no attribute default_quant_desc_weight

Changed to _QuantLinear.default_quant_desc_weight, matching the non-grouped path.

Validation — grouped vs non-grouped, end to end

Nemotron-Nano-3, W4A16 NVFP4 four_over_six, TEGroupedMLP, 4x GB200:

stage	result
PTQ	6382 quantizers inserted, checkpoint saved
QAD	200 iters, logits-distillation loss 3.37e-2 -> 1.91e-2, router seq_load_balancing_loss steady ~1.013
HF export	52 shards, 18487 keys, All enabled NVFP4 weight quantizers have calibrated scales
compressed-tensors	lm_head.weight_packed U8 [131072, 1344], 72 exclusions
serving	loads on stock vLLM 0.26.0 at TP=2 and generates coherently

The grouped export is structurally indistinguishable from the SequentialMLP one:

• identical key count (18487) and routed-expert key count (17920)
• identical quant config: W4A16_NVFP4, num_bits=4, block=16, dynamic=False, ignore=72
• routed-expert scales stay per-block, not per-tensor —
  experts.0.up_proj.weight_scale [1856, 168] and down_proj.weight_scale [2688, 116]
  (168 = 2688/16, 116 = 1856/16)

Accuracy of the grouped path vs SequentialMLP has not been measured yet; only structural
equivalence and serving are verified here.

Stacking

This branch is stacked and should merge after both:

• Support per expert weight quantizer in TEGroupedMLP #1550 (jennifchen/te_per_expert) — the per-expert GroupedQuantizer this builds on
• feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support #2055 — the untied-lm_head quantization fix for Megatron-Bridge

The diff against main therefore includes commits from both. Our own commits are:

• 72d3cb5e7 — from feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support #2055
• c53ee3041 — merge of Support per expert weight quantizer in TEGroupedMLP #1550, resolving two conflicts (below)
• 5f2f3de3c — the default_quant_desc_weight fix

Merge conflict resolutions

• quantization/plugins/megatron.py — Support per expert weight quantizer in TEGroupedMLP #1550 registers ModelOpt extra-state callbacks for every
  QuantModule, which subsumes the narrower output_layer special-case from feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support #2055. Took Support per expert weight quantizer in TEGroupedMLP #1550's
  version; the _resolve_output_layer_untied helper and the sharded_state_dict buffer
  materialisation from feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support #2055 are unaffected and still used.
• examples/megatron_bridge/distill.py — kept both validation blocks, and took Support per expert weight quantizer in TEGroupedMLP #1550's
  most_recent_k=args.checkpoint_keep_last. Note its default is 5; checkpoints here are ~413 GB
  each, so pass --checkpoint_keep_last 2 on constrained filesystems.

Usage

# PTQ / export with fused (grouped) experts
python examples/megatron_bridge/quantize.py --grouped_experts ...
python examples/megatron_bridge/export_quantized_megatron_to_hf.py --grouped_experts ...
# QAD: simply omit --student_nongrouped_experts

Testing

Full PTQ -> QAD -> export -> compressed-tensors -> vLLM serving on Nemotron-Nano-3, as tabulated
above. The _QuantTEGroupedLinear change is exercised by every grouped-experts run; without it,
_setup() raises before any model can be built.

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes — the grouped path is opt-in via --grouped_experts; the default remains non-grouped.
• Did you write any new necessary tests?: No — relies on the tests added in Support per expert weight quantizer in TEGroupedMLP #1550.
• Did you add or update any necessary documentation?: No
• Did you update Changelog?: No — covered by Support per expert weight quantizer in TEGroupedMLP #1550's entry.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added independent per-expert weight quantization for Transformer Engine grouped MoE layers.
  • Added optional compiled execution for grouped expert quantization; eager execution remains the default.
  • Added grouped-expert and randomized calibration options to quantization workflows.
  • Added improved support for exporting quantized grouped MoE checkpoints to Hugging Face and Megatron formats.
  • Added distillation support for SFT-masked data and quantized MoE students.

• Bug Fixes

  • Improved restoration and validation of quantization scales in distributed and sharded checkpoints.
  • Added clearer handling for uncalibrated NVFP4 exports.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change adds independent per-expert TEGroupedMLP quantization, optional compiled weight processing, per-expert checkpoint and export handling, grouped-quantizer calibration support, randomized calibration offsets, NVFP4 validation, and expanded QAD/SFT distillation tooling.

Changes

Per-expert quantization workflow

Layer / File(s)	Summary
Grouped quantizer contracts and calibration modelopt/torch/quantization/nn/*, modelopt/torch/quantization/model_calib.py, modelopt/torch/quantization/plugins/custom.py, modelopt/torch/quantization/utils/core_utils.py, modelopt/torch/quantization/config.py	Adds GroupedQuantizer and AnyQuantizer. Calibration, state restoration, representative lookup, and amax synchronization now support grouped quantizers.
TEGroupedMLP runtime and validation modelopt/torch/quantization/plugins/transformer_engine.py, tests/gpu_megatron/torch/quantization/plugins/test_megatron.py, CHANGELOG.rst	Uses one weight quantizer per expert. An environment variable enables compiled per-expert processing outside calibration. Tests cover amax shapes, compile parity, gradients, expert identity, and quantization loss.
Expert checkpoint state and export modelopt/torch/quantization/plugins/megatron.py, modelopt/torch/export/unified_export_megatron.py, modelopt/torch/export/plugins/vllm_fakequant_hf.py, modelopt/torch/opt/plugins/mcore_dist_checkpointing.py	Preserves per-expert scales and global expert identities across sharded state, extra-state restoration, Megatron export, and grouped quantizer disabling.
Calibration and bridge tooling modelopt/torch/utils/dataset_utils.py, modelopt/torch/utils/plugins/megatron_calibration.py, examples/megatron_bridge/quantize.py, examples/megatron_bridge/export_quantized_megatron_to_hf.py	Adds randomized packed-data offsets, grouped-expert options, MTP configuration handling, and NVFP4 calibration validation with an optional BF16 fallback.
QAD and SFT distillation workflow examples/megatron_bridge/distill.py, modelopt/torch/distill/plugins/megatron.py	Adds checkpoint-initialized students, SFT-masked datasets, quantized-student options, QAD settings, synchronous checkpointing, extended timeouts, and one-time NVFP4 quantizer promotion.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• NVIDIA/Model-Optimizer#1550: Directly covers earlier per-expert TEGroupedMLP quantization changes extended here across export, calibration, checkpointing, and tests.

Suggested labels: cherry-pick-done, cherry-pick-0.46.0

Suggested reviewers: cjluo-nv, shengliangxu

Sequence Diagram(s)

sequenceDiagram
  participant Calibration
  participant GroupedQuantizer
  participant TEGroupedMLP
  participant CheckpointExporter
  Calibration->>GroupedQuantizer: collect and synchronize per-expert amax
  GroupedQuantizer->>TEGroupedMLP: quantize each expert weight
  TEGroupedMLP->>CheckpointExporter: expose per-expert scales and expert identities
  CheckpointExporter-->>GroupedQuantizer: restore grouped quantizer state

Loading

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 65.63% 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.
Title check	✅ Passed	The title clearly describes the main change: adopting TEGroupedMLP per-expert quantizers for the W4A16 NVFP4 four_over_six recipe.
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	The PR adds one weights_only=False load with an inline comment confirming sibling-rank torch.save bytes are not user-supplied; no other flagged patterns or dependency additions were found.

✨ Finishing Touches 💡 2

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch qad/te-per-expert-1550

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch qad/te-per-expert-1550

---

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

[coderabbitai]

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.

👉 Steps to fix this

Actionable comments posted: 14

Caution

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

⚠️ Outside diff range comments (1)

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

632-639: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Flatten nested SequentialQuantizer members before saving rotation state.

GroupedQuantizer accepts SequentialQuantizer members, but _rotate exists only on their TensorQuantizer leaves. Accessing sub._rotate raises AttributeError, and restoring the container would not restore the leaves. Iterate over the leaves when recording, disabling, and restoring rotation state.

🤖 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 - 639,
Update the quantizer handling around SequentialQuantizer and GroupedQuantizer to
flatten nested SequentialQuantizer members to their TensorQuantizer leaves
before recording rotation state. Record each leaf’s _rotate value, disable
rotation on each leaf, and ensure restoration targets those same leaves rather
than container quantizers.

🧹 Nitpick comments (5)

examples/megatron_bridge/distill.py (3)

348-348: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Function-scope imports without a justification comment in examples/megatron_bridge/distill.py and modelopt/torch/distill/plugins/megatron.py. Both new imports sit inside a function body with no stated reason. The coding guidelines allow a local import only for a circular dependency, an optional dependency, or an unusually heavy import, and require a brief comment naming that reason.

• examples/megatron_bridge/distill.py#L348-L348: move from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec to the module header. It runs on every _build_model_provider call but only the branch at line 350 uses it.
• modelopt/torch/distill/plugins/megatron.py#L624-L627: keep the tensor_quantizer import local if it breaks a circular dependency, and add a one-line comment naming that dependency. Otherwise move it to the module header.

As per coding guidelines: "Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment."

🤖 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/megatron_bridge/distill.py` at line 348, Move
get_te_mamba_stack_spec from the local scope of _build_model_provider in
examples/megatron_bridge/distill.py to the module header. In
modelopt/torch/distill/plugins/megatron.py, move the tensor_quantizer import to
the module header unless it is required to avoid a circular dependency; if it
must remain local, add a brief comment identifying that dependency.

Source: Coding guidelines

---

65-73: 📐 Maintainability & Code Quality | 🔵 Trivial

TODO tracked: replace the DistillationProvider.provide patch.

The TODO states that this class-level patch must go away once Megatron-Bridge exposes a student-initialization hook. The registry is keyed by id(distill_provider), so the workaround also depends on the provider object staying alive for the process lifetime.

Do you want me to open an issue to track the upstream hook and the removal of this patch?

🤖 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/megatron_bridge/distill.py` around lines 65 - 73, Track the upstream
hook request as an issue, documenting that the class-level
DistillationProvider.provide patch and _MEGATRON_STUDENT_CKPT_PATHS registry
should be removed once Megatron-Bridge exposes student initialization before
distillation conversion.

---

365-368: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Read FORCE_NO_PER_TOKEN_LOSS once.

The same environment variable is read in two places with inverted conditions. The two settings must stay consistent: calculate_per_token_loss and average_in_collective are a matched pair. A future edit to one site can silently desynchronize them.

Compute the value once in main and use it at both sites.

♻️ Proposed fix

 def main(args: argparse.Namespace):
+    # SFT with CP>1 needs per-token loss on the provider and average_in_collective=False on DDP.
+    # Read the override once so the two settings cannot desynchronize.
+    per_token_loss = args.sft and os.environ.get("FORCE_NO_PER_TOKEN_LOSS", "0") != "1"
     checkpoint_dir = os.path.join(args.output_dir, "checkpoints")

         if args.sft:
-            # Finetuning (SFT) with context parallel (CP>1) requires per-token loss so the
-            # response loss-mask reduces correctly across the CP ranks.
-            provider.calculate_per_token_loss = os.environ.get("FORCE_NO_PER_TOKEN_LOSS", "0") != "1"
+            # Finetuning (SFT) with context parallel (CP>1) requires per-token loss so the
+            # response loss-mask reduces correctly across the CP ranks.
+            provider.calculate_per_token_loss = per_token_loss

-            # Finetuning (SFT) with CP>1 requires per-token loss (set on the provider) and
-            # average_in_collective=False (the per-token loss is summed, not averaged, in the collective).
-            average_in_collective=(not args.sft) or os.environ.get("FORCE_NO_PER_TOKEN_LOSS", "0") == "1",
+            # The per-token loss is summed, not averaged, in the collective.
+            average_in_collective=not per_token_loss,

Also applies to: 484-486

🤖 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/megatron_bridge/distill.py` around lines 365 - 368, Read
FORCE_NO_PER_TOKEN_LOSS once in main, store the resulting boolean, and reuse it
when assigning both provider.calculate_per_token_loss and the corresponding
average_in_collective setting. Remove the second direct environment lookup while
preserving the existing inverted relationship between these matched settings.

modelopt/torch/distill/plugins/megatron.py (1)

644-655: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use rank-aware logging for the promotion messages.

Both messages run on every rank, so a large job emits one copy per rank. setup_distillation_config in this same file already guards logger.info with a rank check at the tensor-and-context-parallel rank. Apply the same guard here, or use print_rank_0 / warn_rank_0.

♻️ Proposed fix

                 if amax is None:
                     # Uncalibrated: leave it alone rather than silently changing precision.
-                    logger.warning(f"NVFP4 weight quantizer {name} has no _amax; not promoted.")
+                    if parallel_state.get_tensor_and_context_parallel_rank() == 0:
+                        logger.warning(f"NVFP4 weight quantizer {name} has no _amax; not promoted.")
                     n_skipped += 1
                     continue

-            if n_promoted or n_skipped:
+            if (n_promoted or n_skipped) and parallel_state.get_tensor_and_context_parallel_rank() == 0:
                 logger.info(

As per coding guidelines: "Develop distributed code with rank-aware logging such as print_rank_0 or warn_rank_0".

🤖 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/distill/plugins/megatron.py` around lines 644 - 655, Update
the NVFP4 promotion logging near StaticBlockScaleQuantizer.from_tensor_quantizer
so both the missing-_amax warning and promotion summary are emitted only from
the tensor-and-context-parallel rank, using the existing rank-check pattern from
setup_distillation_config or the available print_rank_0/warn_rank_0 helpers.

Source: Coding guidelines

modelopt/torch/utils/dataset_utils.py (1)

765-765: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Document random_offset in the docstring Args list.

get_dataset_dataloader is a public API and every other parameter is documented. Add an entry for random_offset describing the extra sequence-length of headroom and the dropped leading tokens.

As per coding guidelines: "Document public and higher-level APIs with docstrings, including examples when useful".

🤖 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` at line 765, Update the docstring Args
section for get_dataset_dataloader to document random_offset, describing that it
enables extra sequence-length headroom and drops the corresponding leading
tokens. Keep the documentation consistent with the existing parameter entries.

Source: Coding guidelines

🤖 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/megatron_bridge/distill.py`:
- Around line 311-314: Resolve the conflict between the validation and fallback
around the argument-parsing logic: either remove the
hf_export_path/student_hf_model ValueError so the existing student_hf_path
fallback remains usable for exports, or retain the validation and update the
--student_hf_model help text to explicitly state it is required with
--hf_export_path. Ensure the chosen behavior is consistent and the fallback is
not misleading.
- Line 348: Move the get_te_mamba_stack_spec import from _build_model_provider
to the module-level imports in examples/megatron_bridge/distill.py, unless
modelopt.torch.nas.plugins.megatron is intentionally optional; in that case,
retain the local import and add a brief comment explaining the
optional-dependency justification.
- Around line 192-198: Update the --sft argument help text in
parser.add_argument to name FinetuningDatasetConfig instead of
GPTSFTDatasetConfig, matching the configuration class constructed by the SFT
branch while preserving the rest of the description.

In `@examples/megatron_bridge/export_quantized_megatron_to_hf.py`:
- Around line 136-142: Update the resolved configuration construction to treat
explicit null values in model_cfg as missing, so defaults such as
mtp_num_layers=0 are preserved. Move the isinstance(model_cfg, dict) check
outside the comprehension, and only read keys from model_cfg when it is a
dictionary; otherwise retain defaults unchanged.
- Around line 125-127: Update the run_config resolution near run_config so it
uses the checkpoint iteration/path selected by the loader, including tracker,
ckpt_step, direct iteration, or checkpoint state metadata. Remove the
lexicographic sorted(...)[0] selection and derive run_config.yaml from the same
resolved checkpoint location used for loading. Preserve the root-level fallback
only when the loader resolves the root checkpoint itself.

In `@modelopt/torch/distill/plugins/megatron.py`:
- Around line 630-634: Update the quantizer scan near the loop over named
modules to traverse only the student model rather than self, preventing teacher
submodules from being promoted. Preserve the existing TensorQuantizer and
StaticBlockScaleQuantizer filtering and promotion behavior for student weight
quantizers.

In `@modelopt/torch/export/unified_export_megatron.py`:
- Around line 1062-1065: Update the assertion message in the grouped_wq
validation to reference pre-0.46 single-quantizer checkpoints instead of
pre-0.47, matching the documented per-expert layout break and re-quantization
release. Leave the assertion condition and type reporting unchanged.
- Around line 1147-1149: Update the grouped-expert handling around
_get_quantized_state and _record_excluded_module so unquantized experts are
added to exclude_modules when qformat is None. For each expert, format the
unformatted prefix with its global_id, append ".", and record that resulting
module name, while preserving existing behavior for quantized experts.
- Around line 1142-1145: Restrict the temporary amax assignment in the
weight-quantizer handling to instances of TensorQuantizer, preserving the
existing _amax and enabled checks. Import TensorQuantizer alongside
GroupedQuantizer, and ensure only TensorQuantizer objects are appended to
temp_amax_wqs so delegated SequentialQuantizer members are not modified or
reset.

In `@modelopt/torch/quantization/nn/modules/tensor_quantizer.py`:
- Around line 1870-1879: Update GroupedQuantizer.__init__ to reject an empty
quantizers argument by raising ValueError before representative access can
occur, while preserving the existing member-type validation. Add a CPU
regression test that verifies constructing GroupedQuantizer without members
raises ValueError.

In `@modelopt/torch/quantization/plugins/custom.py`:
- Around line 143-155: Update _has_complete_static_nvfp4_weight_state and the
surrounding restore path to inspect every member of a GroupedQuantizer against
its matching expert weight, rather than only member zero; ensure recalibration
and restoration preserve each member’s distinct amax state and prevent later
expert processing from skipping invalid members. Add a regression test covering
at least two experts with different amax states, while preserving backward
compatibility for serialized ModeloptBaseConfig and QuantizeConfig objects.

In `@modelopt/torch/quantization/plugins/megatron.py`:
- Around line 241-253: Restrict the missing-state skip in the quantizer
restoration flow to names matching the per-expert `weight_quantizer.<i>` pattern
handled by the fallback in `quantizer_substate`. For non-expert names such as
`input_quantizer` and `output_quantizer`, preserve the previous loud failure
behavior instead of continuing with default properties.

In `@modelopt/torch/utils/dataset_utils.py`:
- Around line 715-718: Update the random_offset handling in
get_megatron_calibration_dataloader to derive the leading-token offset from a
fixed, shared seed rather than the unseeded global random.randint call. Reuse
the existing deterministic seed approach used by the multi-source shuffle so
every rank produces the identical token_stream and runs remain reproducible.

In `@tests/gpu_megatron/torch/quantization/plugins/test_megatron.py`:
- Around line 1026-1027: Deep-copy quant_cfg independently before each
mtq.quantize call in the affected test and in
_test_te_grouped_sharded_state_dict_global_expert_identity_helper, so neither
invocation nor later parametrized tests share a potentially mutated
configuration. Follow the existing copy.deepcopy pattern used by the nearby
quantization call sites.

---

Outside diff comments:
In `@modelopt/torch/export/plugins/vllm_fakequant_hf.py`:
- Around line 632-639: Update the quantizer handling around SequentialQuantizer
and GroupedQuantizer to flatten nested SequentialQuantizer members to their
TensorQuantizer leaves before recording rotation state. Record each leaf’s
_rotate value, disable rotation on each leaf, and ensure restoration targets
those same leaves rather than container quantizers.

---

Nitpick comments:
In `@examples/megatron_bridge/distill.py`:
- Line 348: Move get_te_mamba_stack_spec from the local scope of
_build_model_provider in examples/megatron_bridge/distill.py to the module
header. In modelopt/torch/distill/plugins/megatron.py, move the tensor_quantizer
import to the module header unless it is required to avoid a circular
dependency; if it must remain local, add a brief comment identifying that
dependency.
- Around line 65-73: Track the upstream hook request as an issue, documenting
that the class-level DistillationProvider.provide patch and
_MEGATRON_STUDENT_CKPT_PATHS registry should be removed once Megatron-Bridge
exposes student initialization before distillation conversion.
- Around line 365-368: Read FORCE_NO_PER_TOKEN_LOSS once in main, store the
resulting boolean, and reuse it when assigning both
provider.calculate_per_token_loss and the corresponding average_in_collective
setting. Remove the second direct environment lookup while preserving the
existing inverted relationship between these matched settings.

In `@modelopt/torch/distill/plugins/megatron.py`:
- Around line 644-655: Update the NVFP4 promotion logging near
StaticBlockScaleQuantizer.from_tensor_quantizer so both the missing-_amax
warning and promotion summary are emitted only from the
tensor-and-context-parallel rank, using the existing rank-check pattern from
setup_distillation_config or the available print_rank_0/warn_rank_0 helpers.

In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the docstring Args section for get_dataset_dataloader to
document random_offset, describing that it enables extra sequence-length
headroom and drops the corresponding leading tokens. Keep the documentation
consistent with the existing parameter entries.

🪄 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: a8ec06a9-7a9f-42e4-9786-0b06303a155c

📥 Commits

Reviewing files that changed from the base of the PR and between fed1980 and 5f2f3de.

📒 Files selected for processing (20)

• CHANGELOG.rst
• examples/megatron_bridge/distill.py
• examples/megatron_bridge/export_quantized_megatron_to_hf.py
• examples/megatron_bridge/quantize.py
• modelopt/torch/distill/plugins/megatron.py
• modelopt/torch/export/plugins/vllm_fakequant_hf.py
• modelopt/torch/export/unified_export_megatron.py
• modelopt/torch/opt/plugins/mcore_dist_checkpointing.py
• modelopt/torch/quantization/config.py
• modelopt/torch/quantization/model_calib.py
• modelopt/torch/quantization/nn/__init__.py
• modelopt/torch/quantization/nn/modules/tensor_quantizer.py
• modelopt/torch/quantization/plugins/custom.py
• modelopt/torch/quantization/plugins/megatron.py
• modelopt/torch/quantization/plugins/transformer_engine.py
• modelopt/torch/quantization/utils/core_utils.py
• modelopt/torch/utils/dataset_utils.py
• modelopt/torch/utils/plugins/megatron_calibration.py
• tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
• tests/unit/torch/quantization/test_tensor_quantizer_cpu.py

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Help text names a config class the code does not use.

The --sft help says GPTSFTDatasetConfig, but the SFT branch at line 426 builds a FinetuningDatasetConfig. Use the actual class name so users can match the flag to the code.

📝 Proposed fix

         help="SFT-masked distillation: read raw prompt-completion jsonl from --sft_dataset_root and "
         "mask the loss to the completion (assistant response) tokens. Uses GPTSFTDatasetConfig + the "
+        "real (HuggingFace) tokenizer instead of the pretraining GPTDataset + NullTokenizer.",

Replace with:

        help="SFT-masked distillation: read raw prompt-completion jsonl from --sft_dataset_root and "
        "mask the loss to the completion (assistant response) tokens. Uses FinetuningDatasetConfig + "
        "the real (HuggingFace) tokenizer instead of the pretraining GPTDataset + NullTokenizer.",

🤖 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/megatron_bridge/distill.py` around lines 192 - 198, Update the --sft
argument help text in parser.add_argument to name FinetuningDatasetConfig
instead of GPTSFTDatasetConfig, matching the configuration class constructed by
the SFT branch while preserving the rest of the description.

[coderabbitai]

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

The new check makes the existing --student_hf_model fallback unreachable for export.

Line 311 rejects --hf_export_path without --student_hf_model. Line 313 then defaults student_hf_model to student_hf_path. The default can never apply to an export run, because line 311 already raised. Users who previously relied on the fallback for export now get an error.

Choose one behavior. Either keep the fallback and drop the check, or keep the check and state in the --student_hf_model help that it is required with --hf_export_path.

🤖 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/megatron_bridge/distill.py` around lines 311 - 314, Resolve the
conflict between the validation and fallback around the argument-parsing logic:
either remove the hf_export_path/student_hf_model ValueError so the existing
student_hf_path fallback remains usable for exports, or retain the validation
and update the --student_hf_model help text to explicitly state it is required
with --hf_export_path. Ensure the chosen behavior is consistent and the fallback
is not misleading.

[coderabbitai]

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

Move get_te_mamba_stack_spec to the top-level imports.

This import sits inside _build_model_provider with no justification comment, and it runs on every provider build even though only the branch at line 350 uses it. The coding guidelines allow a local import only for a circular dependency, an optional dependency, or an unusually heavy import, and require a brief comment naming the reason.

If modelopt.torch.nas.plugins.megatron is an optional dependency here, keep the local import and add the comment. Otherwise move it to the module header.

♻️ Proposed fix

         provider.mtp_num_layers = 0
-        from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec
-
         if quantized and args.student_nongrouped_experts:

Add at the module header instead:

from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec

As per coding guidelines: "Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment."

🤖 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/megatron_bridge/distill.py` at line 348, Move the
get_te_mamba_stack_spec import from _build_model_provider to the module-level
imports in examples/megatron_bridge/distill.py, unless
modelopt.torch.nas.plugins.megatron is intentionally optional; in that case,
retain the local import and add a brief comment explaining the
optional-dependency justification.

Source: Coding guidelines

[coderabbitai]

🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Description: Inspect the iter_* resolution used by the Megatron checkpoint loader.
set -euo pipefail

fd -t f 'mbridge.py' modelopt | xargs -r rg -n -C 5 '_get_modelopt_checkpoint_path|iter_'
rg -n -C 5 'def _get_modelopt_checkpoint_path' --glob '*.py'

Repository: NVIDIA/Model-Optimizer

Length of output: 2401

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- exporter ---'
cat -n examples/megatron_bridge/export_quantized_megatron_to_hf.py | sed -n '95,145p'

printf '%s\n' '--- loader call sites and definitions ---'
rg -n -C 12 'load_modelopt_megatron_checkpoint|_get_modelopt_checkpoint_path|def _get_modelopt_checkpoint_path' \
  modelopt examples --glob '*.py'

printf '%s\n' '--- iteration directory fixtures/usages ---'
rg -n -C 4 'iter_[0-9]+|run_config\.yaml' modelopt examples --glob '*.py' --glob '*.yaml'

Repository: NVIDIA/Model-Optimizer

Length of output: 29061

---

🌐 Web query:

Megatron-Bridge _get_modelopt_checkpoint_path source iter_ latest checkpoint

💡 Result:

The function _get_modelopt_checkpoint_path is a utility in the Megatron-Bridge library (specifically within bridge.training.post_training.checkpointing) used to determine the correct path for ModelOpt operations [1][2][3]. This function effectively handles iteration directories by checking the provided checkpoint_path for the presence of ModelOpt states [1][4]. It specifically looks for modelopt_state directories either within specific iteration subdirectories (which follow the iter_* naming convention) or within the root checkpoint directory [1][2]. In the context of Megatron-LM and Megatron-Bridge checkpointing: - Megatron-Bridge typically uses a base directory structure where iteration-specific checkpoints are stored in iter_N/ subdirectories [5][6]. - When determining the "latest" checkpoint, Megatron Bridge relies on a tracker file named latest_train_state.pt located in the root directory [5][7]. - The _get_modelopt_checkpoint_path utility abstracts the path resolution to ensure ModelOpt-specific data is correctly identified and retrieved, regardless of whether it is stored at the root or within an iteration-specific folder [1][3]. For users needing to load a specific checkpoint iteration, the ckpt_step parameter is used to override the default tracker behavior and target a specific iter_N folder directly [5][7]. If you are performing manual conversion or export tasks, logic similar to _get_modelopt_checkpoint_path often involves scanning the checkpoint_path for iter_* directories, sorting them numerically, and selecting the highest index to identify the latest iteration [8].

Citations:

• 1: https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.post_training.checkpointing.html
• 2: https://docs.nvidia.com/nemo/megatron-bridge/0.5.0/apidocs/bridge/bridge.training.post_training.checkpointing.html
• 3: https://docs.nvidia.com/nemo/megatron-bridge/nightly/apidocs/bridge/bridge.training.post_training.checkpointing.html
• 4: https://docs.nvidia.com/nemo/megatron-bridge/0.4.0/apidocs/bridge/bridge.training.post_training.checkpointing.html
• 5: https://docs.nvidia.com/nemo/megatron-bridge/latest/training/checkpointing.html
• 6: https://docs.nvidia.com/nemo/megatron-bridge/0.4.1/training/checkpointing.html
• 7: https://docs.nvidia.com/nemo/megatron-bridge/nightly/training/checkpointing.html
• 8: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py

---

🌐 Web query:

site:github.com/NVIDIA-NeMo/Megatron-Bridge "_get_modelopt_checkpoint_path"

💡 Result:

The term get_modelopt_checkpoint_path does not appear to be a standard, publicly documented public API function within the NVIDIA Model Optimizer (ModelOpt) or the NVIDIA NeMo/Megatron-Bridge libraries. Based on an analysis of the NVIDIA NeMo/Megatron-Bridge repository [1][2][3], interactions with ModelOpt checkpoints typically involve: 1. Manual Directory Specification: Scripts generally expect users to provide a path to the checkpoint directory directly via arguments (e.g., --megatron-save-path or --pretrained-checkpoint) [3][4][5]. 2. Internal Logic: The codebase contains logic within internal utility modules (such as src/megatron/bridge/training/post_training/checkpointing.py) that handles the identification and selection of ModelOpt checkpoints [6]. This logic often involves checking for specific configuration files like run_config.yaml or iterating through subdirectories labeled iter [1][2]. 3. Model Configuration: Quantization configurations are typically handled through functions like get_modelopt_torch_quantization_config [4] or by defining model specifications using tools like get_gpt_modelopt_spec [7]. If you encountered this name in a specific codebase, it is likely an internal helper function or a custom utility specific to that project rather than a public function provided by NVIDIA ModelOpt. You may wish to search for the definition of this function within your local project files to understand its specific implementation and usage.

Citations:

• 1: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py
• 2: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/auto_bridge.py
• 3: https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/quantization
• 4: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/quantization/quantize.py
• 5: [OMNIML-2910] ModelOpt QAT support NVIDIA-NeMo/Megatron-Bridge#873
• 6: NVIDIA-NeMo/Megatron-Bridge@190e6e4
• 7: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/gpt_provider.py

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

url='https://raw.githubusercontent.com/NVIDIA-NeMo/Megatron-Bridge/main/src/megatron/bridge/training/post_training/checkpointing.py'
curl -fsSL "$url" | rg -n -C 20 '_get_modelopt_checkpoint_path|iter_|latest_train_state|run_config'

Repository: NVIDIA/Model-Optimizer

Length of output: 5010

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

url='https://raw.githubusercontent.com/NVIDIA-NeMo/Megatron-Bridge/main/src/megatron/bridge/training/checkpointing.py'
curl -fsSL "$url" | rg -n -C 25 '_resolve_checkpoint_iteration|_DIRECT_ITERATION_DIR_SENTINEL|latest_train_state|ckpt_step'

printf '%s\n' '--- exporter path construction ---'
cat -n examples/megatron_bridge/export_quantized_megatron_to_hf.py | sed -n '145,180p'
rg -n -C 8 'megatron_path|checkpoint_export_paths' examples/megatron_bridge/export_quantized_megatron_to_hf.py

Repository: NVIDIA/Model-Optimizer

Length of output: 28745

---

Resolve run_config.yaml from the checkpoint iteration selected by the loader. sorted(...)[0] selects the first lexicographic match, while the loader uses the tracker, ckpt_step, direct iteration paths, or checkpoint state metadata. Do not replace this with reverse sorting; reuse the loader’s resolution logic so mtp_num_layers matches the loaded checkpoint.

🤖 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/megatron_bridge/export_quantized_megatron_to_hf.py` around lines 125
- 127, Update the run_config resolution near run_config so it uses the
checkpoint iteration/path selected by the loader, including tracker, ckpt_step,
direct iteration, or checkpoint state metadata. Remove the lexicographic
sorted(...)[0] selection and derive run_config.yaml from the same resolved
checkpoint location used for loading. Preserve the root-level fallback only when
the loader resolves the root checkpoint itself.

[coderabbitai]

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

An explicit null in the YAML defeats the default.

If run_config.yaml contains mtp_num_layers: null, model_cfg.get(key, default) returns None, and the following merge keeps None because resolved overrides defaults. The provider then receives mtp_num_layers=None instead of 0.

The if isinstance(model_cfg, dict) filter is also loop-invariant. Move it out of the comprehension.

♻️ Proposed fix

     model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg
-    resolved = {
-        key: model_cfg.get(key, default)
-        for key, default in defaults.items()
-        if isinstance(model_cfg, dict)
-    }
-    resolved = {**defaults, **resolved}
+    resolved = dict(defaults)
+    if isinstance(model_cfg, dict):
+        for key, default in defaults.items():
+            value = model_cfg.get(key)
+            resolved[key] = default if value is None else value

📝 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

	model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg
	resolved = {
	key: model_cfg.get(key, default)
	for key, default in defaults.items()
	if isinstance(model_cfg, dict)
	}
	resolved = {**defaults, **resolved}
	model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg
	resolved = dict(defaults)
	if isinstance(model_cfg, dict):
	for key, default in defaults.items():
	value = model_cfg.get(key)
	resolved[key] = default if value is None else value

🤖 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/megatron_bridge/export_quantized_megatron_to_hf.py` around lines 136
- 142, Update the resolved configuration construction to treat explicit null
values in model_cfg as missing, so defaults such as mtp_num_layers=0 are
preserved. Move the isinstance(model_cfg, dict) check outside the comprehension,
and only read keys from model_cfg when it is a dictionary; otherwise retain
defaults unchanged.

[coderabbitai]

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Reject empty GroupedQuantizer instances.

__init__ accepts zero members, but forward and representative property access use self[0]. An empty container fails later with IndexError. Raise ValueError during construction and add a CPU regression test.

🤖 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/quantization/nn/modules/tensor_quantizer.py` around lines 1870
- 1879, Update GroupedQuantizer.__init__ to reject an empty quantizers argument
by raising ValueError before representative access can occur, while preserving
the existing member-type validation. Add a CPU regression test that verifies
constructing GroupedQuantizer without members raises ValueError.

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Preserve each grouped weight quantizer during restore.

_QuantTEGroupedLinear.modelopt_post_restore calls this generic restore path. The path checks only member zero. If member zero requires recalibration, GroupedQuantizer.reset_amax() clears every member, but max_calibrate(grouped, ...) forwards through only member zero. The later per-expert loop then skips members with _amax is None.

Restore and validate each grouped weight quantizer against its matching expert weight. Add a checkpoint restore regression test with at least two experts and different amax states.

As per coding guidelines, “Preserve backward compatibility for serialized Pydantic-based ModelOpt configuration and checkpoint objects such as ModeloptBaseConfig and QuantizeConfig.”

🤖 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/quantization/plugins/custom.py` around lines 143 - 155, Update
_has_complete_static_nvfp4_weight_state and the surrounding restore path to
inspect every member of a GroupedQuantizer against its matching expert weight,
rather than only member zero; ensure recalibration and restoration preserve each
member’s distinct amax state and prevent later expert processing from skipping
invalid members. Add a regression test covering at least two experts with
different amax states, while preserving backward compatibility for serialized
ModeloptBaseConfig and QuantizeConfig objects.

Source: Coding guidelines

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Restrict the silent skip to per-expert quantizer names.

The comment justifies the fallback for weight_quantizer.<i> names only. The continue at Line 253 applies to every quantizer name. A missing input_quantizer or output_quantizer entry, which previously failed loudly, is now skipped without any signal, and the module restores partially with default quantizer properties.

Keep the fallback for per-expert names and keep a loud failure for the rest.

🐛 Proposed fix

                 quantizer_substate = quantizer_state.get(name)
                 if quantizer_substate is None:
                     # Per-expert quantizers ("weight_quantizer.<i>") are saved per EP rank, so a
                     # module loaded at smaller EP (e.g. EP1 export from an EP16 ckpt) has more
                     # experts than the saved state. Per-expert properties are uniform across
                     # experts (amax rides separately as globally-indexed sharded tensors), so
                     # fall back to expert 0's state. Rewrite only the expert index right after
                     # weight_quantizer, preserving deeper suffixes (e.g. a SequentialQuantizer level
                     # weight_quantizer.<i>.<lvl>); non-expert names are left unchanged.
                     fallback = re.sub(r"(weight_quantizer)\.\d+", r"\1.0", name)
+                    if fallback == name:
+                        raise KeyError(
+                            f"Quantizer state for {name} is missing from the checkpoint."
+                        )
                     quantizer_substate = quantizer_state.get(fallback)
                 if quantizer_substate is None:
                     continue

📝 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

	quantizer_substate = quantizer_state.get(name)
	if quantizer_substate is None:
	# Per-expert quantizers ("weight_quantizer.<i>") are saved per EP rank, so a
	# module loaded at smaller EP (e.g. EP1 export from an EP16 ckpt) has more
	# experts than the saved state. Per-expert properties are uniform across
	# experts (amax rides separately as globally-indexed sharded tensors), so
	# fall back to expert 0's state. Rewrite only the expert index right after
	# weight_quantizer, preserving deeper suffixes (e.g. a SequentialQuantizer level
	# weight_quantizer.<i>.<lvl>); non-expert names are left unchanged.
	fallback = re.sub(r"(weight_quantizer)\.\d+", r"\1.0", name)
	quantizer_substate = quantizer_state.get(fallback)
	if quantizer_substate is None:
	continue
	quantizer_substate = quantizer_state.get(name)
	if quantizer_substate is None:
	# Per-expert quantizers ("weight_quantizer.<i>") are saved per EP rank, so a
	# module loaded at smaller EP (e.g. EP1 export from an EP16 ckpt) has more
	# experts than the saved state. Per-expert properties are uniform across
	# experts (amax rides separately as globally-indexed sharded tensors), so
	# fall back to expert 0's state. Rewrite only the expert index right after
	# weight_quantizer, preserving deeper suffixes (e.g. a SequentialQuantizer level
	# weight_quantizer.<i>.<lvl>); non-expert names are left unchanged.
	fallback = re.sub(r"(weight_quantizer)\.\d+", r"\1.0", name)
	if fallback == name:
	raise KeyError(
	f"Quantizer state for {name} is missing from the checkpoint."
	)
	quantizer_substate = quantizer_state.get(fallback)
	if quantizer_substate is None:
	continue

🤖 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/quantization/plugins/megatron.py` around lines 241 - 253,
Restrict the missing-state skip in the quantizer restoration flow to names
matching the per-expert `weight_quantizer.<i>` pattern handled by the fallback
in `quantizer_substate`. For non-expert names such as `input_quantizer` and
`output_quantizer`, preserve the previous loud failure behavior instead of
continuing with default properties.

[coderabbitai]

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Seed the offset RNG so ranks agree and runs reproduce.

random.randint uses the unseeded global RNG. get_megatron_calibration_dataloader builds this dataset independently on every rank and then shards it with a DistributedSampler that assumes an identical dataset on all ranks. With random_offset=True, each rank drops a different number of leading tokens, so the same sampler index maps to different rows on different ranks. The calibration set then contains duplicated and missing windows, and the run is not reproducible.

Derive the offset from a fixed seed, as the multi-source shuffle at line 870 already does.

🛠️ Proposed fix

     if random_offset:
         max_off = min(seq_length, max(0, len(token_stream) - num_rows * seq_length))
         if max_off > 0:
-            token_stream = token_stream[random.randint(0, max_off):]
+            # Fixed seed: every rank must derive the same offset, otherwise the
+            # DistributedSampler shards a different token grid per rank.
+            token_stream = token_stream[random.Random(0).randint(0, max_off) :]

Note: the avoid-random-python static-analysis hint on this line is a false positive. This offset is calibration sampling, not a security value.

📝 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

	if random_offset:
	max_off = min(seq_length, max(0, len(token_stream) - num_rows * seq_length))
	if max_off > 0:
	token_stream = token_stream[random.randint(0, max_off):]
	if random_offset:
	max_off = min(seq_length, max(0, len(token_stream) - num_rows * seq_length))
	if max_off > 0:
	# Fixed seed: every rank must derive the same offset, otherwise the
	# DistributedSampler shards a different token grid per rank.
	token_stream = token_stream[random.Random(0).randint(0, max_off) :]

🧰 Tools

🪛 ast-grep (0.45.0)

[info] 717-717: use secrets package over random package
Context: random.randint(0, max_off)
Note: [CWE-330] Use of Insufficiently Random Values.

(avoid-random-python)

🤖 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 715 - 718, Update the
random_offset handling in get_megatron_calibration_dataloader to derive the
leading-token offset from a fixed, shared seed rather than the unseeded global
random.randint call. Reuse the existing deterministic seed approach used by the
multi-source shuffle so every rank produces the identical token_stream and runs
remain reproducible.

Source: Linters/SAST tools

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Deep-copy quant_cfg before each mtq.quantize call.

quant_cfg is a module-level constant supplied by @pytest.mark.parametrize (mtq.FP8_DEFAULT_CFG / mtq.NVFP4_DEFAULT_CFG). The other call sites in this file pass copy.deepcopy(...) for that reason (Line 749, Line 817, Line 884). Here the same object is passed twice and is shared with every other parametrized test in the session. If mtq.quantize mutates the config, the second call and later tests observe the mutation.

The same applies to Line 1129 in _test_te_grouped_sharded_state_dict_global_expert_identity_helper.

🐛 Proposed fix

-    mtq.quantize(te_grouped, quant_cfg, forward)
-    mtq.quantize(sequential, quant_cfg, forward)
+    mtq.quantize(te_grouped, copy.deepcopy(quant_cfg), forward)
+    mtq.quantize(sequential, copy.deepcopy(quant_cfg), forward)

🤖 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_megatron/torch/quantization/plugins/test_megatron.py` around lines
1026 - 1027, Deep-copy quant_cfg independently before each mtq.quantize call in
the affected test and in
_test_te_grouped_sharded_state_dict_global_expert_identity_helper, so neither
invocation nor later parametrized tests share a potentially mutated
configuration. Follow the existing copy.deepcopy pattern used by the nearby
quantization call sites.

[jenchen13]

this is a copy from my PR #1550, you can remove this as you will rebase on my later

[jenchen13]

please remove any duplicate code between this PR and #1550 as you will rebase on top of it

[jenchen13]

duplicate code copied from #1550