feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support - #2055
feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support#2055yueshen2016 wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughMegatron distillation now restores student checkpoint state before conversion, supports SFT and QAD settings, validates NVFP4 calibration state, tracks output-layer quantization state, and adds optional random-offset calibration packing. ChangesMegatron QAD distillation and quantization
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant DistillationProvider.provide
participant Megatron student
participant Megatron checkpoint
participant Non-quantized teacher
participant ModelOpt distillation conversion
participant AutoBridge.export_ckpt
DistillationProvider.provide->>Megatron student: Build student
DistillationProvider.provide->>Megatron checkpoint: Restore Megatron and ModelOpt state
DistillationProvider.provide->>Non-quantized teacher: Build teacher
DistillationProvider.provide->>ModelOpt distillation conversion: Apply KD conversion
ModelOpt distillation conversion->>AutoBridge.export_ckpt: Export student checkpoint
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
c643835 to
f40d6ed
Compare
yueshen2016
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🧹 Nitpick comments (7)
modelopt/torch/distill/plugins/megatron.py (1)
623-627: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the justification comment for the local import.
The coding guidelines allow local imports only for circular dependencies, optional dependencies, or unusually heavy imports, and they require a brief explanatory comment. Add one line that states the reason, or move the import to the top of the module if no reason applies.
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 `@modelopt/torch/distill/plugins/megatron.py` around lines 623 - 627, Add a brief explanatory comment immediately before the local import of StaticBlockScaleQuantizer and TensorQuantizer in the _modelopt_nvfp4_promoted branch, documenting the applicable circular, optional, or heavy-import reason; if none applies, move these imports to module scope instead.Source: Coding guidelines
examples/megatron_bridge/export_quantized_megatron_to_hf.py (2)
151-164: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the NVFP4 detection predicate into one shared helper. Both sites reimplement the same test:
TensorQuantizerinstance check,StaticBlockScaleQuantizerexclusion,_num_bits == (2, 1),block_sizes["scale_bits"] == (4, 3), and an_amaxpresence check. The two copies have already diverged in their treatment of_global_amax, and the exclusion rule is wrong in one of them. A single helper in the quantization package keeps the classification rules in one place.
examples/megatron_bridge/export_quantized_megatron_to_hf.py#L151-L164: replace the inline predicate with a call to the shared helper, and keep only the export-specific reporting here.modelopt/torch/distill/plugins/megatron.py#L630-L655: replace the inline predicate with the same helper, and keep only the promotion call and counters here.🤖 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 151 - 164, Extract the duplicated NVFP4 classification logic into one shared helper in the quantization package, covering the TensorQuantizer type check, StaticBlockScaleQuantizer exclusion, _num_bits, scale_bits, and _amax rules consistently. In examples/megatron_bridge/export_quantized_megatron_to_hf.py lines 151-164, replace the inline predicate with the helper and retain only export-specific uncalibrated reporting. In modelopt/torch/distill/plugins/megatron.py lines 630-655, use the same helper and retain only the promotion call and counters.
166-182: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueDocument
MODELOPT_ALLOW_UNCALIBRATED_NVFP4inexamples/megatron_bridge/README.md. Describe the BF16 fallback and the required environment variable. The partial disable path already exports consistent quantization metadata.🤖 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 166 - 182, Document MODELOPT_ALLOW_UNCALIBRATED_NVFP4 in examples/megatron_bridge/README.md, including that uncalibrated NVFP4 weights fall back to BF16 and export requires setting this environment variable to 1. Note that the partial disable path preserves consistent quantization metadata.examples/megatron_bridge/distill.py (3)
308-308: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a justification comment for the local import, or move it to the top of the file.
from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_specat Line 326 is a local import inside_build_model_provider, which runs once per provider (student and teacher). No comment explains why it is local. If it is local because of an optional/heavy dependency (e.g. Mamba), state that explicitly; otherwise move the import to the top of the file with the other imports.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."
♻️ Proposed fix
provider.mtp_num_layers = 0 - from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec + # Local import: modelopt.torch.nas.plugins.megatron pulls in optional Mamba stack + # utilities that are heavy/optional dependencies, not needed by every caller of this script. + from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_specAlso applies to: 325-327
🤖 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 308, Update the local import of get_te_mamba_stack_spec inside _build_model_provider: move it to the file-level imports unless it is intentionally local for an optional or unusually heavy dependency; in that case, keep it local and add a brief comment explaining the dependency and reason for deferred loading.Source: Coding guidelines
343-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate the
FORCE_NO_PER_TOKEN_LOSSoverride logic.The same environment-variable check appears twice: once at Line 346 to set
provider.calculate_per_token_loss, and again at Line 463 to computeaverage_in_collective. Both must stay logically inverse for correctness (the comment at Line 462 confirms this dependency). Compute the flag once and reuse it at both sites to remove the duplication and prevent future drift between the two checks.♻️ Proposed fix
+ use_per_token_loss = os.environ.get("FORCE_NO_PER_TOKEN_LOSS", "0") != "1" + def _build_model_provider(hf_path, load_weights=True, quantized=True): ... 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" + provider.calculate_per_token_loss = use_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", + average_in_collective=(not args.sft) or not use_per_token_loss,Also applies to: 461-463
🤖 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 343 - 346, Compute the FORCE_NO_PER_TOKEN_LOSS-derived flag once in the surrounding distillation setup, then reuse that variable for both provider.calculate_per_token_loss and the average_in_collective calculation near the existing line-462 logic. Preserve their required inverse relationship and remove both duplicated environment-variable checks.
55-128: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy liftDocument the Megatron-Bridge compatibility boundary.
_super_classis an internal implementation detail, not a stable API. This workaround targetsnemo:26.06and older containers, while newer Megatron-Bridge versions move KD conversion to_convert_hook. Document the supported version and update this patch when that API is required.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/megatron_bridge/distill.py` around lines 55 - 128, Document near _distill_provide_with_megatron_student that this _super_class-based workaround is supported only for nemo:26.06 and older containers. Note that newer Megatron-Bridge versions perform KD conversion through _convert_hook, and state that this patch must be updated when that API is required.modelopt/torch/utils/dataset_utils.py (1)
875-875: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winDocument or reject
random_offsetwithout packing.The forwarding call at Line 875 is reached only when
pack=True. Withpack=False,random_offset=Truesilently has no effect. Add anArgsentry that states thepack=Truerequirement, or raiseValueErrorat the interface.As per coding guidelines, public and higher-level APIs must be documented with docstrings.
🤖 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 875, Update the public API docstring near the forwarding call to document that random_offset is effective only when pack=True, or validate the interface by raising ValueError when random_offset=True and pack=False. Ensure the behavior is explicit for callers and preserve the existing packed-data flow.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/export_quantized_megatron_to_hf.py`:
- Around line 124-127: Update the argument parser and model-provider
configuration in export_quantized_megatron_to_hf.py to expose mtp_num_layers and
moe_grouped_gemm as optional command-line arguments. Preserve the provider
defaults when these options are omitted, and pass explicitly supplied values
through so MTP-enabled and grouped-expert checkpoints build matching model
structures for load_modelopt_megatron_checkpoint.
- Around line 151-164: Update the quantizer scan over
unwrapped_model.named_modules() to include StaticBlockScaleQuantizer instances,
while retaining the TensorQuantizer type guard and NVFP4 detection. Check
is_enabled before adding entries so disabled quantizers are excluded, and ensure
enabled static-block quantizers with a missing _global_amax are reported
alongside missing _amax cases.
In `@modelopt/torch/quantization/plugins/megatron.py`:
- Around line 433-452: Update the guard in the weight-quantizer materialization
block around _wq and _block to validate self.weight.shape[1] is divisible by
_block, matching the later _process_quantizer_amax view requirement rather than
checking total element count. When _block is missing or this input-dimension
divisibility check fails, emit a warning identifying that buffer materialization
was skipped; keep the existing buffer initialization unchanged for valid block
sizes.
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the public get_dataset_dataloader signature so random_offset
follows all existing parameters and is keyword-only, preserving the current
positional order and preventing distributed and sampler_kwargs from being
rebound for existing callers.
---
Nitpick comments:
In `@examples/megatron_bridge/distill.py`:
- Line 308: Update the local import of get_te_mamba_stack_spec inside
_build_model_provider: move it to the file-level imports unless it is
intentionally local for an optional or unusually heavy dependency; in that case,
keep it local and add a brief comment explaining the dependency and reason for
deferred loading.
- Around line 343-346: Compute the FORCE_NO_PER_TOKEN_LOSS-derived flag once in
the surrounding distillation setup, then reuse that variable for both
provider.calculate_per_token_loss and the average_in_collective calculation near
the existing line-462 logic. Preserve their required inverse relationship and
remove both duplicated environment-variable checks.
- Around line 55-128: Document near _distill_provide_with_megatron_student that
this _super_class-based workaround is supported only for nemo:26.06 and older
containers. Note that newer Megatron-Bridge versions perform KD conversion
through _convert_hook, and state that this patch must be updated when that API
is required.
In `@examples/megatron_bridge/export_quantized_megatron_to_hf.py`:
- Around line 151-164: Extract the duplicated NVFP4 classification logic into
one shared helper in the quantization package, covering the TensorQuantizer type
check, StaticBlockScaleQuantizer exclusion, _num_bits, scale_bits, and _amax
rules consistently. In
examples/megatron_bridge/export_quantized_megatron_to_hf.py lines 151-164,
replace the inline predicate with the helper and retain only export-specific
uncalibrated reporting. In modelopt/torch/distill/plugins/megatron.py lines
630-655, use the same helper and retain only the promotion call and counters.
- Around line 166-182: Document MODELOPT_ALLOW_UNCALIBRATED_NVFP4 in
examples/megatron_bridge/README.md, including that uncalibrated NVFP4 weights
fall back to BF16 and export requires setting this environment variable to 1.
Note that the partial disable path preserves consistent quantization metadata.
In `@modelopt/torch/distill/plugins/megatron.py`:
- Around line 623-627: Add a brief explanatory comment immediately before the
local import of StaticBlockScaleQuantizer and TensorQuantizer in the
_modelopt_nvfp4_promoted branch, documenting the applicable circular, optional,
or heavy-import reason; if none applies, move these imports to module scope
instead.
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 875: Update the public API docstring near the forwarding call to document
that random_offset is effective only when pack=True, or validate the interface
by raising ValueError when random_offset=True and pack=False. Ensure the
behavior is explicit for callers and preserve the existing packed-data flow.
🪄 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: 48705570-8935-47c6-831c-f0454fc8f152
📒 Files selected for processing (7)
examples/megatron_bridge/distill.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/quantize.pymodelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/megatron_calibration.py
| "mtp_num_layers": 0, # QAD ckpt has MTP dropped | ||
| }, | ||
| init_model_parallel=True, | ||
| moe_grouped_gemm=False, # QAD ckpt is non-grouped (per-block NVFP4) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Do not hardcode the QAD checkpoint shape for every export.
mtp_num_layers: 0 and moe_grouped_gemm=False are now fixed for all inputs. This script also exports checkpoints produced by quantize.py, which can contain MTP layers, and it supports --export_extra_modules for MTP right below at Line 196. A grouped-expert checkpoint also no longer matches the built model structure, so load_modelopt_megatron_checkpoint sees a different parameter layout.
Expose both values as command-line arguments, and keep the provider defaults when the user does not set them.
🤖 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 124
- 127, Update the argument parser and model-provider configuration in
export_quantized_megatron_to_hf.py to expose mtp_num_layers and moe_grouped_gemm
as optional command-line arguments. Preserve the provider defaults when these
options are omitted, and pass explicitly supplied values through so MTP-enabled
and grouped-expert checkpoints build matching model structures for
load_modelopt_megatron_checkpoint.
| include_labels: bool = False, | ||
| apply_chat_template: bool = False, | ||
| pack: bool = False, | ||
| random_offset: bool = False, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python - <<'PY'
import ast
from pathlib import Path
for path in Path(".").rglob("*.py"):
try:
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
except (OSError, SyntaxError):
continue
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
if isinstance(node.func, ast.Name):
name = node.func.id
elif isinstance(node.func, ast.Attribute):
name = node.func.attr
else:
name = None
if name == "get_dataset_dataloader" and len(node.args) >= 10:
print(f"{path}:{node.lineno}: {len(node.args)} positional arguments")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 160
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- function definition ---'
rg -n -A35 -B8 'def get_dataset_dataloader' modelopt/torch/utils/dataset_utils.py
printf '%s\n' '--- all repository call sites ---'
rg -n -C3 'get_dataset_dataloader\s*\(' --glob '*.py' .
printf '%s\n' '--- public exports and module context ---'
rg -n -C3 '__all__|get_dataset_dataloader' modelopt/torch/utils/dataset_utils.py modelopt/torch/utils --glob '*.py' | head -240
printf '%s\n' '--- changed-file summary ---'
git diff --stat -- modelopt/torch/utils/dataset_utils.pyRepository: NVIDIA/Model-Optimizer
Length of output: 36217
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("modelopt/torch/utils/dataset_utils.py")
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "get_dataset_dataloader":
positional = list(node.args.posonlyargs) + list(node.args.args)
print("current positional parameters:", [arg.arg for arg in positional])
print("current keyword-only parameters:", [arg.arg for arg in node.args.kwonlyargs])
print("distributed positional index:", next(i for i, arg in enumerate(positional) if arg.arg == "distributed"))
print("sampler_kwargs positional index:", next(i for i, arg in enumerate(positional) if arg.arg == "sampler_kwargs"))
break
print("repository calls with positional arguments:")
for candidate in Path(".").rglob("*.py"):
try:
candidate_tree = ast.parse(candidate.read_text(encoding="utf-8"), filename=str(candidate))
except (OSError, SyntaxError):
continue
for node in ast.walk(candidate_tree):
if not isinstance(node, ast.Call):
continue
func = node.func
name = func.id if isinstance(func, ast.Name) else func.attr if isinstance(func, ast.Attribute) else None
if name == "get_dataset_dataloader" and node.args:
print(f"{candidate}:{node.lineno}: {len(node.args)} positional arguments")
PY
printf '%s\n' '--- parent signature ---'
git show HEAD^:modelopt/torch/utils/dataset_utils.py 2>/dev/null \
| sed -n '/^def get_dataset_dataloader(/,/^ \) -> DataLoader:/p' \
|| true
printf '%s\n' '--- current status and recent commit summary ---'
git status --short
git log -3 --oneline -- modelopt/torch/utils/dataset_utils.pyRepository: NVIDIA/Model-Optimizer
Length of output: 917
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- parent function signature ---'
git show HEAD^:modelopt/torch/utils/dataset_utils.py | sed -n '755,775p'
printf '%s\n' '--- parent declaration location ---'
git grep -n -A18 -B2 'def get_dataset_dataloader' HEAD^ -- modelopt/torch/utils/dataset_utils.py || true
printf '%s\n' '--- commit diff summary for the declaration ---'
git diff --unified=12 HEAD^ HEAD -- modelopt/torch/utils/dataset_utils.py | sed -n '/get_dataset_dataloader/,+55p'Repository: NVIDIA/Model-Optimizer
Length of output: 5894
Preserve the existing positional argument order.
get_dataset_dataloader is a public function. Move random_offset after the existing parameters and make it keyword-only to avoid rebinding positional distributed and sampler_kwargs arguments.
Proposed signature change
pack: bool = False,
- random_offset: bool = False,
distributed: bool = False,
sampler_kwargs: dict | None = None,
+ *,
+ random_offset: bool = False,🤖 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 public
get_dataset_dataloader signature so random_offset follows all existing
parameters and is keyword-only, preserving the current positional order and
preventing distributed and sampler_kwargs from being rebound for existing
callers.
|
/claude review |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
♻️ Duplicate comments (1)
modelopt/torch/utils/dataset_utils.py (1)
765-765:⚠️ Potential issue | 🟠 MajorPreserve the existing positional argument order.
random_offsetnow appears beforedistributedandsampler_kwargs. Existing positional callers can bind values to the wrong parameters. Moverandom_offsetafter the existing parameters and make it keyword-only. This repeats the unresolved compatibility issue from the previous review.Proposed signature
pack: bool = False, - random_offset: bool = False, distributed: bool = False, sampler_kwargs: dict | None = None, + *, + random_offset: bool = False,🤖 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 affected function signature containing random_offset so distributed and sampler_kwargs retain their existing positional order; move random_offset after those parameters and make it keyword-only, preserving compatibility for existing positional callers.
🧹 Nitpick comments (1)
modelopt/torch/utils/dataset_utils.py (1)
765-765: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument
random_offsetin the public API.The
Argssection does not describe the new parameter. Document that it applies topack=True, defaults toFalse, and shifts packed row boundaries.Suggested docstring entry
+ random_offset: If True and ``pack=True``, randomly shifts the packed row + boundaries by up to one sequence length. Defaults to False.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 public API docstring for the function exposing random_offset to add an Args entry describing that it applies when pack=True, defaults to False, and shifts packed row boundaries; keep the parameter declaration and existing behavior unchanged.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 `@modelopt/torch/utils/dataset_utils.py`:
- Around line 715-718: Update the dataset-building flow around the random_offset
handling to generate a single offset when distributed mode is enabled, broadcast
it to all ranks before packing, and apply that shared value to token_stream on
every rank. Preserve the existing local random offset behavior for
non-distributed execution and retain the current max_off bounds.
---
Duplicate comments:
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the affected function signature containing random_offset so
distributed and sampler_kwargs retain their existing positional order; move
random_offset after those parameters and make it keyword-only, preserving
compatibility for existing positional callers.
---
Nitpick comments:
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the public API docstring for the function exposing
random_offset to add an Args entry describing that it applies when pack=True,
defaults to False, and shifts packed row boundaries; keep the parameter
declaration and existing behavior unchanged.
🪄 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: 7896fa12-797e-4dca-b790-458901b3bb22
📒 Files selected for processing (7)
examples/megatron_bridge/distill.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/quantize.pymodelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/megatron_calibration.py
🚧 Files skipped from review as they are similar to previous changes (6)
- modelopt/torch/utils/plugins/megatron_calibration.py
- modelopt/torch/distill/plugins/megatron.py
- modelopt/torch/quantization/plugins/megatron.py
- examples/megatron_bridge/distill.py
- examples/megatron_bridge/export_quantized_megatron_to_hf.py
- examples/megatron_bridge/quantize.py
| 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):] |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 4 'random\.seed|torch\.manual_seed|set_random_seed|DistributedSampler|get_megatron_calibration_dataloader|random_offset' .Repository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- dataset_utils structure ---'
ast-grep outline modelopt/torch/utils/dataset_utils.py | sed -n '1,180p'
printf '%s\n' '--- packing and dataloader implementation ---'
sed -n '660,890p' modelopt/torch/utils/dataset_utils.py
printf '%s\n' '--- Megatron calibration implementation ---'
sed -n '1,155p' modelopt/torch/utils/plugins/megatron_calibration.py
printf '%s\n' '--- targeted callers and distributed initialization ---'
rg -n -C 3 \
'get_megatron_calibration_(dataloader|forward_loop)|get_dataset_dataloader\(|init_process_group|initialize_model_parallel|set_random_seed|random\.seed|torch\.manual_seed|DistributedSampler' \
modelopt examples tests \
-g '*.py' \
| rg 'dataset_utils|megatron_calibration|megatron|calibration|initialize_model_parallel|set_random_seed|random\.seed|torch\.manual_seed|DistributedSampler' \
| sed -n '1,260p'Repository: NVIDIA/Model-Optimizer
Length of output: 41330
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- model-parallel initialization and seed handling ---'
sed -n '60,120p' modelopt/torch/utils/plugins/mbridge.py
rg -n -C 8 \
'initialize_model_parallel\(|provider\.initialize_model_parallel|set_random_seed\(|random\.seed\(' \
examples/megatron_bridge modelopt/torch/utils/plugins tests/_test_utils/torch/megatron \
-g '*.py' \
| sed -n '1,320p'
printf '%s\n' '--- calibration example setup ---'
sed -n '320,400p' examples/megatron_bridge/quantize.py
sed -n '400,465p' examples/megatron_bridge/prune_minitron.py
printf '%s\n' '--- distributed helper seed and rank behavior ---'
rg -n -C 5 \
'def .*seed|seed\s*=|rank\(|get_data_parallel_rank|get_world_size|init_process_group' \
modelopt/torch/utils/distributed.py \
modelopt/torch/utils/plugins/mbridge.py \
-g '*.py' \
| sed -n '1,280p'
printf '%s\n' '--- relevant tests and random-offset coverage ---'
rg -n -C 6 \
'random_offset|pack=True|DistributedSampler|calibration_dataloader' \
tests/unit/torch/utils tests \
-g '*.py' \
| sed -n '1,320p'Repository: NVIDIA/Model-Optimizer
Length of output: 27383
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- Megatron example initialization order ---'
rg -n -C 12 \
'load_mbridge_model_from_hf|init_model_parallel|initialize_model_parallel|forward_loop\s*=|get_megatron_calibration_forward_loop' \
examples/megatron_bridge/quantize.py \
examples/megatron_bridge/prune_minitron.py \
| sed -n '1,360p'
printf '%s\n' '--- all direct callers of the calibration APIs ---'
rg -n -C 8 \
'get_megatron_calibration_dataloader|get_megatron_calibration_forward_loop' \
. -g '*.py' \
| sed -n '1,360p'
printf '%s\n' '--- read-only ordering and RNG behavior probe ---'
python3 - <<'PY'
from pathlib import Path
import random
dataset = Path("modelopt/torch/utils/dataset_utils.py").read_text()
pack = dataset.index("def _pack_documents_into_rows")
sampler = dataset.index("DistributedSampler(dataset")
offset = dataset.index("random.randint(0, max_off)", pack)
assert offset < sampler, "random offset is not selected before sampler construction"
print("random offset selection precedes DistributedSampler construction: yes")
def rank_offsets(seeds):
return [random.Random(seed).randint(0, 8) for seed in seeds]
same_seed = rank_offsets([0, 0, 0, 0])
different_seeds = rank_offsets([0, 1, 2, 3])
print("same Python seed offsets:", same_seed)
print("different Python seed offsets:", different_seeds)
print("different Python RNG states can produce different offsets:",
len(set(different_seeds)) > 1)
PYRepository: NVIDIA/Model-Optimizer
Length of output: 33876
🌐 Web query:
Megatron Bridge provider.initialize_model_parallel(seed=0) Python random seed set_random_seed implementation
💡 Result:
In Megatron-LM, random seed management is handled primarily by the _set_random_seed function (typically defined in megatron/training/initialize.py) [1][2][3]. While initialize_model_parallel sets up the necessary parallel process groups (TP, PP, DP, CP, EP) [4][5][6], it does not automatically configure the random seeds; this is a separate initialization step that must be called afterward to ensure reproducibility across different parallel ranks [2][3]. The _set_random_seed implementation performs the following steps: 1. Seed Adjustment: It modifies the base seed provided by the user (args.seed) to ensure that different parallel ranks receive unique seeds, preventing identical dropout patterns or data sampling across ranks [1][3]. - Pipeline Parallel (PP) ranks: Seed = base_seed + (100 * pp_rank) [1][3]. - Data Parallel (DP) ranks (if enabled): Seed = seed + (10 * dp_rank) [1][3]. 2. Standard RNG Seeding: It sets the seeds for Python's native random module, numpy.random, and PyTorch's CPU/CUDA manual seeds using the adjusted seed value [1][2][3]. 3. Model Parallel RNG Tracking: It calls tensor_parallel.model_parallel_cuda_manual_seed(seed,...), which replaces standard torch.cuda.manual_seed in Megatron-LM environments [7][8]. This function initializes the CudaRNGStatesTracker, which maintains separate RNG states for different parallel regions to ensure correct, reproducible behavior [7][2][8]: - Default State (Data Parallel): Uses the base seed (or DP-adjusted seed), shared among model-parallel GPUs but unique across different model-parallel groups [7][8]. - Tensor-Parallel (TP) State: Uses an offset seed (seed + 2718 + tp_rank), unique to each GPU within a TP group to ensure distinct dropout in TP regions [7][8]. - Expert-Parallel (EP) State (for MoE): Uses a distinct offset calculation involving EP and TP ranks [7][8]. Crucially, users should not call torch.cuda.manual_seed directly after _set_random_seed, as the custom RNG tracker is designed to manage these states specifically for Megatron-LM's parallel architecture [7][8][9].
Citations:
- 1: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/initialize.py
- 2: https://deepwiki.com/NVIDIA/Megatron-LM/5.2-initialization-and-setup
- 3: https://github.com/NVIDIA/Megatron-LM/blob/23ba3570/megatron/training/initialize.py
- 4: https://github.com/NVIDIA/Megatron-LM/blob/40b1fd37/megatron/core/parallel_state.py
- 5: https://docs.nvidia.com/megatron-core/developer-guide/0.15.0/apidocs/core/core.parallel_state.html
- 6: https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.parallel_state.html.md
- 7: https://github.com/NVIDIA/Megatron-LM/blob/dddecd19/megatron/core/tensor_parallel/random.py
- 8: https://docs.nvidia.com/megatron-core/developer-guide/0.17.1/apidocs/core/core.tensor_parallel.random.html
- 9: [BUG] Mismatch Between Docstring and Behavior in core.tensor_parallel.random.model_parallel_cuda_manual_seed Megatron-LM#858
🌐 Web query:
site:github.com/NVIDIA/Megatron-Bridge "def initialize_model_parallel" seed
💡 Result:
The relevant Megatron Bridge API is:
initialize_model_parallel(
seed: int | None = None,
seed_kwargs: dict | None = None,
**model_parallel_kwargs,
) -> Noneseed sets the model-parallel RNG seed, while seed_kwargs passes additional arguments to model_parallel_cuda_manual_seed. [1]
[1] (docs.nvidia.com)
Citations:
🌐 Web query:
Megatron Bridge ModelProvider initialize_model_parallel source model_parallel_cuda_manual_seed seed_kwargs
💡 Result:
In the Megatron-LM ecosystem, the components you referenced relate to distributed state initialization and random seed management. initialize_model_parallel This function, located in megatron/core/parallel_state.py, is the primary entry point for establishing distributed training dimensions [1][2][3]. It creates the necessary process groups for tensor parallelism (TP), pipeline parallelism (PP), data parallelism (DP), context parallelism (CP), and expert parallelism (EP) [2]. It is typically called during the initialization phase of a Megatron training job, often via higher-level helpers like initialize_megatron in megatron/training/initialize.py [4][3]. model_parallel_cuda_manual_seed Located in megatron/core/tensor_parallel/random.py, this function acts as a replacement for torch.cuda.manual_seed in a model-parallel environment [5][6]. It must be called after initialize_model_parallel [5][6]. It ensures that random number generator (RNG) states are correctly tracked and synchronized across different parallelism dimensions (e.g., separate seeds for data-parallel vs. tensor-parallel regions to ensure correct dropout behavior) [5][6]. Bridge ModelProvider The "Bridge" refers to the Megatron Bridge library (often used for converting Hugging Face models) [7][8]. The ModelProvider pattern is used to configure and instantiate models within this framework [7][9]. - Provider Workflow: You obtain a provider instance (e.g., via bridge.to_megatron_provider), configure parallelism attributes (such as tensor_model_parallel_size and pipeline_model_parallel_size) on the provider instance, call finalize to validate these settings, and finally invoke provide_distributed_model to create the model [7][10][8]. - Seed Kwargs: While initialize_model_parallel and model_parallel_cuda_manual_seed handle low-level state and seeds, the Bridge's provide_distributed_model method manages the integration of these configurations, ensuring the model is correctly initialized within the distributed environment defined by the provider [9]. In summary, the recommended workflow is to initialize the parallel environment, configure your model provider, finalize it, and then instantiate the model using that provider [7][8]. Seeding should follow the standard Megatron-LM requirement: call initialize_model_parallel, followed by model_parallel_cuda_manual_seed, before model construction [5][11][12].
Citations:
- 1: https://github.com/NVIDIA/Megatron-LM/blob/40b1fd37/megatron/core/parallel_state.py
- 2: https://deepwiki.com/NVIDIA/Megatron-LM/4.1-parallel-state-management
- 3: https://deepwiki.com/NVIDIA/Megatron-LM/4-parallelism-strategies
- 4: https://github.com/NVIDIA/Megatron-LM/blob/3e6e32b5/megatron/training/initialize.py
- 5: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/random.py
- 6: https://github.com/NVIDIA/Megatron-LM/blob/dddecd19/megatron/core/tensor_parallel/random.py
- 7: https://docs.nvidia.com/nemo/megatron-bridge/latest/bridge-guide.html
- 8: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/README.md
- 9: https://docs.nvidia.com/nemo/megatron-bridge/0.3.0/apidocs/bridge/bridge.models.model_provider.html
- 10: https://docs.nvidia.com/nemo/megatron-bridge/0.4.1/bridge-guide.html
- 11: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/QuickStart.md
- 12: https://github.com/NVIDIA/Megatron-LM/blob/main/examples/run_simple_mcore_train_loop.py
Use one shared random offset across distributed ranks.
With distributed=True and random_offset=True, random.randint() uses each process's Python RNG state before DistributedSampler shards the dataset. Megatron-Bridge's model-parallel seed does not synchronize this Python RNG state, so ranks can shard different packed streams. Generate the offset once and broadcast it before packing.
🧰 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
dataset-building flow around the random_offset handling to generate a single
offset when distributed mode is enabled, broadcast it to all ranks before
packing, and apply that shared value to token_stream on every rank. Preserve the
existing local random offset behavior for non-distributed execution and retain
the current max_off bounds.
There was a problem hiding this comment.
♻️ Duplicate comments (2)
modelopt/torch/utils/dataset_utils.py (2)
765-765: 🎯 Functional Correctness | 🟠 MajorPreserve the existing positional argument order.
random_offsetis inserted beforedistributedandsampler_kwargsat Line 765. Existing positional callers now bind those values to the wrong parameters. Keep the existing parameters in place and makerandom_offsetkeyword-only.This is the previously reported
get_dataset_dataloadercompatibility issue, and it remains unresolved.Proposed signature change
apply_chat_template: bool = False, pack: bool = False, - random_offset: bool = False, distributed: bool = False, sampler_kwargs: dict | None = None, + *, + random_offset: bool = False, ) -> DataLoader:Verify positional call sites
#!/bin/bash set -euo pipefail python3 - <<'PY' import ast from pathlib import Path for path in Path(".").rglob("*.py"): try: tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) except (OSError, SyntaxError): continue for node in ast.walk(tree): if not isinstance(node, ast.Call): continue if isinstance(node.func, ast.Name): name = node.func.id elif isinstance(node.func, ast.Attribute): name = node.func.attr else: name = None if name == "get_dataset_dataloader" and len(node.args) >= 10: print(f"{path}:{node.lineno}: {len(node.args)} positional arguments") PY🤖 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 get_dataset_dataloader signature so existing positional parameters, including distributed and sampler_kwargs, retain their original order; make random_offset keyword-only by placing it after the positional parameter section. Preserve compatibility for existing positional callers and ensure random_offset is supplied explicitly by keyword.
715-718: 🗄️ Data Integrity & Integration | 🟡 MinorSynchronize the random offset across distributed ranks.
If
distributed=Trueandrandom_offset=True, each process executesrandom.randint()independently at Line 718. Each rank can therefore create different packed rows beforeDistributedSamplershards the dataset. Broadcast one offset, or derive it from a shared seed, before packing.This is the previously reported distributed-offset issue, and it remains unresolved.
Verify distributed callers and RNG setup
#!/bin/bash set -euo pipefail rg -n -C 8 \ 'get_megatron_calibration_dataloader|get_dataset_dataloader|random_offset|DistributedSampler|random\.seed|set_random_seed' \ modelopt/torch -g '*.py'🤖 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, Synchronize the random offset in the dataset-building flow around the random_offset branch before token_stream is sliced. When distributed is enabled, select one offset consistently across all ranks using the existing distributed communication or shared-seed utilities, then apply that same value on every process; preserve the current local random behavior when distributed is disabled.
🧹 Nitpick comments (1)
modelopt/torch/utils/dataset_utils.py (1)
765-765: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument
random_offsetin the public API.The
get_dataset_dataloaderdocstring omits the new parameter. Add anArgsentry and state that the flag is ignored whenpack=False.Proposed documentation addition
pack: If True, use global-stream document packing (Megatron-LM pretraining style): all raw samples are concatenated into one EOS-separated token stream and sliced into uniform-length rows. Rows can (and usually do) start mid-document — this matches the distribution Megatron's blended ``.bin``/``.idx`` pretraining uses, so the trained model has seen this pattern extensively. Non-final rows are fully real tokens (no pad); only the trailing partial row (when the stream runs out before reaching ``num_samples`` rows) is padded. Default ``False`` for backwards-compatibility with the prior one-doc-per-row tokenize-and-pad behavior; calibration callers should pass ``True``. + random_offset: If True and ``pack=True``, shift the packed token stream by a + random prefix within one sequence-length window. Ignored when ``pack=False``. distributed: If True, shard the dataset across ranks with a ``DistributedSampler``As per coding guidelines, public and higher-level APIs must be documented with docstrings.
🤖 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 get_dataset_dataloader docstring to add an Args entry for random_offset, describing its behavior and explicitly stating that the flag is ignored when pack=False.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.
Duplicate comments:
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the get_dataset_dataloader signature so existing positional
parameters, including distributed and sampler_kwargs, retain their original
order; make random_offset keyword-only by placing it after the positional
parameter section. Preserve compatibility for existing positional callers and
ensure random_offset is supplied explicitly by keyword.
- Around line 715-718: Synchronize the random offset in the dataset-building
flow around the random_offset branch before token_stream is sliced. When
distributed is enabled, select one offset consistently across all ranks using
the existing distributed communication or shared-seed utilities, then apply that
same value on every process; preserve the current local random behavior when
distributed is disabled.
---
Nitpick comments:
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the get_dataset_dataloader docstring to add an Args entry for
random_offset, describing its behavior and explicitly stating that the flag is
ignored when pack=False.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 15015317-5ae1-4e56-856e-0143268081af
📒 Files selected for processing (7)
examples/megatron_bridge/distill.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/quantize.pymodelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/megatron_calibration.py
🚧 Files skipped from review as they are similar to previous changes (6)
- modelopt/torch/distill/plugins/megatron.py
- modelopt/torch/utils/plugins/megatron_calibration.py
- modelopt/torch/quantization/plugins/megatron.py
- examples/megatron_bridge/export_quantized_megatron_to_hf.py
- examples/megatron_bridge/quantize.py
- examples/megatron_bridge/distill.py
f40d6ed to
11ab426
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if not isinstance(module, TensorQuantizer) or isinstance( | ||
| module, StaticBlockScaleQuantizer | ||
| ): | ||
| continue | ||
| block_sizes = getattr(module, "_block_sizes", None) | ||
| is_nvfp4 = getattr(module, "_num_bits", None) == (2, 1) and ( | ||
| isinstance(block_sizes, dict) and block_sizes.get("scale_bits") == (4, 3) | ||
| ) | ||
| if not is_nvfp4: | ||
| continue | ||
| amax = getattr(module, "_amax", None) | ||
| 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.") | ||
| n_skipped += 1 | ||
| continue | ||
| StaticBlockScaleQuantizer.from_tensor_quantizer( | ||
| module, global_amax=amax.detach().float().abs().max() | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Algorithm] This ad-hoc promotion predicate is missing the is_static_block_quant condition, so it will promote dynamic NVFP4 quantizers to static block scaling.
The issue. The inline check is _num_bits == (2, 1) and block_sizes["scale_bits"] == (4, 3). The canonical predicate is TensorQuantizer.is_nvfp4_static (modelopt/torch/quantization/nn/modules/tensor_quantizer.py:573), which additionally requires is_static_block_quant — i.e. block_sizes.get("type") != "dynamic" and _fake_quant. A dynamic NVFP4 weight quantizer ({"type": "dynamic", "scale_bits": (4, 3)}, _num_bits == (2, 1)) satisfies the inline predicate but not the real one.
Why it matters. Dynamic NVFP4 weight quantizers do carry _amax: max_calibrate deliberately runs weight calibration on the weight tensor directly "so every weight quantizer gets _amax" (modelopt/torch/quantization/model_calib.py:339-341). So the amax is None skip does not protect them — they pass the guard and get converted via from_tensor_quantizer, permanently switching the layer from per-block dynamic scaling to static block scaling frozen at a stale calibration amax. This is silent numerical corruption of every MoE/linear weight in the student.
This is reachable on the default path: --student_nongrouped_experts defaults off, which the help text describes as the mode for "dynamic NVFP4 / grouped-expert checkpoints (e.g. Nemotron-3-Nano)". The measured run cited in the comment (already promoted 460, converted 1) was a static-NVFP4 recipe, so it would not have surfaced this.
Secondly, global_amax=amax.detach().float().abs().max() bypasses shared-state tying. promote_static_block_weight_quantizers (modelopt/torch/quantization/utils/core_utils.py:1016-1035) checks whether the quantizer belongs to a SharedWeightGlobalAmaxState group and ties it to the group's canonical buffer, raising if the group was never populated — precisely so fusible siblings (q/k/v, gate/up) share one FP8 grid. Promoting here gives each sibling an independent global_amax, reintroducing the inconsistency _check_grouped_weight_global_amax_synced exists to catch. Today only output_layer (no siblings) converts, but the loop is generic and will mis-promote siblings whenever the "everything else is already promoted" assumption breaks (older checkpoint, different restore path).
Suggested fix: delegate to the existing helper rather than reimplementing it — it already handles the format check, shared-state tying, and the reduce_amax global:
if not getattr(self, "_modelopt_nvfp4_promoted", False):
from modelopt.torch.quantization.utils.core_utils import (
promote_static_block_weight_quantizers,
)
n_promoted = promote_static_block_weight_quantizers(self)
if n_promoted:
logger.info(f"Promoted {n_promoted} static-block weight quantizer(s).")
self._modelopt_nvfp4_promoted = TrueIf the inline loop must stay, gate it on getattr(module, "is_nvfp4_static", False) instead of the hand-rolled tuple comparison.
| _wq = getattr(self, "weight_quantizer", None) | ||
| if _wq is not None and getattr(_wq, "is_enabled", False): | ||
| _block_sizes = getattr(_wq, "_block_sizes", None) or {} | ||
| _block = _block_sizes.get(-1) or _block_sizes.get(1) | ||
| if _block and self.weight.numel() % int(_block) == 0: | ||
| if getattr(_wq, "_amax", None) is None: | ||
| _wq.amax = torch.zeros( | ||
| self.weight.numel() // int(_block), | ||
| 1, | ||
| dtype=torch.float32, | ||
| device=self.weight.device, | ||
| ) | ||
| # register_buffer directly: the ``global_amax`` property lives on | ||
| # StaticBlockScaleQuantizer, and on restore this is still a plain | ||
| # TensorQuantizer (promotion happens later), so the setter is unavailable. | ||
| if getattr(_wq, "_global_amax", None) is None: | ||
| _wq.register_buffer( | ||
| "_global_amax", | ||
| torch.zeros((), dtype=torch.float32, device=self.weight.device), | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Algorithm] The materialized _amax shape is wrong for a TP-sharded output_layer, and the divisibility guard checks the wrong quantity.
1. numel() is the local shard, but the block dim is the input dim. _amax is allocated as [self.weight.numel() // block, 1], and _process_quantizer_amax (line 374-380) later reshapes it with v.view(self.weight.shape[0], -1). For that view to yield the [out_features, blocks_per_row] layout the checkpoint expects, the flat length must be exactly weight.shape[0] * (weight.shape[1] // block). Using numel() // block only coincides with that when shape[1] % block == 0.
output_layer is a ColumnParallelLinear, so with TP>1 weight.shape is [vocab/TP, hidden] — shape[1] is the full hidden size, unsharded, so the two agree there. But the guard permits cases where they don't: it tests weight.numel() % block == 0, which is satisfied whenever the product is divisible even if shape[1] is not. E.g. shape = [128, 24] with block = 16: numel() = 3072, divisible by 16, so a [192, 1] buffer is allocated; but view(128, -1) needs 128 | 192, which fails, and _process_quantizer_amax raises a shape error at save time — after the load plan is already built.
2. _block_sizes.get(-1) or _block_sizes.get(1) — -1 and 1 are distinct axes in general (-1 is the last/input axis for a 2D weight, 1 is the same axis only for 2D). Falling through from -1 to 1 is fine for 2D, but combined with (1) the resulting _block is used against numel() rather than against the axis it actually blocks.
Suggested fix — derive the length from the input dim and validate that dim:
_wq = getattr(self, "weight_quantizer", None)
if _wq is not None and getattr(_wq, "is_enabled", False):
_block_sizes = getattr(_wq, "_block_sizes", None) or {}
_block = _block_sizes.get(-1) or _block_sizes.get(1)
_out, _in = self.weight.shape[0], self.weight.shape[1]
if not _block or _in % int(_block) != 0:
warn_rank_0(
f"{prefix}: skipping weight-quantizer buffer materialization "
f"(block_sizes={_block_sizes}, weight shape={tuple(self.weight.shape)}); "
"output_layer quantizer scales may not restore."
)
else:
if getattr(_wq, "_amax", None) is None:
_wq.amax = torch.zeros(
_out * (_in // int(_block)), 1,
dtype=torch.float32, device=self.weight.device,
)
...Note this also silently no-ops today when the guard fails (no else), which reproduces exactly the silent-BF16 failure mode this PR is fixing — a warning is important here.
3. is_enabled on a pre-replacement module. Per the sibling change at line 305-320, sharded_state_dict can also run when weight_quantizer does not exist yet. getattr(_wq, "is_enabled", False) returns False then, so nothing is materialized and the load silently skips the keys again. Worth confirming sharded_state_dict is only ever reached post-replacement; if not, this branch needs the same untied fallback the callback registration now uses.
| # Wrap into DistillationProvider | ||
| kd_config = ModelOptDistillConfig( | ||
| skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale | ||
| ) | ||
|
|
||
| # VLM detection convention: HF VLM configs expose a ``vision_config``, and Megatron-Bridge nests | ||
| # the text model under the ``language_model`` submodule (used as ``distill_submodule`` below). If a | ||
| # future model breaks either convention, the ``getattr(model, "language_model")`` in the provider | ||
| # will error loudly rather than silently distilling the wrong module. | ||
| is_vlm = hasattr( | ||
| AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), | ||
| "vision_config", | ||
| ) | ||
|
|
||
| if is_vlm: | ||
| warn_rank_0( | ||
| "VLM detected: distilling model.language_model only (vision tower / projector untouched). " | ||
| "To export megatron non-quantized checkpoint, use export_distilled_megatron_to_hf.py" | ||
| ) | ||
| distill_provider = convert_to_distillation_provider( | ||
| student_provider, | ||
| teacher_provider, | ||
| kd_config, | ||
| distill_submodule="language_model" if is_vlm else None, | ||
| student_provider, teacher_provider, kd_config | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Compatibility] This PR removes VLM distillation support that exists at its own merge base, along with several other unrelated regressions. These look like an accidental revert rather than intentional changes — none is mentioned in the PR description.
I verified against the PR's base blob (dfd404fab30, the distill.py this diff is computed against), which contains all of the following:
1. VLM distillation is gone. Base has is_vlm detection via AutoConfig's vision_config, passes distill_submodule="language_model" to convert_to_distillation_provider, and takes a dedicated save_vlm_to_hf export branch that re-exports the distilled language model back into the full VLM. The diff deletes all of it (base lines 258-272, 380-388). A VLM student now distills the whole model root instead of language_model, and exports through the LLM-only path. The helpers are still shipped and still document this contract — examples/megatron_bridge/export_distilled_megatron_to_hf.py:26 states "these two helpers are also reused by distill.py for its VLM/LLM export paths" — so that comment is now stale and save_vlm_to_hf has no caller in distill.py.
2. --checkpoint_keep_last is removed and hardcoded. Base exposes the flag (plus a >= -1 validation) and passes most_recent_k=args.checkpoint_keep_last. Line 491 now hardcodes most_recent_k=2 with the comment each ~413GB here; fs1 near quota — a site-specific detail from the author's cluster baked into a shared example, and a removed CLI flag that breaks existing scripts.
3. async_save=True → False (line 493), justified by async writer repeatedly corrupted iter-400 ckpt. That's a real symptom worth reporting, but flipping the default for every user of this example on the basis of one run is a significant throughput regression. If the async writer is genuinely broken, that deserves its own issue; otherwise make it a flag.
4. --student_hf_model becomes mandatory rather than defaulting. Base: if args.student_hf_model is None: args.student_hf_model = args.student_hf_path, documented as "Defaults to --student_hf_path, which is correct for homogeneous students". Line 295-296 now raises when --hf_export_path is set without it, so every previously-working homogeneous-student invocation fails. The help text also loses the Puzzletron/NAS-template explanation.
5. --validate_only / eval_iters validation present in base (if args.validate_only and args.eval_iters == 0) is dropped.
Suggested fix: rebase onto current main and re-apply only the intended additions (--student_nongrouped_experts, --sft, the checkpoint-init workaround, MoE aux-loss config). The branch is already flagged as having a merge conflict, which is consistent with these deletions being stale-base artifacts.
| continue | ||
| if getattr(module, "_amax", None) is None: | ||
| uncalibrated.append((name, "_amax")) | ||
| elif block_sizes.get("type") == "static" and getattr(module, "_global_amax", None) is None: |
There was a problem hiding this comment.
[IMPORTANT Export] The export guard never checks is_enabled, so it raises on disabled quantizers — turning an intentionally-excluded layer into a hard export failure.
The comment above states the guard targets "an enabled NVFP4 weight quantizer", and the error message says "enabled NVFP4 weight quantizer(s) are missing calibrated scales" — but this loop only filters on type and format. A recipe that deliberately excludes a layer ("*lm_head*": {"enable": False}, or any of the 72 exclusions this PR's own results table mentions) leaves a TensorQuantizer in place with NVFP4 _num_bits/_block_sizes, _disabled = True, and no _amax — because it was never calibrated, correctly. That module lands in uncalibrated and export raises.
The fallback path then calls .disable() on modules that are already disabled, which is itself the tell that disabled quantizers reach here.
Separately, the loop excludes StaticBlockScaleQuantizer instances (isinstance(module, StaticBlockScaleQuantizer) → continue), which is backwards for this guard's purpose. A promoted static-block quantizer is exactly the case that needs both _amax and _global_amax, and the NVFP4QTensor.quantize broadcast failure described in the comment above occurs on the promoted class. As written the guard only inspects unpromoted quantizers, so a promoted-but-half-restored output_layer — this PR's motivating bug — is skipped entirely.
Suggested fix: filter on is_enabled, and use the canonical is_nvfp4_static property (modelopt/torch/quantization/nn/modules/tensor_quantizer.py:573) instead of reimplementing the format check, so promoted quantizers are covered:
for name, module in unwrapped_model.named_modules():
if not isinstance(module, TensorQuantizer) or not module.is_enabled:
continue
if not getattr(module, "is_nvfp4_static", False):
continue
if getattr(module, "_amax", None) is None:
uncalibrated.append((name, "_amax"))
elif getattr(module, "_global_amax", None) is None:
uncalibrated.append((name, "_global_amax"))is_nvfp4_static already implies block_sizes.get("type") != "dynamic", so the separate block_sizes.get("type") == "static" test on this line becomes unnecessary. Worth noting that test is also currently too narrow: static block configs commonly omit "type" rather than setting it to "static", in which case a missing _global_amax is silently accepted — the exact failure the comment says it wants to catch.
| def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None: | ||
| """Whether ``output_layer`` weights are untied from the input embeddings, or None if unknown. | ||
|
|
||
| Megatron-Core models carry ``share_embeddings_and_output_weights`` (Megatron-Bridge sets it | ||
| from the HF config, Megatron-LM from ``--untie-embeddings-and-output-weights``), so reading | ||
| it off the model works under both frameworks. ``megatron.training.get_args()`` does not: | ||
| Bridge has no global args store, and defaulting to "tied" there silently drops the | ||
| ``output_layer`` weight-quantizer state from the sharded checkpoint. | ||
| """ | ||
| for _, module in model.named_modules(): | ||
| shared = getattr(module, "share_embeddings_and_output_weights", None) | ||
| if shared is not None: | ||
| return not bool(shared) | ||
| return None | ||
|
|
There was a problem hiding this comment.
[IMPORTANT Algorithm] _resolve_output_layer_untied returns the first module carrying share_embeddings_and_output_weights in named_modules() order, which is not necessarily the module that owns output_layer.
named_modules() yields the root first, so for a plain GPTModel/MambaModel this picks the right flag. But for the composite models this plugin explicitly handles it can pick the wrong one:
- VLM. The loop below (line 331-340) skips
vision_modelwhen setting config flags, so a VLM root is expected here. If the VLM root exposesshare_embeddings_and_output_weights(or a vision submodule is visited beforelanguage_model), the value read may not be the language model's. The resolved flag is then written to every non-visionMegatronModule's config, so one wrong read propagates to all of them. - PP>1 / multiple model chunks.
megatron_replace_quant_module_hookis called per model; with severalMegatronModulechildren the single root-level scan assumes they all agree.
Because a wrong True here makes sharded_state_dict take the quantized branch for a genuinely tied output_layer, this can advertise (and materialize, per the buffer block at line 433) quantizer keys for a weight shared with the input embedding — producing a checkpoint whose keys don't match what a correctly-resolved model would advertise. The failure is silent and, unlike the bug being fixed, it corrupts rather than omits.
Suggested fix: resolve from the module that actually owns output_layer, and only fall back to a scan if that fails:
def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None:
for _, module in model.named_modules():
if hasattr(module, "output_layer") and (
(shared := getattr(module, "share_embeddings_and_output_weights", None)) is not None
):
return not bool(shared)
return NoneSetting modelopt_output_layer_untied per-owning-module (rather than broadcasting one value to every config) would also make the PP / VLM cases correct by construction. A warn_rank_0 when candidate modules disagree would surface the remaining ambiguity rather than silently taking the first.
| defaults = {"mtp_num_layers": 0} | ||
| run_config = next(iter(sorted(pathlib.Path(megatron_path).glob("*/run_config.yaml"))), None) | ||
| if run_config is None: | ||
| run_config = pathlib.Path(megatron_path) / "run_config.yaml" |
There was a problem hiding this comment.
[IMPORTANT Compatibility] Hardcoding mtp_num_layers=0 and moe_grouped_gemm=False in the shared exporter breaks export of checkpoints that do have MTP layers or grouped experts.
Both comments describe one specific checkpoint ("QAD ckpt has MTP dropped", "QAD ckpt is non-grouped"), but this script is the general Megatron-Bridge → HF quantized exporter, not a QAD-only tool. load_mbridge_model_from_hf defaults moe_grouped_gemm=True (modelopt/torch/utils/plugins/mbridge.py:47), and the model structure built here must match the checkpoint being loaded by load_modelopt_megatron_checkpoint. Forcing both values means:
- A grouped-expert quantized checkpoint (the common PTQ output —
quantize.pyonly switches to non-grouped for static-block NVFP4) is now loaded into a non-grouped structure. Given this PR's own central finding — that the dist-checkpoint loader silently skips keys the model doesn't advertise — a structural mismatch here is liable to drop expert weights silently rather than error. - An MTP-enabled checkpoint loses its MTP layers on export.
Same concern for quantize.py:292,299, which hardcodes the identical pair on the calibration path. There mtp_num_layers=0 is defensible (MTP genuinely isn't supported during calibration), but moe_grouped_gemm=False is only needed for static-block NVFP4 recipes — forcing it for every recipe changes the layer spec for dynamic-NVFP4/FP8 users who previously calibrated with grouped GEMM, which is both a perf regression and a checkpoint-structure change.
Suggested fix: expose both as CLI flags with the current defaults preserved, so the QAD path opts in explicitly:
parser.add_argument("--mtp_num_layers", type=int, default=None,
help="Override provider mtp_num_layers (0 for QAD checkpoints with MTP dropped).")
parser.add_argument("--moe_grouped_gemm", action=argparse.BooleanOptionalAction, default=True,
help="Use grouped GEMM for MoE. Pass --no-moe-grouped-gemm for "
"static-block NVFP4 checkpoints built with non-grouped experts.")then pass through only when set, leaving the provider default otherwise. Note the mirrored --student_nongrouped_experts flag added in distill.py gets this right — it defaults off precisely to preserve committed grouped behavior. These two sites should follow the same principle.
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 5
♻️ Duplicate comments (1)
examples/megatron_bridge/export_quantized_megatron_to_hf.py (1)
194-206: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe NVFP4 calibration guard still skips
StaticBlockScaleQuantizer, the exact case it documents.Line 195 still reads
if not isinstance(module, TensorQuantizer) or isinstance(module, StaticBlockScaleQuantizer): continue. This is the same code flagged in a prior review round on this file (marked "Addressed in commit f40d6ed"), but the code under review has not changed: everyStaticBlockScaleQuantizerinstance is skipped before reaching the_global_amaxcheck at Line 205. Since_global_amaxis defined onStaticBlockScaleQuantizer, and the comment at Lines 186-189 states a missing_global_amaxotherwise fails later insideNVFP4QTensor.quantizewith a broadcast shape mismatch, this is precisely the failure mode the guard is meant to catch, and it currently cannot. Theelifbranch at Line 205 (block_sizes.get("type") == "static") is effectively dead for this class. Theis_enabledcheck requested previously (so disabled quantizers are not reported) is also still absent.Because this reproduces a previously-flagged defect that was reported as fixed, verify whether the fix was reverted or never merged into this branch.
🐛 Proposed fix
uncalibrated: list[tuple[str, str]] = [] for name, module in unwrapped_model.named_modules(): - if not isinstance(module, TensorQuantizer) or isinstance(module, StaticBlockScaleQuantizer): + if not isinstance(module, TensorQuantizer) or not module.is_enabled: continue block_sizes = getattr(module, "_block_sizes", None) is_nvfp4 = getattr(module, "_num_bits", None) == (2, 1) and ( isinstance(block_sizes, dict) and block_sizes.get("scale_bits") == (4, 3) ) if not is_nvfp4: continue if getattr(module, "_amax", None) is None: uncalibrated.append((name, "_amax")) - elif block_sizes.get("type") == "static" and getattr(module, "_global_amax", None) is None: + elif ( + isinstance(module, StaticBlockScaleQuantizer) + or block_sizes.get("type") == "static" + ) and getattr(module, "_global_amax", None) is None: uncalibrated.append((name, "_global_amax"))🤖 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 194 - 206, Update the NVFP4 calibration guard in the loop over unwrapped_model.named_modules so StaticBlockScaleQuantizer instances are not excluded before validation. Include the requested is_enabled check so disabled quantizers are skipped, then preserve the _amax and static _global_amax validation for enabled NVFP4 quantizers.
🧹 Nitpick comments (1)
modelopt/torch/utils/dataset_utils.py (1)
765-765: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument
random_offsetin the public API.At Line [765], the new parameter has no entry in
get_dataset_dataloader'sArgssection. Document that it applies only whenpack=Trueand describe the random-prefix behavior.As per coding guidelines, document public and higher-level APIs with docstrings.
🤖 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 get_dataset_dataloader docstring’s Args section to document random_offset, including that it is effective only when pack=True and controls adding a randomly sized prefix before packed samples.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/export_quantized_megatron_to_hf.py`:
- Around line 193-222: Synchronize the uncalibrated-quantizer condition across
ranks before any raise or later collective, using the existing has_extra_modules
all_reduce pattern and ReduceOp.MAX. Update the validation block around
uncalibrated and the RuntimeError so every rank raises consistently when any
rank finds missing scales; ranks without local findings should use an
abort-context message, while preserving the existing allow-environment warning
and local disable behavior.
- Around line 116-144: Update _provider_overrides_from_checkpoint to select the
latest iter_* checkpoint directory when megatron_path contains multiple
checkpoints, rather than choosing the first sorted run_config.yaml; retain the
root run_config.yaml fallback. After yaml.safe_load, validate that cfg is a dict
before calling cfg.get, and use defaults for list or scalar YAML roots.
In `@examples/megatron_bridge/quantize.py`:
- Around line 185-189: Update the argument handling in the quantization flow
around --calib_random_offset and get_megatron_vlm_calibration_forward_loop so
VLM calibration cannot silently ignore the option: either reject its combination
with VLM calibration via a clear validation error, or explicitly document and
enforce that the flag is text-only. Preserve the existing random-offset behavior
for text calibration.
- Around line 177-184: Update the quantization setup in the main flow of
quantize.py to resolve mtq_config before calling load_mbridge_model_from_hf().
When args.grouped_experts is enabled, validate that expert quantization uses
per-tensor scaling and reject per-block configurations before constructing the
model; leave compatible configurations and the non-grouped path unchanged.
- Around line 291-300: Update _hf_config_has_mtp to read MTP fields from both
attribute-backed objects and mapping-backed configs, including nested
text_config mappings, so existing MTP detection remains accurate before line 312
drops the heads. Add tests covering object-backed and mapping-backed top-level
and text_config configurations.
---
Duplicate comments:
In `@examples/megatron_bridge/export_quantized_megatron_to_hf.py`:
- Around line 194-206: Update the NVFP4 calibration guard in the loop over
unwrapped_model.named_modules so StaticBlockScaleQuantizer instances are not
excluded before validation. Include the requested is_enabled check so disabled
quantizers are skipped, then preserve the _amax and static _global_amax
validation for enabled NVFP4 quantizers.
---
Nitpick comments:
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the get_dataset_dataloader docstring’s Args section to
document random_offset, including that it is effective only when pack=True and
controls adding a randomly sized prefix before packed samples.
🪄 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: 97013c63-11ad-4614-9fdb-322db58c2e97
📒 Files selected for processing (7)
examples/megatron_bridge/distill.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/quantize.pymodelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/megatron_calibration.py
🚧 Files skipped from review as they are similar to previous changes (4)
- modelopt/torch/distill/plugins/megatron.py
- modelopt/torch/utils/plugins/megatron_calibration.py
- modelopt/torch/quantization/plugins/megatron.py
- examples/megatron_bridge/distill.py
| def _provider_overrides_from_checkpoint(megatron_path: str) -> dict: | ||
| """Read ``mtp_num_layers`` from the checkpoint so the exporter matches how it was saved. | ||
|
|
||
| Only ``mtp_num_layers`` is taken from here. ``moe_grouped_gemm`` is deliberately NOT derived: | ||
| for a ``MambaModelProvider`` the expert layout is set by ``mamba_stack_spec``, so a checkpoint | ||
| saved with non-grouped experts still records ``moe_grouped_gemm: true`` and trusting it would | ||
| build a mismatched model. | ||
| """ | ||
| defaults = {"mtp_num_layers": 0} | ||
| run_config = next(iter(sorted(pathlib.Path(megatron_path).glob("*/run_config.yaml"))), None) | ||
| if run_config is None: | ||
| run_config = pathlib.Path(megatron_path) / "run_config.yaml" | ||
| if not run_config.exists(): | ||
| print_rank_0(f"No run_config.yaml under {megatron_path}; using defaults {defaults}.") | ||
| return defaults | ||
| try: | ||
| cfg = yaml.safe_load(run_config.read_text()) or {} | ||
| except Exception as exc: | ||
| print_rank_0(f"Could not parse {run_config} ({exc}); using defaults {defaults}.") | ||
| return defaults | ||
| 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} | ||
| print_rank_0(f"Model shape from {run_config.name}: {resolved}") | ||
| return resolved |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect how run_config.yaml is written/read for Megatron checkpoints.
rg -n -C5 'run_config\.yaml' --type=pyRepository: NVIDIA/Model-Optimizer
Length of output: 160
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate files ---'
git ls-files | grep -E '(^|/)(export_quantized_megatron_to_hf\.py|.*run.*config.*\.py|.*megatron.*\.py)$' | head -200
printf '%s\n' '--- references ---'
rg -n -C4 'run_config|_provider_overrides_from_checkpoint' . --glob '*.py' --glob '*.yaml' --glob '*.yml' || true
printf '%s\n' '--- target function ---'
target=$(git ls-files | grep 'examples/megatron_bridge/export_quantized_megatron_to_hf.py$' | head -1)
if [ -n "$target" ]; then
sed -n '100,155p' "$target"
fiRepository: NVIDIA/Model-Optimizer
Length of output: 11868
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- expression behavior ---'
python3 - <<'PY'
cases = [None, [], ["x"], 1, "text", {}, {"model": {"mtp_num_layers": 2}}, {"mtp_num_layers": 3}]
for raw in cases:
cfg = raw or {}
try:
model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg
print(f"{raw!r} -> {model_cfg!r}")
except Exception as exc:
print(f"{raw!r} -> {type(exc).__name__}: {exc}")
PY
printf '%s\n' '--- checkpoint-path and exporter references ---'
rg -n -C5 'export_quantized_megatron_to_hf|megatron_path|checkpoints/|iter_[0-9].*run_config|run_config\.yaml' . \
--glob '*.py' --glob '*.md' --glob '*.yaml' --glob '*.yml' --glob '*.sh' || true
printf '%s\n' '--- nearby exporter tests ---'
test_file=$(git ls-files | grep 'tests/examples/megatron_bridge/test_quantize_export.py$' | head -1)
if [ -n "$test_file" ]; then
sed -n '1,260p' "$test_file"
fiRepository: NVIDIA/Model-Optimizer
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- quantized exporter interface ---'
file=examples/megatron_bridge/export_quantized_megatron_to_hf.py
sed -n '1,110p' "$file"
printf '%s\n' '--- checkpoint resolution helper ---'
file=$(git ls-files | grep 'modelopt.*/mbridge.py$' | head -1)
rg -n -C8 '_get_modelopt_checkpoint_path|latest_checkpointed_iteration|iter_' "$file"
printf '%s\n' '--- quantization output and save arguments ---'
rg -n -C5 'export_megatron_path|save.*checkpoint|checkpoint_dir|iter_|latest_checkpointed_iteration' \
examples/megatron_bridge/quantize.py tests/examples/megatron_bridge/test_quantize_export.py \
--glob '*.py' | head -300
printf '%s\n' '--- focused run_config references ---'
rg -l 'run_config\.yaml' . --glob '*.py' --glob '*.md' --glob '*.yaml' --glob '*.yml' |
while IFS= read -r f; do
printf '%s\n' "--- $f"
rg -n -C3 'run_config\.yaml' "$f"
doneRepository: NVIDIA/Model-Optimizer
Length of output: 4523
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- checkpoint resolution helper ---'
sed -n '115,155p' modelopt/torch/utils/plugins/mbridge.py
rg -n -C10 '_get_modelopt_checkpoint_path|latest_checkpointed_iteration|iter_' \
modelopt/torch/utils/plugins/mbridge.py
printf '%s\n' '--- quantization output and save arguments ---'
rg -n -C5 'export_megatron_path|save.*checkpoint|checkpoint_dir|iter_|latest_checkpointed_iteration' \
examples/megatron_bridge/quantize.py tests/examples/megatron_bridge/test_quantize_export.py \
--glob '*.py' | head -300
printf '%s\n' '--- all run_config references ---'
rg -l 'run_config\.yaml' . --glob '*.py' --glob '*.md' --glob '*.yaml' --glob '*.yml' |
while IFS= read -r f; do
printf '%s\n' "--- $f"
rg -n -C3 'run_config\.yaml' "$f"
doneRepository: NVIDIA/Model-Optimizer
Length of output: 12424
Guard non-dict YAML roots and select the latest checkpoint. Check isinstance(cfg, dict) before calling cfg.get; the proposed check runs too late for list or scalar roots. When megatron_path is a parent containing multiple iter_* directories, resolve the latest iteration instead of using sorted(...)[0], which can select an older 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 116
- 144, Update _provider_overrides_from_checkpoint to select the latest iter_*
checkpoint directory when megatron_path contains multiple checkpoints, rather
than choosing the first sorted run_config.yaml; retain the root run_config.yaml
fallback. After yaml.safe_load, validate that cfg is a dict before calling
cfg.get, and use defaults for list or scalar YAML roots.
| uncalibrated: list[tuple[str, str]] = [] | ||
| for name, module in unwrapped_model.named_modules(): | ||
| if not isinstance(module, TensorQuantizer) or isinstance(module, StaticBlockScaleQuantizer): | ||
| continue | ||
| block_sizes = getattr(module, "_block_sizes", None) | ||
| is_nvfp4 = getattr(module, "_num_bits", None) == (2, 1) and ( | ||
| isinstance(block_sizes, dict) and block_sizes.get("scale_bits") == (4, 3) | ||
| ) | ||
| if not is_nvfp4: | ||
| continue | ||
| if getattr(module, "_amax", None) is None: | ||
| uncalibrated.append((name, "_amax")) | ||
| elif block_sizes.get("type") == "static" and getattr(module, "_global_amax", None) is None: | ||
| uncalibrated.append((name, "_global_amax")) | ||
|
|
||
| if uncalibrated: | ||
| detail = ", ".join(f"{name}.{attr}" for name, attr in uncalibrated[:8]) | ||
| if len(uncalibrated) > 8: | ||
| detail += ", ..." | ||
| message = ( | ||
| f"{len(uncalibrated)} enabled NVFP4 weight quantizer(s) are missing calibrated scales " | ||
| f"after loading {args.megatron_path}: {detail}. These weights would be exported as " | ||
| "BF16 instead of NVFP4. Re-run PTQ with a ModelOpt that saves and restores this " | ||
| "quantizer state, or set MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1 to export them as BF16." | ||
| ) | ||
| if os.environ.get("MODELOPT_ALLOW_UNCALIBRATED_NVFP4") != "1": | ||
| raise RuntimeError(message) | ||
| print_rank_0(f"WARNING (MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1): {message}") | ||
| for name, _ in uncalibrated: | ||
| unwrapped_model.get_submodule(name).disable() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
RuntimeError on one rank can hang the rest of the job under pipeline parallelism.
uncalibrated is built from unwrapped_model.named_modules(), which is per-rank-local when pp_size > 1 (each pipeline stage holds a different subset of layers, and lm_head lives only on the last stage). If only the rank owning the affected layer(s) finds missing scales, only that rank executes raise RuntimeError(message) at Line 219. Every other rank has no way to learn about the failure and continues past this block to the torch.distributed.all_reduce call at Line 236, which then hangs waiting on the crashed rank. This script already implements the correct cross-rank pattern for exactly this situation a few lines below (has_extra_modules all_reduce with ReduceOp.MAX); mirror that pattern here so the decision to raise (or warn/disable) is made consistently on every rank before any rank exits or hangs on a later collective call. This scenario is directly relevant to this PR, which targets fixing lm_head quantization/tiedness handling.
has_uncalibrated = bool(uncalibrated)
if torch.distributed.is_initialized():
flag = torch.tensor(
[int(has_uncalibrated)], dtype=torch.int, device=torch.cuda.current_device()
)
torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MAX)
has_uncalibrated = bool(flag.item())
if uncalibrated or has_uncalibrated:
if not uncalibrated:
message = (
"A peer rank found enabled NVFP4 weight quantizer(s) missing calibrated scales "
f"after loading {args.megatron_path}. Aborting on all ranks to avoid a hang."
)
else:
... # existing local `detail`/`message` construction
if os.environ.get("MODELOPT_ALLOW_UNCALIBRATED_NVFP4") != "1":
raise RuntimeError(message)
print_rank_0(f"WARNING (MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1): {message}")
for name, _ in uncalibrated:
unwrapped_model.get_submodule(name).disable()
else:
print_rank_0("All enabled NVFP4 weight quantizers have calibrated scales.")🤖 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 193
- 222, Synchronize the uncalibrated-quantizer condition across ranks before any
raise or later collective, using the existing has_extra_modules all_reduce
pattern and ReduceOp.MAX. Update the validation block around uncalibrated and
the RuntimeError so every rank raises consistently when any rank finds missing
scales; ranks without local findings should use an abort-context message, while
preserving the existing allow-environment warning and local disable behavior.
| parser.add_argument( | ||
| "--grouped_experts", | ||
| action="store_true", | ||
| help="Build MoE experts grouped (GroupedMLP). Default is non-grouped (SequentialMLP), " | ||
| "which per-block NVFP4 requires because TEGroupedLinear can only represent per-tensor " | ||
| "scales. Set this for per-tensor recipes on MoE models, where grouped GEMM is faster. " | ||
| "The export must use the matching layout.", | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Expect: an explicit fail-fast validation for grouped experts and per-block recipes.
rg -n -C 6 \
'grouped_experts|moe_grouped_gemm|get_quant_config|NVFP4|block_sizes|scale_bits' \
. --glob '*.py' --glob '*.yaml'Repository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== quantize.py structure and relevant ranges =="
ast-grep outline examples/megatron_bridge/quantize.py
sed -n '280,380p' examples/megatron_bridge/quantize.py
sed -n '380,440p' examples/megatron_bridge/quantize.py
echo "== focused grouped-layout and recipe references =="
rg -n -C 5 \
'grouped_experts|moe_grouped_gemm|get_quant_config|per.?block|NVFP4|block_sizes|TEGroupedLinear|GroupedMLP|SequentialMLP' \
examples modelopt tests \
--glob '*.py' --glob '*.yaml' \
| head -n 1200Repository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== get_quant_config and parser definitions =="
sed -n '90,270p' examples/megatron_bridge/quantize.py
sed -n '1,95p' examples/megatron_bridge/quantize.py
echo "== all grouped-layout references outside the noisy kernel matches =="
rg -n -C 8 \
'moe_grouped_gemm|grouped_experts|TEGroupedLinear|GroupedMLP|SequentialMLP' \
examples modelopt tests \
--glob '*.py' \
--glob '!**/kernels/**'
echo "== recipe definitions and exporter references =="
rg -n -C 8 \
'NVFP4|nvfp4|per.?block|block_sizes|recipe|export' \
examples/megatron_bridge \
--glob '*.py' --glob '*.yaml' \
| head -n 1600Repository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== exact config resolution and save order =="
sed -n '215,285p' examples/megatron_bridge/quantize.py
sed -n '345,430p' examples/megatron_bridge/quantize.py
sed -n '430,495p' examples/megatron_bridge/quantize.py
echo "== grouped-linear implementation constraints =="
rg -n -C 12 \
'GroupedLinear|TEGroupedLinear|per-tensor|per_tensor|blockwise|block.?size|block_sizes|raise (ValueError|AssertionError)|assert' \
modelopt/torch/quantization/plugins modelopt/torch/export \
--glob '*.py' \
| head -n 1800
echo "== direct tests for grouped plus blockwise recipes =="
rg -n -C 15 \
'moe_grouped_gemm.*True|grouped.*True|TEGroupedMLP only supports|INT4_BLOCKWISE|blockwise' \
tests/gpu_megatron tests \
--glob '*.py' \
| head -n 1200Repository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("examples/megatron_bridge/quantize.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in {"main", "get_quant_config"}:
calls = []
for child in ast.walk(node):
if isinstance(child, ast.Call):
fn = child.func
if isinstance(fn, ast.Name):
name = fn.id
elif isinstance(fn, ast.Attribute):
name = f"{ast.unparse(fn.value)}.{fn.attr}"
else:
name = ast.unparse(fn)
if name in {"get_quant_config", "load_mbridge_model_from_hf", "mtq.quantize"}:
calls.append((child.lineno, name))
print(node.name, sorted(calls))
print("main source ranges:")
lines = path.read_text().splitlines()
for start, end in [(220, 270), (300, 375), (400, 490)]:
print(f"--- {start}:{end} ---")
for i in range(start, end + 1):
print(f"{i}: {lines[i-1]}")
PY
echo "== concise grouped/blockwise evidence =="
rg -n -C 4 \
'TEGroupedMLP only supports per-tensor quantization|INT4_BLOCKWISE_WEIGHT_ONLY_CFG|NVFP4_DEFAULT_CFG|moe_grouped_gemm' \
tests/gpu_megatron/torch/quantization/plugins/test_megatron.py \
tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.pyRepository: NVIDIA/Model-Optimizer
Length of output: 41021
Validate grouped-layout compatibility before model construction.
The grouped MoE path supports only per-tensor quantization. Resolve mtq_config before load_mbridge_model_from_hf() and reject per-block expert quantization when args.grouped_experts is enabled.
🤖 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/quantize.py` around lines 177 - 184, Update the
quantization setup in the main flow of quantize.py to resolve mtq_config before
calling load_mbridge_model_from_hf(). When args.grouped_experts is enabled,
validate that expert quantization uses per-tensor scaling and reject per-block
configurations before constructing the model; leave compatible configurations
and the non-grouped path unchanged.
| parser.add_argument( | ||
| "--calib_random_offset", | ||
| action="store_true", | ||
| help="Drop a random leading-token offset before packing calib windows (Megatron-LM --calib-use-random-offset).", | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject or document --calib_random_offset for VLM calibration.
Line [408] is reached only for text calibration. The image-text branch builds get_megatron_vlm_calibration_forward_loop without random_offset. A VLM run can therefore accept the flag and silently ignore it.
Reject this combination or state clearly that the option is text-only.
Proposed fail-fast check
use_image_calib = args.calib_dataset_name in get_supported_vlm_datasets()
+ if args.calib_random_offset and use_image_calib:
+ raise ValueError(
+ "--calib_random_offset is supported only for text calibration."
+ )Also applies to: 408-408
🤖 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/quantize.py` around lines 185 - 189, Update the
argument handling in the quantization flow around --calib_random_offset and
get_megatron_vlm_calibration_forward_loop so VLM calibration cannot silently
ignore the option: either reject its combination with VLM calibration via a
clear validation error, or explicitly document and enforce that the flag is
text-only. Preserve the existing random-offset behavior for text calibration.
| _MTP_HF_CONFIG_FIELDS = ("num_nextn_predict_layers", "mtp_num_hidden_layers", "mtp_num_layers") | ||
|
|
||
|
|
||
| def _hf_config_has_mtp(hf_cfg) -> bool: | ||
| """Whether an HF config declares MTP heads (checked top-level and under ``text_config``).""" | ||
| return any( | ||
| cfg is not None and getattr(cfg, field, 0) | ||
| for cfg in (getattr(hf_cfg, "text_config", None), hf_cfg) | ||
| for field in _MTP_HF_CONFIG_FIELDS | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Expect: tests or normalization code cover both mapping-backed and object-backed text_config values.
rg -n -C 4 \
'text_config|_hf_config_has_mtp|num_nextn_predict_layers|mtp_num_hidden_layers|mtp_num_layers' \
. --glob '*.py'Repository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="$(fd -t f '^quantize\.py$' examples modelopt tests | head -n 1)"
printf '%s\n' "FILE=$file"
wc -l "$file"
cat -n "$file" | sed -n '1,40p;270,345p'
printf '\n-- relevant definitions and calls --\n'
rg -n -C 8 \
'_hf_config_has_mtp|_MTP_HF_CONFIG_FIELDS|mtp_num_layers|MTP|text_config' \
"$file"
printf '\n-- nearby tests --\n'
rg -n -C 6 \
'hf_config_has_mtp|mtp_num_layers|num_nextn_predict_layers|text_config' \
tests examples --glob '*.py' \
| head -n 300Repository: NVIDIA/Model-Optimizer
Length of output: 5558
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="examples/megatron_bridge/quantize.py"
test -f "$file"
wc -l "$file"
cat -n "$file" | sed -n '1,55p;270,345p'
printf '\n-- relevant definitions and calls --\n'
rg -n -C 10 \
'_hf_config_has_mtp|_MTP_HF_CONFIG_FIELDS|mtp_num_layers|MTP|text_config' \
"$file"
printf '\n-- targeted tests and normalization code --\n'
rg -n -C 6 \
'_hf_config_has_mtp|num_nextn_predict_layers|mtp_num_hidden_layers|mtp_num_layers|text_config' \
tests examples --glob '*.py' \
| rg -v '(^|/)(alpamayo|speculative_decoding|gpu_megatron)' \
| head -n 300Repository: NVIDIA/Model-Optimizer
Length of output: 36728
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '-- loader definition and config flow --'
rg -n -C 12 \
'def load_mbridge_model_from_hf|load_mbridge_model_from_hf\(' \
. --glob '*.py' \
| head -n 240
printf '%s\n' '-- dependency versions and relevant config construction --'
rg -n -C 8 \
'transformers|Qwen2VLConfig|Qwen3VLConfig|text_config' \
pyproject.toml requirements*.txt setup.cfg setup.py tests/_test_utils/torch/transformers_models.py \
2>/dev/null | head -n 260
printf '%s\n' '-- current import and mapping support --'
sed -n '20,80p' examples/megatron_bridge/quantize.py
rg -n -C 5 \
'collections\.abc|Mapping|isinstance\(.*dict|isinstance\(.*Mapping' \
examples/megatron_bridge modelopt tests --glob '*.py' \
| head -n 220Repository: NVIDIA/Model-Optimizer
Length of output: 32967
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="modelopt/torch/utils/plugins/mbridge.py"
cat -n "$file" | sed -n '35,180p'
printf '\n-- config access in loader --\n'
rg -n -C 10 \
'AutoConfig|config|text_config|from_pretrained|from_hf' \
"$file" | head -n 260
printf '\n-- exact current helper behavior for object and mapping inputs --\n'
python3 - <<'PY'
from collections.abc import Mapping
from types import SimpleNamespace
fields = ("num_nextn_predict_layers", "mtp_num_hidden_layers", "mtp_num_layers")
def current(hf_cfg):
return any(
cfg is not None and getattr(cfg, field, 0)
for cfg in (getattr(hf_cfg, "text_config", None), hf_cfg)
for field in fields
)
def mapping_aware(hf_cfg):
def value(cfg, field):
return cfg.get(field, 0) if isinstance(cfg, Mapping) else getattr(cfg, field, 0)
return any(
cfg is not None and value(cfg, field)
for cfg in (getattr(hf_cfg, "text_config", None), hf_cfg)
for field in fields
)
cases = {
"object nested MTP": SimpleNamespace(
text_config=SimpleNamespace(mtp_num_layers=1)
),
"mapping nested MTP": SimpleNamespace(
text_config={"mtp_num_layers": 1}
),
"top-level MTP": SimpleNamespace(mtp_num_layers=1),
"no MTP": SimpleNamespace(text_config={"num_hidden_layers": 2}),
}
for name, cfg in cases.items():
print(name, "current=", current(cfg), "mapping_aware=", mapping_aware(cfg))
PYRepository: NVIDIA/Model-Optimizer
Length of output: 8804
🌐 Web query:
Transformers Qwen3VLConfig __init__ text_config dict normalized PretrainedConfig AutoBridge.from_hf_pretrained bridge.hf_pretrained.config
💡 Result:
The query involves two distinct technical domains: the configuration management for the Qwen3-VL model within the Hugging Face Transformers library, and the model conversion workflow using the NVIDIA Megatron-Bridge library [1][2]. Qwen3VLConfig Initialization The Qwen3VLConfig class serves as the configuration container for Qwen3-VL models. During initialization, the text_config and vision_config parameters allow for either passing a configuration dictionary or a PreTrainedConfig object [3]. In the implementation, these sub-configurations are normalized during the post-initialization phase (typically via post_init) [1][4]. If a parameter is passed as a dict, the configuration class automatically instantiates the appropriate sub-config class (Qwen3VLTextConfig or Qwen3VLVisionConfig) using the dictionary as kwargs [1][5]. If a parameter is None, the class defaults to an empty initialization of the corresponding sub-config [1][6]. Megatron-Bridge AutoBridge The AutoBridge.from_hf_pretrained method is a high-level API in the Megatron-Bridge library designed to load a Hugging Face model and prepare it for conversion to Megatron format [2][7]. When you call bridge = AutoBridge.from_hf_pretrained(model_path), the method: 1. Loads the model configuration in a thread-safe manner (often using safe_load_config_with_retry) [8][9]. 2. Validates the configuration against supported architectures [8]. 3. Returns an AutoBridge instance [7]. The bridge object contains an hf_pretrained attribute, which holds the loaded Hugging Face model or configuration. You can access the underlying Hugging Face configuration object via bridge.hf_pretrained.config [8][7]. This enables programmatic access to the model's structural parameters after the bridge has been instantiated for conversion tasks [10].
Citations:
- 1: https://github.com/huggingface/transformers/blob/v5.12.0/src/transformers/models/qwen3_vl/configuration_qwen3_vl.py
- 2: https://docs.nvidia.com/nemo/megatron-bridge/latest/bridge-guide.html
- 3: https://huggingface.co/docs/transformers/en/model_doc/qwen3_vl
- 4: https://github.com/huggingface/transformers/blob/c472755e/src/transformers/models/qwen3_vl/configuration_qwen3_vl.py
- 5: https://github.com/huggingface/transformers/blob/1f8daee0/src/transformers/models/qwen3_vl/configuration_qwen3_vl.py
- 6: https://github.com/huggingface/transformers/blob/v4.57.2/src/transformers/models/qwen3_vl/configuration_qwen3_vl.py
- 7: https://docs.nvidia.com/nemo/megatron-bridge/nightly/apidocs/bridge/bridge.models.conversion.auto_bridge.html
- 8: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/auto_bridge.py
- 9: https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.hf_pretrained.safe_config_loader.html
- 10: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-tech-details.md
🌐 Web query:
site:github.com/NVIDIA-NeMo/Megatron-Bridge AutoBridge from_hf_pretrained hf_pretrained config PretrainedConfig
💡 Result:
AutoBridge is a component of the NVIDIA-NeMo Megatron-Bridge library designed to facilitate bidirectional conversion between Hugging Face (HF) and Megatron-Core model formats [1][2]. The method AutoBridge.from_hf_pretrained(path, kwargs) is the primary entry point for loading models [1][3]. It automatically detects the model architecture from the provided Hugging Face model path (or ID), selects the appropriate bridge, and validates the configuration to ensure compatibility for conversion operations [1][3][4]. Key technical aspects include: - Configuration Handling: The method loads the Hugging Face configuration, performing architecture validation before proceeding with further operations [3]. It also handles thread-safe configuration loading and supports the trust_remote_code parameter for loading custom models [3]. - PretrainedConfig Interaction: While AutoBridge.from_hf_pretrained loads from a path or model ID, the library also provides AutoBridge.from_hf_config(config: PretrainedConfig), which accepts an already-loaded PretrainedConfig object [1][4]. However, documentation advises using from_hf_pretrained for export workflows to ensure that all necessary artifacts, such as tokenizers, are correctly included [1]. - Workflow: Once an AutoBridge instance is created, it can be used to convert models to a Megatron provider via.to_megatron_provider, which then allows for the instantiation of distributed Megatron models [5][4]. The library is optimized for handling parallelism-aware conversion (TP/PP/EP) and utilizes the MappingRegistry to define weight mappings between HF and Megatron parameters [4][6][2].
Citations:
- 1: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-guide.md
- 2: http://github.com/NVIDIA-NeMo/Megatron-Bridge
- 3: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/auto_bridge.py
- 4: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-tech-details.md
- 5: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/README.md
- 6: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/model_bridge.py
🌐 Web query:
Transformers configuration text_config isinstance PretrainedConfig dict all multimodal config classes Qwen3VL Qwen2VL Gemma3 Mistral3
💡 Result:
In the Hugging Face Transformers library, multimodal models use a composite configuration pattern to manage separate sub-architectures (like vision and text encoders) [1][2]. The text_config attribute in these configuration classes is typically a PreTrainedConfig object (or a dictionary that is subsequently converted into one) [3][4][5][6]. Multimodal configuration classes—such as Qwen2VLConfig, Qwen3VLConfig, Gemma3Config, and Mistral3Config—commonly use an post_init method to check the type of text_config [4][5][6][7]. When initializing these models, if text_config is passed as a dict, the configuration class uses its internal mapping or a specialized configuration class to instantiate it as a proper PreTrainedConfig object [3][4][8][5]. If it is already a PreTrainedConfig instance, it is assigned directly [9][6]. Summary of specified models: - Qwen2VLConfig: Uses a sub_configs mapping to define text_config as Qwen2VLTextConfig [5]. - Qwen3VLConfig: Defines text_config as Qwen3VLTextConfig [10][6]. - Gemma3Config: Defines text_config as Gemma3TextConfig [4]. - Mistral3Config: Uses AutoConfig to handle the text_config [7]. Additionally, the base PreTrainedConfig class provides a get_text_config method, which is designed to return the relevant text configuration from a composite model, helping to standardize access across different generations of multimodal architectures [1][2].
Citations:
- 1: https://huggingface.co/docs/transformers/main_classes/configuration
- 2: https://github.com/huggingface/transformers/blob/main/src/transformers/configuration_utils.py
- 3: https://github.com/huggingface/transformers/blob/v5.8.1/src/transformers/models/paligemma/configuration_paligemma.py
- 4: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma3/configuration_gemma3.py
- 5: https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py
- 6: https://github.com/huggingface/transformers/blob/v5.3.0/src/transformers/models/qwen3_vl/configuration_qwen3_vl.py
- 7: https://github.com/huggingface/transformers/blob/v5.8.1/src/transformers/models/mistral3/configuration_mistral3.py
- 8: https://github.com/huggingface/transformers/blob/v5.8.1/src/transformers/models/blip/configuration_blip.py
- 9: https://github.com/huggingface/transformers/blob/v5.5.3/src/transformers/models/vision_text_dual_encoder/configuration_vision_text_dual_encoder.py
- 10: https://huggingface.co/docs/transformers/main/en/model_doc/qwen3_vl
🌐 Web query:
site:github.com/NVIDIA-NeMo/Megatron-Bridge "text_config" "AutoBridge" model bridge
💡 Result:
AutoBridge is a high-level API provided by the NVIDIA NeMo Megatron-Bridge library, designed to facilitate bidirectional model conversion and interoperability between Hugging Face (transformers) and NVIDIA Megatron Core formats [1][2]. Key aspects of AutoBridge include: Unified Entry Point: AutoBridge serves as the primary interface for users, automatically detecting the model architecture from a Hugging Face checkpoint (hub or local path) and selecting the appropriate architecture-specific bridge class to manage the conversion process [1][3][4]. Bidirectional Conversion: Hugging Face to Megatron: It allows users to import Hugging Face models into Megatron for distributed training (handling tensor, pipeline, expert, and sequence parallelism) [1][2][3]. Megatron to Hugging Face: It enables exporting trained Megatron models back into Hugging Face format for deployment or use with other inference engines [1][2][3]. Workflow Integration: It simplifies workflows by providing high-level methods such as from_hf_pretrained() for initialization, to_megatron_provider() for configuring model parallelism, and to_megatron_model() for instantiation [1][3][4]. The bridge is parallelism-aware, managing weight mapping and tensor distribution across different Megatron parallelism configurations (TP/PP/VPP/CP/EP/ETP) [1][2]. It supports memory-efficient operations through per-parameter streaming using safetensors [1][2]. As part of the NeMo Framework ecosystem, Megatron-Bridge ensures checkpoint integrity and accuracy through built-in verification mechanisms during the conversion process [2]. It is widely used for scaling language, vision-language, and multimodal models to large-scale distributed environments [2][5].
Citations:
- 1: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-guide.md
- 2: https://github.com/NVIDIA-NeMo/Megatron-Bridge
- 3: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/auto_bridge.py
- 4: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-tech-details.md
- 5: kimi k2.5 support NVIDIA-NeMo/Megatron-Bridge#2743
🏁 Script executed:
#!/bin/bash
set -euo pipefail
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
for spec in \
"transformers_qwen3_vl.py|https://raw.githubusercontent.com/huggingface/transformers/v5.3.0/src/transformers/models/qwen3_vl/configuration_qwen3_vl.py" \
"transformers_qwen2_vl.py|https://raw.githubusercontent.com/huggingface/transformers/v4.57.2/src/transformers/models/qwen2_vl/configuration_qwen2_vl.py" \
"transformers_gemma3.py|https://raw.githubusercontent.com/huggingface/transformers/v4.57.2/src/transformers/models/gemma3/configuration_gemma3.py" \
"transformers_mistral3.py|https://raw.githubusercontent.com/huggingface/transformers/v5.8.1/src/transformers/models/mistral3/configuration_mistral3.py" \
"bridge_auto.py|https://raw.githubusercontent.com/NVIDIA-NeMo/Megatron-Bridge/main/src/megatron/bridge/models/conversion/auto_bridge.py"
do
name="${spec%%|*}"
url="${spec#*|}"
curl --fail --silent --show-error --location "$url" > "$tmp/$name"
printf '\n-- %s --\n' "$name"
rg -n -C 8 'text_config|hf_pretrained|from_hf_pretrained|PretrainedConfig' "$tmp/$name" | head -n 180 || true
doneRepository: NVIDIA/Model-Optimizer
Length of output: 27932
🏁 Script executed:
#!/bin/bash
set -euo pipefail
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
url="https://raw.githubusercontent.com/NVIDIA-NeMo/Megatron-Bridge/main/src/megatron/bridge/models/conversion/auto_bridge.py"
curl --fail --silent --show-error --location "$url" > "$tmp/auto_bridge.py"
cat -n "$tmp/auto_bridge.py" | sed -n '90,160p;480,620p'Repository: NVIDIA/Model-Optimizer
Length of output: 10980
Handle mapping-backed text_config values
If a mapping reaches _hf_config_has_mtp, getattr misses all nested MTP fields and suppresses the warning even though line 312 drops the MTP heads. Use a mapping-aware accessor and add tests for object-backed and mapping-backed configurations.
🤖 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/quantize.py` around lines 291 - 300, Update
_hf_config_has_mtp to read MTP fields from both attribute-backed objects and
mapping-backed configs, including nested text_config mappings, so existing MTP
detection remains accurate before line 312 drops the heads. Add tests covering
object-backed and mapping-backed top-level and text_config configurations.
Source: MCP tools
…QAD support
Also fixes a ModelOpt bug that affects ANY model with an untied lm_head quantized through
Megatron-Bridge, not just Nemotron: the output layer was silently exported as BF16.
An untied `lm_head` (`output_layer`) was silently exported as BF16 instead of NVFP4
whenever the model was built with Megatron-Bridge, even though the recipe enabled
`*output_layer*weight_quantizer`. The same recipe under Megatron-LM quantized it
correctly, so this was not a configuration problem.
`_MegatronParallelLinear.sharded_state_dict()` special-cases `output_layer` and asks
`megatron.training.get_args()` whether embeddings are untied. Megatron-Bridge has no
global args store, so the call raises and the handler falls back to "tied", taking the
early return that drops all quantizer state. Fixing that alone is not sufficient: the
dist-checkpoint loader silently skips any checkpoint key the model does not advertise,
and `sharded_state_dict()` can only advertise a buffer that already exists, so the
calibrated scales in the checkpoint had nowhere to land.
* `_resolve_output_layer_untied()` reads `share_embeddings_and_output_weights` off the
model, which Megatron-Core carries under both frameworks, and records it on the
config so `sharded_state_dict()` can consult it. `get_args()` remains the fallback,
so Megatron-LM behavior is unchanged, and an unknown result still means "tied".
* Materialize missing weight-quantizer scale buffers before the load plan is built.
`_amax` must be allocated flat as `[numel // block, 1]`: `_process_quantizer_amax`
exposes it to the checkpoint as an `[out_features, blocks]` view over the same
storage, so the loader writes straight through. Allocating the viewed shape loads
successfully but leaves the wrong in-memory rank, which then breaks the exporter's
scale math. `_global_amax` is registered directly because its property lives on
`StaticBlockScaleQuantizer` and the module is still a plain `TensorQuantizer` here.
The export example now fails loudly instead of quietly emitting BF16 when an enabled
NVFP4 weight quantizer is missing either scale, naming the module and the attribute.
Static-block NVFP4 needs both; an `_amax`-only check let a half-restored quantizer
through, which failed much later inside `NVFP4QTensor.quantize` where `scale * scale_2`
broadcasts `[N, 1]` against `[N]`. `MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1` restores the
previous behavior, reported rather than silent.
Also included: Megatron-Bridge PTQ/QAD enablement for Nemotron-style hybrid MoE models
(gate the non-grouped MoE spec to the quantized student so the BF16 teacher is built
with its natural spec, SFT-masked distillation, student initialization from a Megatron
checkpoint, calibration random offset).
Verified end to end against a Megatron-LM-produced reference: `lm_head.weight` is now
`U8 [131072, 1344]` with `weight_scale` and `weight_scale_2`, 18487 keys and 72
excluded modules, matching the reference exactly (previously BF16 `[131072, 2688]`,
18485 keys, 73 excluded with `lm_head` among them).
Two Megatron-Bridge compatibility shims were removed after being shown unnecessary:
a `DistillationProvider.to_cfg_dict` monkeypatch (a 5-iteration distillation run trains
and checkpoints cleanly without it) and an `InferenceCudaGraphScope` enum stub added for
a Megatron-LM-PTQ import path that is not used (zero occurrences across a full
PTQ/QAD/export run).
Signed-off-by: James Shen <yueshen@nvidia.com>
The two example scripts no longer hard-code model-shape assumptions. MoE expert grouping is a
`--grouped_experts` flag on both `quantize.py` and the exporter, defaulting to non-grouped so
existing behavior is unchanged; per-block NVFP4 requires non-grouped because TEGroupedLinear can
only represent a per-tensor scale, while per-tensor recipes can now opt into faster grouped GEMM.
The exporter reads `mtp_num_layers` from the checkpoint's run_config.yaml instead of assuming 0,
and `quantize.py` now warns when it drops MTP heads, matching prune_minitron.py. Expert grouping is
deliberately NOT derived from run_config.yaml: a MambaModelProvider sets the layout via
mamba_stack_spec, so a non-grouped checkpoint still records `moe_grouped_gemm: true` and trusting it
would build a mismatched model.
Signed-off-by: James Shen <yueshen@nvidia.com>
Review fixes: the export guard no longer excludes StaticBlockScaleQuantizer, which is the class that
owns `_global_amax` -- excluding it skipped exactly the case the guard exists to catch -- and it now
filters on `is_enabled` to match the message it prints. The scale-buffer materialization checks
`in_features % block_size`, matching the `view(weight.shape[0], -1)` performed later rather than
total element count, and warns instead of silently leaving the buffers unallocated.
Signed-off-by: James Shen <yueshen@nvidia.com>
11ab426 to
72d3cb5
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
modelopt/torch/utils/dataset_utils.py (1)
765-765: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new
random_offsetargument.
get_dataset_dataloaderis a public API, but itsArgssection omitsrandom_offset. Document that the option applies only whenpack=True, shifts the token stream by a bounded prefix, and is ignored whenpack=False.Proposed docstring addition
pack: If True, use global-stream document packing ... + random_offset: If True with ``pack=True``, shift the packed token stream + by a bounded random prefix. Ignored when ``pack=False``. distributed: If True, shard the dataset across ranks ...As per coding guidelines, "Document public and higher-level APIs with docstrings, including examples when useful; keep internal helpers self-documenting."
🤖 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 get_dataset_dataloader docstring’s Args section to document random_offset, stating that it applies only when pack=True, shifts the token stream by a bounded prefix, and is ignored when pack=False.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.
Nitpick comments:
In `@modelopt/torch/utils/dataset_utils.py`:
- Line 765: Update the get_dataset_dataloader docstring’s Args section to
document random_offset, stating that it applies only when pack=True, shifts the
token stream by a bounded prefix, and is ignored when pack=False.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8e520174-a9a5-4406-81a5-fc6eea0ed10e
📒 Files selected for processing (7)
examples/megatron_bridge/distill.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/quantize.pymodelopt/torch/distill/plugins/megatron.pymodelopt/torch/quantization/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/megatron_calibration.py
🚧 Files skipped from review as they are similar to previous changes (6)
- examples/megatron_bridge/export_quantized_megatron_to_hf.py
- examples/megatron_bridge/quantize.py
- modelopt/torch/distill/plugins/megatron.py
- modelopt/torch/quantization/plugins/megatron.py
- examples/megatron_bridge/distill.py
- modelopt/torch/utils/plugins/megatron_calibration.py
1 participant
What does this PR do?
Type of change: Bug fix + new feature
Adds Megatron-Bridge support for W4A16 NVFP4 (
four_over_six) PTQ and QAD on Nemotron-stylehybrid Mamba/attention MoE models, and fixes a ModelOpt bug found while doing so.
The bug
A recipe enabling
*output_layer*weight_quantizerproducedlm_head.weightas BF16[vocab, hidden]instead of packed NVFP4. The same recipe file under Megatron-LM quantized itcorrectly, so this was never a configuration problem. The result is a checkpoint that does not match
its own recipe, is larger than intended, and keeps the output projection — on the hot path for every
generated token — at full precision.
Root cause
_MegatronParallelLinear.sharded_state_dict()special-casesoutput_layerand asksMegatron-LM's global args whether embeddings are untied:
Megatron-Bridge has no global args store, so the call raises and the handler concludes "tied" for
a model that is untied. Observed directly:
Fixing only that is not sufficient: the dist-checkpoint loader silently skips any checkpoint key
the model does not advertise, and
sharded_state_dict()can only advertise a buffer that alreadyexists. On restore the quantizer exists and is enabled but its scale buffers do not, so the
calibrated values in the checkpoint have nowhere to land.
The fix
_resolve_output_layer_untied(model)readsshare_embeddings_and_output_weightsoff themodel — Megatron-Core carries it under both frameworks (Bridge sets it from the HF config,
Megatron-LM from
--untie-embeddings-and-output-weights) — and records it on the config sosharded_state_dict()can consult it.get_args()remains the fallback, so Megatron-LMbehavior is unchanged, and an unknown result still means "tied" (the conservative branch).
loader has somewhere to write. Two details that are easy to get wrong:
_amaxmust be allocated flat as[numel // block, 1]._process_quantizer_amaxexposesit to the checkpoint as an
[out_features, blocks]view over the same storage, so theloader writes straight through. Allocating the viewed shape loads successfully but leaves the
wrong in-memory rank, which then breaks the exporter's scale math.
_global_amaxis registered withregister_bufferdirectly: its property lives onStaticBlockScaleQuantizer, and on restore the module is still a plainTensorQuantizer.Export guard
examples/megatron_bridge/export_quantized_megatron_to_hf.pynow raises instead of quietlyemitting BF16 when an enabled NVFP4 weight quantizer is missing either scale, naming the module and
the attribute. Static-block NVFP4 needs both; an
_amax-only check let a half-restored quantizerthrough and failed much later inside
NVFP4QTensor.quantize, wherescale * scale_2broadcasts[N, 1]against[N]into anN x Nallocation.MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1restores theprevious behavior, reported rather than silent.
A silent fallback is what let this ship unnoticed; making it loud is arguably the more important
half of the change.
Megatron-Bridge PTQ/QAD enablement
previously applied to the BF16 teacher too, misplacing its MoE expert weights so the teacher
predicted ground truth far worse than its own student, and QAD distilled toward a broken teacher.
--student_nongrouped_experts(default off, so existing grouped behavior is preserved)--sft,--sft_dataset_root) — loss masked to completion tokensDistillationProvider.provideis wrapped; marked TODO for upstream replacement)--calib_random_offset, matching Megatron-LM's--calib-use-random-offsetTesting
Verified end to end on a hybrid Mamba/attention MoE model, against a Megatron-LM-produced
reference checkpoint:
lm_head.weight[131072, 2688]U8 [131072, 1344]U8 [131072, 1344]lm_head.weight_scale/_scale_2hf_quant_configexclusionslm_head)lm_head)lm_head)Full pipeline exercised: PTQ → QAD (200 iters, 0 NaN) → HF export → compressed-tensors conversion →
vLLM serving (coherent generation) → 4 accuracy benchmarks, with the quantized
lm_headsurvivingevery stage. The fix also propagates through QAD: the QAD checkpoint now contains
output_layer.weight_quantizer._amax/._global_amax, which it previously lacked entirely.No new unit tests. Reproducing this needs a Megatron-Bridge model, a quantized checkpoint and a
dist-checkpoint round-trip, which does not fit the existing unit-test harness. Guidance on where a
regression test belongs would be welcome — the highest-value assertion is that
sharded_state_dict()advertises the
output_layerweight-quantizer keys for an untied model under Bridge.Before your PR is "Ready for review"
get_args()path; the new flagdefaults off; unknown tiedness still resolves to the previous conservative branch. The export guard
is the one behavior change (silent BF16 → error), reversible via
MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1.no new dependencies.
happy to add an entry; this is arguably a critical bug fix and probably warrants one.
Additional Information
Two Megatron-Bridge compatibility shims were removed after being shown unnecessary rather than
assumed so: a
DistillationProvider.to_cfg_dictmonkeypatch (a 5-iteration distillation run trainsand checkpoints cleanly without it) and an
InferenceCudaGraphScopeenum stub for an unusedMegatron-LM-PTQ import path (zero occurrences across a full PTQ/QAD/export run).
Summary by CodeRabbit
New Features
Bug Fixes