feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support #2055
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -174,6 +174,19 @@ def get_args() -> argparse.Namespace: | |
| parser.add_argument( | ||
| "--calib_num_samples", type=int, default=1024, help="Number of samples for calibration" | ||
| ) | ||
| 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.", | ||
| ) | ||
|
Comment on lines
+177
to
+184
Contributor
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 🤖 Prompt for AI Agents |
||
| 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).", | ||
| ) | ||
|
Comment on lines
+185
to
+189
Contributor
There was a problem hiding this comment. 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win Reject or document Line [408] is reached only for text calibration. The image-text branch builds 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 |
||
| parser.add_argument("--calib_batch_size", type=int, default=1, help="Calibration batch size") | ||
| parser.add_argument( | ||
| "--seq_length", | ||
|
|
@@ -275,6 +288,18 @@ def get_quant_config(args: argparse.Namespace) -> dict: | |
| return mtq_config | ||
|
|
||
|
|
||
| _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 | ||
| ) | ||
|
Comment on lines
+291
to
+300
Contributor
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:
💡 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:
🌐 Web query:
💡 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:
🌐 Web query:
💡 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:
🌐 Web query:
💡 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 Citations:
🏁 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 If a mapping reaches 🤖 Prompt for AI AgentsSource: MCP tools |
||
|
|
||
|
|
||
| def main(args: argparse.Namespace): | ||
| bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf( | ||
| hf_model_name_or_path=args.hf_model_name_or_path, | ||
|
|
@@ -284,14 +309,27 @@ def main(args: argparse.Namespace): | |
| "pipeline_model_parallel_size": args.pp_size, | ||
| "expert_model_parallel_size": args.ep_size, | ||
| "context_parallel_size": args.cp_size, | ||
| "mtp_num_layers": 0, # MTP not supported during calibration | ||
| "expert_tensor_parallel_size": 1, # Expert tensor parallelism is not supported | ||
| "pipeline_dtype": torch.bfloat16, | ||
| "seq_length": args.seq_length, | ||
| "gradient_accumulation_fusion": False, # not supported | ||
| }, | ||
| init_model_parallel=True, | ||
| # Default non-grouped: per-block NVFP4 needs it (TEGroupedLinear is per-tensor only). | ||
| # Opt into grouped for per-tensor recipes, where grouped GEMM is faster. | ||
| moe_grouped_gemm=args.grouped_experts, | ||
| ) | ||
|
|
||
| # `mtp_num_layers=0` above drops MTP heads: calibration does not support them. Say so rather | ||
| # than silently shipping a checkpoint without a head the model declares. | ||
| if _hf_config_has_mtp(bridge.hf_pretrained.config): | ||
| warn_rank_0( | ||
| "Dropping Multi-Token Prediction (MTP): calibration does not support it. The exported " | ||
| "checkpoint will not contain MTP weights and standard autoregressive inference is " | ||
| "unaffected. To use MTP speculative decoding, run a separate phase with mtp_num_layers>0." | ||
| ) | ||
|
|
||
| # Only the language model is quantized (vision tower + projector stay full precision) | ||
| language_model = getattr(unwrapped_model, "language_model", unwrapped_model) | ||
| is_vlm = language_model is not unwrapped_model | ||
|
|
@@ -367,6 +405,7 @@ def main(args: argparse.Namespace): | |
| seq_length=args.seq_length, | ||
| batch_size=args.calib_batch_size, | ||
| pack=True, # Megatron pretraining-style global-stream document packing | ||
| random_offset=args.calib_random_offset, | ||
| ) | ||
|
|
||
| # Run text prefill on the language model: we quantize the root (a VLM root forward expects | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[IMPORTANT Compatibility] Hardcoding
mtp_num_layers=0andmoe_grouped_gemm=Falsein 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_hfdefaultsmoe_grouped_gemm=True(modelopt/torch/utils/plugins/mbridge.py:47), and the model structure built here must match the checkpoint being loaded byload_modelopt_megatron_checkpoint. Forcing both values means: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.Same concern for
quantize.py:292,299, which hardcodes the identical pair on the calibration path. Theremtp_num_layers=0is defensible (MTP genuinely isn't supported during calibration), butmoe_grouped_gemm=Falseis 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:
then pass through only when set, leaving the provider default otherwise. Note the mirrored
--student_nongrouped_expertsflag added indistill.pygets this right — it defaults off precisely to preserve committed grouped behavior. These two sites should follow the same principle.