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: