From 72d3cb5e714d7d8dfe6aa9adb0fc2a93583ebb64 Mon Sep 17 00:00:00 2001 From: James Shen Date: Mon, 3 Aug 2026 10:42:18 -0700 Subject: [PATCH] feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/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 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 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 --- examples/megatron_bridge/distill.py | 299 +++++++++++++----- .../export_quantized_megatron_to_hf.py | 102 ++++++ examples/megatron_bridge/quantize.py | 39 +++ modelopt/torch/distill/plugins/megatron.py | 46 +++ .../torch/quantization/plugins/megatron.py | 103 +++++- modelopt/torch/utils/dataset_utils.py | 20 +- .../utils/plugins/megatron_calibration.py | 4 + 7 files changed, 509 insertions(+), 104 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index dfd404fab30..66e636c1ca0 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -25,15 +25,18 @@ import os import torch -from _distillation_provider import convert_to_distillation_provider -from export_distilled_megatron_to_hf import export_llm_to_hf, save_vlm_to_hf from megatron.bridge import AutoBridge +from megatron.bridge.models.distillation_provider import ( + DistillationProvider, + convert_to_distillation_provider, +) from megatron.bridge.recipes.utils.optimizer_utils import ( distributed_fused_adam_with_cosine_annealing, ) from megatron.bridge.training.config import ( CheckpointConfig, ConfigContainer, + FinetuningDatasetConfig, GPTDatasetConfig, LoggerConfig, MockGPTDatasetConfig, @@ -46,18 +49,83 @@ from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig -from megatron.core.utils import unwrap_model from transformers import AutoConfig import modelopt.torch.distill as mtd +import modelopt.torch.distill.plugins.megatron as mtd_mcore import modelopt.torch.utils.distributed as dist -from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 +from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.plugins.mbridge import load_modelopt_megatron_checkpoint with contextlib.suppress(ModuleNotFoundError): import modelopt.torch.puzzletron.plugins.mbridge # noqa: F401 +# TODO: Megatron-Bridge does not (yet) expose a hook to initialize the student before the +# knowledge-distillation conversion, so we patch ``DistillationProvider.provide`` to do it. Replace +# this block once a first-class mechanism is available upstream. +# +# Maps id(distill_provider) -> megatron_checkpoint_path for providers whose student should be +# initialized from a Megatron checkpoint. A registry is used (instead of an instance attribute) +# because a DistillationProvider proxies attribute assignment to its teacher once the teacher is +# set, so anything stored on the instance would leak onto the teacher. +_MEGATRON_STUDENT_CKPT_PATHS: dict[int, str] = {} + +_original_distill_provide = DistillationProvider.provide + + +def _distill_provide_with_megatron_student( + self, pre_process=None, post_process=None, vp_stage=None +): + """Replacement for ``DistillationProvider.provide`` that can initialize the student from a ckpt. + + For providers registered in ``_MEGATRON_STUDENT_CKPT_PATHS``, the student is built and its weights + (plus, for a quantized checkpoint, the ModelOpt quantize mode) are restored from the Megatron + checkpoint *before* the knowledge-distillation conversion -- otherwise the quantize mode is lost, + since ``restore_sharded_modelopt_state`` is a no-op once a model is already converted. The rest + mirrors the upstream implementation. Patched at the class level (not the instance) to avoid the + teacher-proxying issue described on ``_MEGATRON_STUDENT_CKPT_PATHS``. + """ + if vp_stage is not None: + raise ValueError("ModelOpt KD currently does not support virtual-pipeline parallel.") + + megatron_path = _MEGATRON_STUDENT_CKPT_PATHS.get(id(self)) + if megatron_path is None: + # If a path was registered (for some provider) but this provide() call doesn't match, + # the provider was likely copied/wrapped between convert_to_distillation_provider() and now, + # so the id()-keyed lookup silently misses. Fail loudly rather than train an uninitialized + # student (this script only ever builds one DistillationProvider). + if _MEGATRON_STUDENT_CKPT_PATHS: + raise RuntimeError( + "DistillationProvider.provide() found no registered Megatron-student checkpoint path " + "for this provider, but one was registered for a different provider id -- the provider " + "was likely copied/wrapped. Update this workaround." + ) + return _original_distill_provide(self, pre_process, post_process, vp_stage) + + student_model = self._super_class.provide(self, pre_process, post_process, vp_stage) + print_rank_0(f"Loading student weights from Megatron checkpoint {megatron_path}") + load_modelopt_megatron_checkpoint([student_model], megatron_path) + # Hack to get teacher's pre-wrap hooks called to potentially load HF weights + teacher_model = self.teacher.provide_distributed_model( + wrap_with_ddp=False, mixed_precision_wrapper=None + )[0] + kd_cfg = mtd_mcore.setup_distillation_config( + self.kd_config, student_model.config, teacher_model.config + ) + modelopt_cfg = { + "teacher_model": teacher_model, + "criterion": kd_cfg.criterion, + "loss_balancer": kd_cfg.loss_balancer, + } + kd_model = mtd.convert(student_model, mode=[("kd_loss", modelopt_cfg)]) + mtd_mcore.adjust_distillation_model_for_mcore(kd_model, kd_cfg) + return kd_model + + +DistillationProvider.provide = _distill_provide_with_megatron_student + + def get_args(): """Parse command-line arguments.""" parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.") @@ -75,6 +143,16 @@ def get_args(): help="HuggingFace model name or path for the teacher (e.g. Qwen/Qwen3-8B)", ) parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code") + parser.add_argument( + "--student_nongrouped_experts", + action="store_true", + help=( + "Build the quantized student with non-grouped MoE experts. Required for STATIC-BLOCK " + "NVFP4 recipes (e.g. four_over_six): TEGroupedLinear only supports per-tensor scales, " + "not per-block. Leave OFF (default) for dynamic NVFP4 / grouped-expert checkpoints " + "(e.g. Nemotron-3-Nano). Never applied to the BF16 teacher." + ), + ) parser.add_argument( "--student_megatron_path", type=str, @@ -110,6 +188,20 @@ def get_args(): parser.add_argument( "--use_mock_data", action="store_true", help="Use mock data instead of --data_paths" ) + parser.add_argument( + "--sft", + action="store_true", + help="SFT-masked distillation: read raw prompt-completion jsonl from --sft_dataset_root and " + "mask the loss to the completion (assistant response) tokens. Uses GPTSFTDatasetConfig + the " + "real (HuggingFace) tokenizer instead of the pretraining GPTDataset + NullTokenizer.", + ) + parser.add_argument( + "--sft_dataset_root", + type=str, + default=None, + help="Directory containing training.jsonl / validation.jsonl with prompt-completion " + '{"input": , "output": } records (used with --sft).', + ) # Training & Eval arguments parser.add_argument( "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving" @@ -188,18 +280,20 @@ def get_args(): type=str, required=False, default=None, - help="Reference HF model with a homogeneous architecture, used as the export template for a " - "heterogeneous (Puzzletron/NAS) student's weights. Defaults to --student_hf_path, which is " - "correct for homogeneous students; unused for VLMs.", + help="HuggingFace model ID to use as template for export (e.g., Qwen/Qwen3-0.6B). " + "Should match the base architecture of the student model if --hf_export_path is provided.", ) args = parser.parse_args() # Sanity checks - if not args.use_mock_data and not args.data_paths: + if args.sft: + if not args.sft_dataset_root: + raise ValueError("--sft requires --sft_dataset_root (dir with training.jsonl/validation.jsonl).") + elif not args.use_mock_data and not args.data_paths: raise ValueError("Must provide either --data_paths or set --use_mock_data.") - if args.student_hf_model is None: - args.student_hf_model = args.student_hf_path + if args.hf_export_path and not args.student_hf_model: + raise ValueError("Must provide --student_hf_model if --hf_export_path is provided.") print_args(args) @@ -211,7 +305,7 @@ def main(args: argparse.Namespace): tensorboard_dir = os.path.join(args.output_dir, "tb_logs") # Build student and teacher model providers - def _build_model_provider(hf_path, load_weights=True): + def _build_model_provider(hf_path, load_weights=True, quantized=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) provider = bridge.to_megatron_provider(load_weights=load_weights) @@ -224,6 +318,32 @@ def _build_model_provider(hf_path, load_weights=True): provider.expert_model_parallel_size = args.ep_size provider.expert_tensor_parallel_size = 1 # Expert tensor parallelism is not supported provider.seq_length = args.seq_length + # Match the PTQ/quantize.py setup: MTP is not supported during QAD, and NVFP4 per-block + # quantization requires non-grouped experts (TEGroupedLinear only supports per-tensor). + # For a hybrid Mamba provider the layer SPEC must be rebuilt with moe_grouped_gemm=False -- + # setting the flag alone does not propagate. Mirror modelopt's load_mbridge_model_from_hf. + provider.mtp_num_layers = 0 + from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec + + if quantized and args.student_nongrouped_experts: + # Static-block NVFP4 students need non-grouped experts (TEGroupedLinear can't do per-block + # scales). OFF by default = grouped = committed behavior (works for dynamic NVFP4 like + # Nano-3). NEVER applied to the BF16 teacher (would misplace its MoE experts). + if hasattr(provider, "mamba_stack_spec"): + provider.mamba_stack_spec = get_te_mamba_stack_spec(moe_grouped_gemm=False) + elif (getattr(provider, "num_moe_experts", 0) or 0) > 0: + provider.moe_grouped_gemm = False + # Regularize the MoE router during QAD so it does not degenerate. Jenny's working Megatron-LM + # QAD uses `--moe-aux-loss-coeff 1e-4 --moe-router-load-balancing-type seq_aux_loss`; without + # it our router weights drifted the most (~8% vs ~1% elsewhere) and the MoE broke. Applies to + # both providers, but only the (trained) student's aux loss affects optimization. + if (getattr(provider, "num_moe_experts", 0) or 0) > 0: + provider.moe_router_load_balancing_type = "seq_aux_loss" + provider.moe_aux_loss_coeff = 1e-4 + 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" if args.recompute_granularity is not None: provider.recompute_granularity = args.recompute_granularity provider.recompute_method = args.recompute_method @@ -245,52 +365,25 @@ def _build_model_provider(hf_path, load_weights=True): # Gradient accumulation fusion is not supported with ModelOpt quantized models. Disable it # before the model is built so the student's linear layers are constructed accordingly. student_provider.gradient_accumulation_fusion = False - teacher_provider = _build_model_provider(args.teacher_hf_path) + teacher_provider = _build_model_provider(args.teacher_hf_path, quantized=False) + # 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 ) if args.student_megatron_path: - # QAD: restore the quantized student weights + ModelOpt state before the KD conversion (a no-op - # once converted). Prepend so this runs before the provider's KD-conversion pre-wrap hook. if student_has_modelopt_state: print_rank_0( f"Detected ModelOpt state in {args.student_megatron_path}; " "restoring quantizers for Quantization Aware Distillation (QAD)." ) - - def _restore_student_hook(model_chunks): - print_rank_0( - f"Loading student weights from Megatron checkpoint {args.student_megatron_path}" - ) - load_modelopt_megatron_checkpoint( - [unwrap_model(model_chunks[0])], args.student_megatron_path - ) - return model_chunks - - distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) + # Register so the patched DistillationProvider.provide initializes this provider's student + # from the Megatron checkpoint (see _distill_provide_with_megatron_student). + _MEGATRON_STUDENT_CKPT_PATHS[id(distill_provider)] = args.student_megatron_path # Build optimizer and scheduler optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing( @@ -301,24 +394,48 @@ def _restore_student_hook(model_chunks): ) # Build dataset config - dataset_kwargs = { - "seq_length": args.seq_length, - "path_to_cache": args.data_path_to_cache, - "random_seed": args.seed, - "reset_attention_mask": False, - "reset_position_ids": False, - "eod_mask_loss": False, - "num_dataset_builder_threads": 1, - "data_sharding": True, - "dataloader_type": "single", - "skip_getting_attention_mask_from_dataset": True, - } - if args.use_mock_data: - dataset_config = MockGPTDatasetConfig(**dataset_kwargs) + if args.sft: + # SFT-masked (Quantization-Aware) distillation via the container's Bridge FinetuningDatasetConfig + # -> NeMo-style GPTSFTDataset. `dataset_root` holds training.jsonl / validation.jsonl with + # {"input": , "output": } records. prompt_template="{input}{output}" tokenizes + # input+output verbatim (adjacent placeholders, no separator) matching the identity-formatted + # source; label_key="output" + answer_only_loss=True mask the loss to the assistant response only + # (answer_start_idx == len(context_ids)); truncation_field="input" truncates the context if needed. + dataset_config = FinetuningDatasetConfig( + seq_length=args.seq_length, + dataset_root=args.sft_dataset_root, + seed=args.seed, + dataloader_type="batch", + do_validation=True, + do_test=False, + dataset_kwargs={ + "prompt_template": "{input}{output}", + "label_key": "output", + "truncation_field": "input", + "answer_only_loss": True, + "add_bos": False, + "add_eos": True, + }, + ) else: - # Convert flat CLI list (e.g. ["1.0", "/path/data"]) to Megatron blend format - blend = get_blend_from_list(args.data_paths) - dataset_config = GPTDatasetConfig(blend=blend, split="99,1,0", **dataset_kwargs) + dataset_kwargs = { + "seq_length": args.seq_length, + "path_to_cache": args.data_path_to_cache, + "random_seed": args.seed, + "reset_attention_mask": False, + "reset_position_ids": False, + "eod_mask_loss": False, + "num_dataset_builder_threads": 1, + "data_sharding": True, + "dataloader_type": "single", + "skip_getting_attention_mask_from_dataset": True, + } + if args.use_mock_data: + dataset_config = MockGPTDatasetConfig(**dataset_kwargs) + else: + # Convert flat CLI list (e.g. ["1.0", "/path/data"]) to Megatron blend format + blend = get_blend_from_list(args.data_paths) + dataset_config = GPTDatasetConfig(blend=blend, split="99,1,0", **dataset_kwargs) # Assemble ConfigContainer and run distillation config = ConfigContainer( @@ -341,7 +458,9 @@ def _restore_student_hook(model_chunks): grad_reduce_in_fp32=True, overlap_grad_reduce=True, overlap_param_gather=True, - average_in_collective=True, + # 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", use_distributed_optimizer=True, ), dataset=dataset_config, @@ -354,22 +473,36 @@ def _restore_student_hook(model_chunks): wandb_entity=args.wandb_entity, # optional wandb_exp_name=args.wandb_exp_name, ), - tokenizer=TokenizerConfig( - tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + tokenizer=( + TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=args.student_hf_path, + hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, + ) + if args.sft + else TokenizerConfig( + tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + ) ), checkpoint=CheckpointConfig( save_interval=args.eval_interval, save=checkpoint_dir, load=checkpoint_dir, # Resume from this directory (if exists) - most_recent_k=5, # Keeps 5 most recent checkpoints (not metric-based) + most_recent_k=2, # Keeps 2 most recent checkpoints (each ~413GB here; fs1 near quota) ckpt_format="torch_dist", - async_save=True, + async_save=False, # sync save: async writer repeatedly corrupted iter-400 ckpt (inline_container) fully_parallel_save=True, ), rng=RNGConfig(seed=args.seed), mixed_precision="bf16_mixed", ) + # QAD with NVFP4 fake-quant makes the first optimizer step (many grad-accum microbatches) + # very slow; raise the NCCL process-group timeout above the default (~10 min) so the initial + # step does not trip the collective watchdog. Guarded in case the config field is renamed. + if hasattr(config, "dist") and hasattr(config.dist, "distributed_timeout_minutes"): + config.dist.distributed_timeout_minutes = 60 + print_rank_0("\nStarting distillation...") distill(config) print_rank_0( @@ -377,22 +510,7 @@ def _restore_student_hook(model_chunks): " in megatron distributed checkpoint format.\n" ) - if args.hf_export_path and is_vlm: - # Only the language model was distilled; export it back into the full VLM. - print_rank_0(f"Exporting distilled VLM to HF format to {args.hf_export_path}") - # ``distill`` tore down the model-parallel groups on exit, so rebuild them. - distill_provider.initialize_model_parallel(seed=args.seed) - full_student = distill_provider.full_model - # Strip the distillation wrapper -> plain trained language model (in place; reassign to be safe). - full_student.language_model = mtd.export(full_student.language_model) - save_vlm_to_hf( - full_student, - args.hf_export_path, - args.student_hf_path, - trust_remote_code=args.trust_remote_code, - ) - print_rank_0(f"Saved distilled VLM to {args.hf_export_path} in HF format") - elif args.hf_export_path: + if args.hf_export_path: print_rank_0(f"Exporting final distilled ckpt to HF format to {args.hf_export_path}") # Save rank before destroying process group (dist.rank() won't work after destruction) is_rank_0 = dist.rank() == 0 @@ -402,13 +520,20 @@ def _restore_student_hook(model_chunks): dist.cleanup() if is_rank_0: - export_llm_to_hf( + export_bridge = AutoBridge.from_hf_pretrained( + args.student_hf_model, trust_remote_code=args.trust_remote_code + ) + # Copy weights and remote code + export_bridge.export_ckpt( megatron_path=f"{checkpoint_dir}/iter_{args.train_iters:07d}", - hf_export_path=args.hf_export_path, - student_hf_path=args.student_hf_path, - template_hf=args.student_hf_model, - trust_remote_code=args.trust_remote_code, + hf_path=args.hf_export_path, + show_progress=True, + strict=True, ) + # Copy config.json from student_hf_path (handles both local paths and HF model IDs) + AutoConfig.from_pretrained( + args.student_hf_path, trust_remote_code=args.trust_remote_code + ).save_pretrained(args.hf_export_path) if __name__ == "__main__": diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 17db5e6da34..a0b2e8be407 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -36,6 +36,9 @@ """ import argparse +import yaml +import pathlib +import os import torch from megatron.bridge.models.hf_pretrained.utils import is_safe_repo @@ -44,6 +47,10 @@ import modelopt.torch.utils.distributed as dist from modelopt.torch.export import export_mcore_gpt_to_hf from modelopt.torch.utils import print_args, print_rank_0 +from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + StaticBlockScaleQuantizer, + TensorQuantizer, +) from modelopt.torch.utils.plugins.mbridge import ( load_mbridge_model_from_hf, load_modelopt_megatron_checkpoint, @@ -52,6 +59,13 @@ def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument( + "--grouped_experts", + action="store_true", + help="Build MoE experts grouped (GroupedMLP). Default is non-grouped, which per-block " + "NVFP4 checkpoints require. Set this only when the checkpoint was saved with grouped " + "experts; the layout must match or the weights will not load.", + ) parser.add_argument( "--hf_model_name_or_path", type=str, @@ -99,7 +113,39 @@ def get_args() -> argparse.Namespace: return args +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 + + def main(args: argparse.Namespace): + _ckpt_shape = _provider_overrides_from_checkpoint(args.megatron_path) trust_remote_code = is_safe_repo( trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path ) @@ -116,8 +162,11 @@ def main(args: argparse.Namespace): "num_layers_in_first_pipeline_stage": args.num_layers_in_first_pipeline_stage, "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage, "pipeline_dtype": torch.bfloat16, + "mtp_num_layers": _ckpt_shape["mtp_num_layers"], }, init_model_parallel=True, + # Default non-grouped, matching quantize.py; the layout must match the checkpoint. + moe_grouped_gemm=args.grouped_experts, load_weights=False, # The weights come from the Megatron checkpoint, so HF weights are not loaded ) @@ -127,6 +176,59 @@ def main(args: argparse.Namespace): load_modelopt_megatron_checkpoint(model, args.megatron_path) unwrapped_model = unwrap_model(model[0]) + # Static-NVFP4 export guard. + # + # An *enabled* NVFP4 weight quantizer that reaches the exporter without its calibrated scales + # means the values stored in the checkpoint were not restored. Exporting such a weight silently + # falls back to BF16: the result is larger than the recipe specifies and no longer matches it, + # with nothing in the logs to say so. Fail loudly instead. + # + # Static-block NVFP4 needs BOTH ``_amax`` (per block) and ``_global_amax`` (per tensor). A + # missing ``_global_amax`` slips past an ``_amax``-only check and then fails much later inside + # ``NVFP4QTensor.quantize``, where ``scale * scale_2`` broadcasts [N, 1] against [N] into an + # N x N allocation. Naming the attribute here turns that into an actionable message. + # + # Set MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1 to keep the previous behavior (disable the quantizer + # and emit BF16), which is then reported rather than silent. + uncalibrated: list[tuple[str, str]] = [] + for name, module in unwrapped_model.named_modules(): + # StaticBlockScaleQuantizer must be INCLUDED: `_global_amax` is defined on it, so excluding + # it would skip exactly the case this guard exists to catch. Only report enabled quantizers, + # matching the message. + if not isinstance(module, TensorQuantizer) or not getattr(module, "is_enabled", False): + 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 ( + isinstance(module, StaticBlockScaleQuantizer) + or 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() + else: + print_rank_0("All enabled NVFP4 weight quantizers have calibrated scales.") + # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is # correct regardless of pipeline placement / global rank ordering. diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 3454da00441..fb898c99933 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -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.", + ) + 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).", + ) 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 + ) + + 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 diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 7e81c21a462..e02d65eeef0 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -608,6 +608,52 @@ def _set_input_tensor(self, input_tensors: list[Tensor]): # HACK: Concatenate output tensors when PP>1 so they can be passed between ranks. def _forward(self, *args, **kwargs): + # Static-block NVFP4: promote the student's weight quantizers once, after the + # checkpoint amax/scales have been loaded, so the training forward takes the + # StaticBlockScaleQuantizer path rather than the generic FP8 (E4M3) path. Promotion + # cannot happen at build time because the scales only exist after the load. + # + # NOTE: in practice this converts exactly ONE module -- ``output_layer``. A measured run + # reports ``already promoted 460, converted 1, skipped 0``: every other quantizer is + # already a StaticBlockScaleQuantizer by the time training starts. So this is a workaround + # for output_layer being the one module the restore path does not promote (the same + # asymmetry behind its weight-quantizer scales not being restored). The better fix is to + # promote it on the normal path; until then, without this block the output projection + # would train through the generic FP8 path instead of static-block NVFP4. + if not getattr(self, "_modelopt_nvfp4_promoted", False): + from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + StaticBlockScaleQuantizer, + TensorQuantizer, + ) + + n_promoted = n_skipped = 0 + for name, module in self.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 + 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() + ) + n_promoted += 1 + if n_promoted or n_skipped: + logger.info( + f"Promoted {n_promoted} NVFP4 weight quantizer(s) to " + f"StaticBlockScaleQuantizer ({n_skipped} skipped)." + ) + self._modelopt_nvfp4_promoted = True with torch.no_grad(): self._teacher_model.eval() teacher_output = self._teacher_model(*args, **kwargs) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..2f742b87c5a 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -248,6 +248,22 @@ def _incompatible_method(self, *args, **kwargs): return _incompatible_method +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 + + def megatron_replace_quant_module_hook(model: torch.nn.Module): """Configure Megatron-Core model quantization support. @@ -261,6 +277,8 @@ def megatron_replace_quant_module_hook(model: torch.nn.Module): 3. For Attention modules, we configure them to use core_attention path for KV cache quantization. """ + untied = _resolve_output_layer_untied(model) + def _configure_attention_for_kv_cache_quant(module: Attention): """Configure Attention module for KV cache quantization compatibility.""" # Disable flash_decode if enabled - it bypasses core_attention (only called during inference) @@ -287,11 +305,19 @@ def _configure_attention_for_kv_cache_quant(module: Attention): def _register_extra_state_callbacks(model: torch.nn.Module): for name, module in model.named_modules(): if type(module) in QuantModuleRegistry: - # Skip output_layer w/o enabled weight_quantizer - if name.endswith("output_layer") and not getattr( - getattr(module, "weight_quantizer", None), "is_enabled", False - ): - continue + # Skip output_layer w/o enabled weight_quantizer. This hook also runs BEFORE + # QuantModule replacement (e.g. on restore), when ``weight_quantizer`` does not + # exist yet -- the old check then always skipped, so output_layer never received + # ModelOpt extra-state callbacks and its quantizer state (promotion to + # StaticBlockScaleQuantizer, ``_amax``, ``_global_amax``) was never restored. + # Fall back to the tying flag: an untied output_layer is quantizable. + if name.endswith("output_layer"): + _wq = getattr(module, "weight_quantizer", None) + _skip = ( + not getattr(_wq, "is_enabled", False) if _wq is not None else not untied + ) + if _skip: + continue register_modelopt_extra_state_callbacks( module, quant_module_get_extra_state, @@ -307,6 +333,10 @@ def _register_extra_state_callbacks(model: torch.nn.Module): if "vision_model" not in name: # We only enable hetereogenous_dist_checkpoint for language model, vision model is not quantized module.config.hetereogenous_dist_checkpoint = True + if untied is not None: + # Carried on the config so _MegatronParallelLinear.sharded_state_dict can read + # it without Megatron-LM global args (absent under Megatron-Bridge). + module.config.modelopt_output_layer_untied = untied _register_extra_state_callbacks(module) @@ -374,18 +404,63 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # output_layer.input_quantizer._amax but TP-only does not. This lead to # state_dict mismatch. if prefix.endswith("output_layer."): - try: - from megatron.training import get_args as _mlm_get_args - - _untied = bool( - getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) - ) - except Exception as e: - warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") - _untied = False + # Prefer the model-derived flag (set by megatron_replace_quant_module_hook); it is the + # only source available under Megatron-Bridge, which has no global args store. + _untied = getattr(self.config, "modelopt_output_layer_untied", None) + if _untied is None: + try: + from megatron.training import get_args as _mlm_get_args + + _untied = bool( + getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) + ) + except Exception as e: + warn_rank_0( + f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}" + ) + _untied = False if not _untied: return super().sharded_state_dict(prefix, sharded_offsets, metadata) + # Materialize missing weight-quantizer scale buffers so their keys appear in the load + # plan -- the dist-checkpoint loader SILENTLY SKIPS any checkpoint key the model does + # not advertise, which leaves output_layer uncalibrated and exports it as BF16. + # ``_amax`` must be allocated FLAT ``[numel // block, 1]``: that is the in-memory + # layout every other block-quantized layer uses, and ``_process_quantizer_amax`` below + # exposes it to the checkpoint as a ``[out_features, blocks]`` VIEW sharing the same + # storage, so the loader writes straight through. Allocating the viewed shape instead + # loads fine but leaves the wrong in-memory shape, which breaks the export scale math. + _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) + # `_process_quantizer_amax` later does `v.view(weight.shape[0], -1)`, which + # requires in_features (not just numel) to divide evenly by the block size. + if _block and self.weight.shape[-1] % 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), + ) + else: + # Leaving the buffers unallocated is the silent-drop failure this block exists + # to prevent, so say so rather than proceeding quietly. + warn_rank_0( + f"{prefix}weight_quantizer: cannot materialize scale buffers " + f"(block_size={_block}, in_features={self.weight.shape[-1]}); its " + "calibrated scales will not be restored from the checkpoint." + ) + quantizer_state_dict = {} for k, v in self.state_dict(prefix="", keep_vars=True).items(): if "_quantizer" in k and "_amax" in k: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index fd9b1e2f55e..e980383e730 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -679,7 +679,11 @@ def __len__(self): def _pack_documents_into_rows( - samples: list[str], tokenizer: "PreTrainedTokenizerBase", seq_length: int, num_rows: int + samples: list[str], + tokenizer: "PreTrainedTokenizerBase", + seq_length: int, + num_rows: int, + random_offset: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Global-stream document packing (Megatron-LM pretraining style). @@ -696,14 +700,23 @@ def _pack_documents_into_rows( eos_id = tokenizer.eos_token_id pad_id = tokenizer.pad_token_id has_eos_sep = eos_id is not None + # With random_offset (Megatron-LM --calib-use-random-offset), build one extra window + # of headroom, then drop a random number of leading tokens so the window grid shifts and + # calibration samples mid-document positions differently (relevant for long-context KV stats). + target_len = num_rows * seq_length + (seq_length if random_offset else 0) token_stream: list[int] = [] for s in samples: token_stream.extend(tokenizer.encode(s, add_special_tokens=False)) if has_eos_sep: token_stream.append(eos_id) - if len(token_stream) >= num_rows * seq_length: + if len(token_stream) >= target_len: break + 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):] + n_full = min(num_rows, len(token_stream) // seq_length) rows_ids: list[list[int]] = [ token_stream[i * seq_length : (i + 1) * seq_length] for i in range(n_full) @@ -749,6 +762,7 @@ def get_dataset_dataloader( include_labels: bool = False, apply_chat_template: bool = False, pack: bool = False, + random_offset: bool = False, distributed: bool = False, sampler_kwargs: dict | None = None, ) -> DataLoader: @@ -858,7 +872,7 @@ def get_dataset_dataloader( if pack: total_rows = sum(num_samples) input_ids, attention_mask = _pack_documents_into_rows( - all_samples, tokenizer, max_sample_length, total_rows + all_samples, tokenizer, max_sample_length, total_rows, random_offset=random_offset ) if input_ids.shape[0] < total_rows: warn_rank_0( diff --git a/modelopt/torch/utils/plugins/megatron_calibration.py b/modelopt/torch/utils/plugins/megatron_calibration.py index 4da38858209..a069f1505a4 100644 --- a/modelopt/torch/utils/plugins/megatron_calibration.py +++ b/modelopt/torch/utils/plugins/megatron_calibration.py @@ -50,6 +50,7 @@ def get_megatron_calibration_dataloader( device: torch.device | str | None = "cuda", apply_chat_template: bool = True, pack: bool = False, + random_offset: bool = False, ) -> torch.utils.data.DataLoader: """Build a DP-sharded calibration dataloader for Megatron-Core models. @@ -76,6 +77,7 @@ def get_megatron_calibration_dataloader( device=device, apply_chat_template=apply_chat_template, pack=pack, + random_offset=random_offset, distributed=dp_size > 1, sampler_kwargs={ "num_replicas": dp_size, @@ -95,6 +97,7 @@ def get_megatron_calibration_forward_loop( device: torch.device | str | None = "cuda", apply_chat_template: bool = True, pack: bool = False, + random_offset: bool = False, ) -> Callable[[torch.nn.Module], None]: """Build a Megatron-Core calibration ``forward_loop(model)``. @@ -116,6 +119,7 @@ def get_megatron_calibration_forward_loop( device=device, apply_chat_template=apply_chat_template, pack=pack, + random_offset=random_offset, ) def _forward_loop(model: torch.nn.Module) -> None: