diff --git a/CHANGELOG.rst b/CHANGELOG.rst index eba77094fee..f20a33f1c17 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,12 @@ Changelog - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint. +*Megatron Framework (M-LM / M-Bridge)* + +- Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences. +- Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``. +- GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models. + *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 0d02ed08d24..333ccbbcb3d 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -178,6 +178,15 @@ def get_args(): parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size") parser.add_argument("--cp_size", type=int, default=1, help="Context parallel size") parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size") + parser.add_argument( + "--dsa_kernel_backend", + choices=["none", "tilelang", "cudnn"], + default=None, + help=( + "Fused kernels for DeepSeek sparse attention (e.g. GLM-5.x). The unfused fallback " + "builds dense [seq, seq] scores per head, which does not fit at long sequence lengths." + ), + ) # Dataset arguments parser.add_argument( @@ -238,6 +247,13 @@ def get_args(): "or the whole conversation. Multi-turn records train all of their responses under " "'assistant'.", ) + parser.add_argument( + "--sft_in_batch_packing", + action="store_true", + help="Pack each micro-batch's --sft_hf_dataset conversations into one THD sequence (attention " + "restarts per conversation) instead of padding every conversation to --seq_length. --mbs " + "then sets the conversations per pack.", + ) # Training & Eval arguments parser.add_argument( "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving" @@ -264,6 +280,13 @@ def get_args(): action="store_true", help="Keep the student's MTP head at its input weights: skip its loss and freeze it.", ) + parser.add_argument( + "--moe_aux_loss_coeff", + type=float, + default=None, + help="Override the MoE load-balancing loss weight (e.g. 0 to drop it). It is computed over every " + "token, so with --sft padding it also sends gradient into the padded positions.", + ) parser.add_argument( "--no_async_save", action="store_true", @@ -421,8 +444,10 @@ def get_args(): ) if args.eval_iters > 0 and not args.sft_hf_validation: raise ValueError("--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation.") - elif args.sft_hf_validation or args.sft_loss_mode != "assistant": - raise ValueError("--sft_hf_validation / --sft_loss_mode require --sft_hf_dataset.") + elif args.sft_hf_validation or args.sft_loss_mode != "assistant" or args.sft_in_batch_packing: + raise ValueError( + "--sft_hf_validation / --sft_loss_mode / --sft_in_batch_packing require --sft_hf_dataset." + ) if args.sft and (args.data_paths or args.use_mock_data): raise ValueError( "--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins " @@ -508,7 +533,11 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=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 + if args.dsa_kernel_backend is not None: + provider.dsa_kernel_backend = args.dsa_kernel_backend set_moe_expert_layout(provider, moe_grouped_gemm) + if args.moe_aux_loss_coeff is not None: + provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff if per_token_loss: provider.calculate_per_token_loss = True if args.recompute_granularity is not None: @@ -632,9 +661,14 @@ def _restore_student_hook(model_chunks): _hf_source(args.sft_hf_validation) if args.sft_hf_validation else None ), preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode), - dataloader_type="batch", + # "batch" collates a rank's whole global batch at once and then splits it into + # micro-batches, which a packed THD batch cannot be; packing collates per micro-batch. + dataloader_type="single" if args.sft_in_batch_packing else "batch", do_validation=args.eval_iters > 0, do_test=False, + enable_in_batch_packing=args.sft_in_batch_packing, + # Each packed conversation must split evenly across the 2 * CP context-parallel chunks. + in_batch_packing_pad_to_multiple_of=2 * args.cp_size if args.cp_size > 1 else 1, ) elif args.sft: # SFT-masked distillation via Bridge's FinetuningDatasetConfig -> NeMo-style GPTSFTDataset, diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index bb087e1c3ab..a035dda679d 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -23,7 +23,7 @@ The HuggingFace unified exporter does not gather tensor-parallel-sharded weights, so this script always loads the checkpoint at tensor_model_parallel_size=1 (re-sharding from whatever TP was used during quantization). Use --pp_size and/or --ep_size to shard a large model across GPUs for -export; --ep_size needs grouped-GEMM experts (not --no_moe_grouped_gemm). +export. Example usage to export an FP8 checkpoint produced by quantize.py: @@ -53,6 +53,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist from modelopt.torch.export import export_mcore_gpt_to_hf +from modelopt.torch.export.plugins.vllm_fakequant_megatron import export_mcore_gpt_to_hf_vllm_fq from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.mlflow import Tool, masked_args from modelopt.torch.utils.plugins.mbridge import ( @@ -106,6 +107,15 @@ def get_args() -> argparse.Namespace: "HuggingFace, in which case SequentialMLP is selected automatically." ), ) + parser.add_argument( + "--vllm_fakequant_export", + action="store_true", + help=( + "Fold the weight quantizers into BF16 weights, for accuracy evaluation in a stock vLLM " + "of formats it has no kernels for (e.g. GGML IQ2). Activation quantizers go to a " + "quantizer_state.pth sidecar." + ), + ) parser.add_argument( "--export_extra_modules", action="store_true", @@ -118,7 +128,7 @@ def get_args() -> argparse.Namespace: "--ep_size", type=int, default=1, - help="Expert parallel size (grouped-GEMM experts only); shards the experts across ranks", + help="Expert parallel size; shards the experts across ranks", ) parser.add_argument( "--num_layers_in_first_pipeline_stage", @@ -199,7 +209,10 @@ def main(args: argparse.Namespace): f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." ) # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet. - export_mcore_gpt_to_hf( + export_fn = ( + export_mcore_gpt_to_hf_vllm_fq if args.vllm_fakequant_export else export_mcore_gpt_to_hf + ) + export_fn( unwrapped_model, args.hf_model_name_or_path, export_extra_modules=export_extra_modules, diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index e12789cb3ef..44bcee9bf00 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -152,6 +152,15 @@ def get_args() -> argparse.Namespace: parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size") parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size") parser.add_argument("--cp_size", type=int, default=1, help="Context parallel size") + parser.add_argument( + "--dsa_kernel_backend", + choices=["none", "tilelang", "cudnn"], + default=None, + help=( + "Fused kernels for DeepSeek sparse attention (e.g. GLM-5.x). The unfused fallback " + "builds dense [seq, seq] scores per head, which does not fit at long sequence lengths." + ), + ) # Quantization arguments parser.add_argument( @@ -345,6 +354,11 @@ def main(args: argparse.Namespace): "pipeline_dtype": torch.bfloat16, "seq_length": args.seq_length, "gradient_accumulation_fusion": False, # not supported + **( + {"dsa_kernel_backend": args.dsa_kernel_backend} + if args.dsa_kernel_backend is not None + else {} + ), }, init_model_parallel=True, ) diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py index c34171cbd80..77c1ad75e1d 100644 --- a/modelopt/torch/export/plugins/mcore_custom.py +++ b/modelopt/torch/export/plugins/mcore_custom.py @@ -394,10 +394,15 @@ def save_safetensors_by_layer_index( } for layer_index in range(total_layers): meta_filename = name_template.format(layer_index + 1, total_layers) + ".json" - with open(save_directory + "/" + meta_filename) as f: - shard = json.load(f) - safetensor_index["metadata"]["total_size"] += shard["metadata"]["total_size"] - safetensor_index["weight_map"].update(shard["weight_map"]) + # Plus the expert shards that EP ranks other than 0 write beside the layer's own. + ep_shards = sorted( + Path(save_directory).glob(meta_filename[: -len(".json")] + "-ep*.json") + ) + for shard_file in [Path(save_directory) / meta_filename, *ep_shards]: + with open(shard_file) as f: + shard = json.load(f) + safetensor_index["metadata"]["total_size"] += shard["metadata"]["total_size"] + safetensor_index["weight_map"].update(shard["weight_map"]) with open(save_directory + "/model.safetensors.index.json", "w") as f: json.dump(safetensor_index, f, indent=4) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 2a4cc264dca..1355b833189 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -322,7 +322,15 @@ def save_pretrained( index: self._get_quantizer_state(state_dict) for index, state_dict in enumerate([*self.layer_state_dicts.values(), self._state_dict]) } + + # The weights are written in BF16, so a source checkpoint's own quantization (e.g. FP8) + # no longer describes them. + if hasattr(self._hf_config, "quantization_config"): + del self._hf_config.quantization_config + super().save_pretrained(save_directory, pretrained_model_name_or_path) + # The base exporter copies a pre-quantized source's hf_quant_config.json along with it. + (Path(save_directory) / "hf_quant_config.json").unlink(missing_ok=True) # Save fresh quantizer files after the base exporter copies source files. recipe = self._extract_quantizer_recipe_markers(quantizer_state_dicts) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 91fc68c9dd2..6f3f128d1fc 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -93,9 +93,11 @@ from megatron.core.models.multimodal.llava_model import LLaVAModel from megatron.core.parallel_state import ( get_data_parallel_rank, + get_expert_data_parallel_rank, get_expert_model_parallel_group, get_expert_model_parallel_rank, get_expert_model_parallel_world_size, + get_expert_tensor_parallel_rank, get_pipeline_model_parallel_rank, get_pipeline_model_parallel_world_size, get_tensor_model_parallel_rank, @@ -161,6 +163,8 @@ def __init__( self._state_dict = OrderedDict() self._layer_state_dicts = OrderedDict() + # Keys of the SequentialMLP experts held by this rank (all ranks' differ at EP>1). + self._local_expert_keys: set[str] = set() self._hf_pretrained_model_name = pretrained_model_name_or_path self._hf_config = transformers.AutoConfig.from_pretrained( pretrained_model_name_or_path, trust_remote_code=trust_remote_code @@ -346,24 +350,24 @@ def save_pretrained( raise NotImplementedError( "Megatron GGML unified export currently requires pipeline model parallel size 1" ) - # SequentialMLP rules index experts by local position, so EP>1 would collide. Agreed across - # ranks so that stages without such experts don't block in the collectives below. + # Packed local-expert rules hand over all of a rank's experts as one, so EP>1 would collide. + # Agreed across ranks so that stages without such experts don't block in the collectives below. if ( torch.distributed.is_initialized() and get_expert_model_parallel_world_size() > 1 + and self.rules.get("use_packed_local_experts", False) and self._any_rank(any(hasattr(m, "local_experts") for m in self.model.modules())) ): raise NotImplementedError( - "Export at expert parallel size > 1 needs grouped-GEMM experts; " - "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1." + "Export at expert parallel size > 1 does not support packed local experts; " + "export at EP=1." ) # One layer-shard writer per pipeline stage: other TP / DP / EP ranks hold the same layers # (EP>1 ranks hold no gathered experts at all), and writing them too would race on the same # files. Unlike is_writer_rank (one global writer of save_directory metadata), this has no # PP term on purpose: each stage owns a disjoint slice of the layer shards. - is_stage_layer_writer = ( - tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0 - ) + ep_rank = get_expert_model_parallel_rank() + is_stage_layer_writer = tp_rank == 0 and get_data_parallel_rank() == 0 and ep_rank == 0 # Main export process layer_state_dicts = self.layer_state_dicts @@ -486,6 +490,10 @@ def save_pretrained( # Bracket the writer's config.json read-modify-write with barriers so peers # never observe a truncated file (also ensures export_dir exists). torch.distributed.barrier() + if is_writer_rank: + # Expert shards left by an earlier export at a larger EP would enter the index. + for stale in Path(save_directory).glob("model-*-of-*-ep*.*"): + stale.unlink() config_json_file = save_directory + "/config.json" if is_writer_rank and self._hf_quant_config and os.path.exists(config_json_file): with open(config_json_file) as f: @@ -497,11 +505,30 @@ def save_pretrained( json.dump(config_dict, f, indent=4) torch.distributed.barrier() + # SequentialMLP experts are not gathered: each other EP rank writes its own to extra shards. + # Expert ownership follows the expert grid (EP x ETP x expert-DP), not the dense TP rank. + writes_own_experts = ( + get_expert_tensor_parallel_rank() == 0 + and get_expert_data_parallel_rank() == 0 + and ep_rank > 0 + and bool(self._local_expert_keys) + ) + name_template = "model-{:05d}-of-{:05d}" + if writes_own_experts: + # Layers without local experts (e.g. dense ones) would write empty shards. + layer_state_dicts = { + i: experts + for i, sd in layer_state_dicts.items() + if (experts := {k: v for k, v in sd.items() if k in self._local_expert_keys}) + } + name_template += f"-ep{ep_rank:03d}" save_safetensors_by_layer_index( - layer_state_dicts=layer_state_dicts if is_stage_layer_writer else {}, + layer_state_dicts=layer_state_dicts + if is_stage_layer_writer or writes_own_experts + else {}, total_layers=self.model.config.num_layers, save_directory=save_directory, - name_template="model-{:05d}-of-{:05d}", + name_template=name_template, ) # Every rank has written its shards; one rank now checks nothing was dropped. The result is @@ -785,13 +812,18 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): ) if hasattr(layer.mlp.experts, "local_experts"): if not self.rules.get("use_packed_local_experts", False): - for expert_id, expert in enumerate(layer.mlp.experts.local_experts): + keys_before = set(self._state_dict) + # Global ids: at EP>1 each rank holds a different slice of the experts. + for expert_id, expert in zip( + layer.mlp.local_expert_indices, layer.mlp.experts.local_experts + ): self.rules["local_experts.linear_fc1"]( expert.linear_fc1, layer_id, expert_id, is_mtp=is_mtp ) self.rules["local_experts.linear_fc2"]( expert.linear_fc2, layer_id, expert_id, is_mtp=is_mtp ) + self._local_expert_keys |= set(self._state_dict) - keys_before else: # For llama 4, in hf unified checkpoint, all local experts share one scale self.rules["local_experts.linear_fc1"]( diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index ddedd453983..83dfc2f00d5 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -154,8 +154,12 @@ def pinned_packed_weight(quantizer, format_name: str) -> torch.Tensor | None: return getattr(quantizer, _pin_name(format_name), None) +# Prefix of the buffer holding a pinned payload; checkpoint code keeps buffers with it. +PINNED_PAYLOAD_PREFIX = "_ggml_pinned_" + + def _pin_name(format_name: str) -> str: - return f"_ggml_pinned_{format_name}" + return f"{PINNED_PAYLOAD_PREFIX}{format_name}" def _matching_rows(table: torch.Tensor, rows: torch.Tensor) -> torch.Tensor | None: diff --git a/modelopt/torch/quantization/plugins/custom.py b/modelopt/torch/quantization/plugins/custom.py index 54dad627457..beeb63b2f97 100644 --- a/modelopt/torch/quantization/plugins/custom.py +++ b/modelopt/torch/quantization/plugins/custom.py @@ -24,6 +24,7 @@ from modelopt.torch.utils.distributed import ParallelState +from ..ggml.common import PINNED_PAYLOAD_PREFIX from ..nn import ( GroupedQuantizer, NVFP4StaticQuantizer, @@ -132,7 +133,10 @@ def modelopt_post_restore(self, prefix: str = ""): def _check_unsupported_states(quantizer: TensorQuantizer): for k in quantizer.state_dict(): - if k not in ["_amax", "_pre_quant_scale", "_global_amax"]: + # GGML payloads pinned by GPTQ are loaded with the weights from the sharded checkpoint. + if k not in ["_amax", "_pre_quant_scale", "_global_amax"] and not k.startswith( + PINNED_PAYLOAD_PREFIX + ): warnings.warn( f"Restore of {k} for {prefix} is not supported. The restore of this layer might be " f"incorrect. Please implement a custom restore for {k}." diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 5e422235e14..a96e7fd2a33 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -58,6 +58,7 @@ from ..algorithms import AutoQuantizeGradientSearcher from ..conversion import maybe_promote_nvfp4_static_quantizer +from ..ggml.common import PINNED_PAYLOAD_PREFIX from ..nn import ( GroupedQuantizer, QuantModule, @@ -564,7 +565,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): self._process_quantizer_amax(k, v, quantizer_state_dict) elif k == "input_quantizer._pre_quant_scale": self._process_activation_quantizer_pre_quant_scale(k, v, quantizer_state_dict) - elif self._parameter_to_keep_in_quantizer_state_dict(k): + elif self._parameter_to_keep_in_quantizer_state_dict(k) or PINNED_PAYLOAD_PREFIX in k: quantizer_state_dict[k] = v elif "quantizer" in k: warn_rank_0( @@ -616,7 +617,8 @@ def _get_shard_axis_dict(self, state_dict): continue if "weight_quantizer." in k: weight_quantizer_axis = self.get_submodule(k.rsplit(".", 1)[0]).axis - if weight_quantizer_axis is not None: + # GGML payloads are ``[out_features, in_features / block, block_bytes]``. + if weight_quantizer_axis is not None or PINNED_PAYLOAD_PREFIX in k: shard_axis_dict[k] = 0 return shard_axis_dict @@ -654,7 +656,7 @@ def _get_shard_axis_dict(self, state_dict): weight_quantizer_axis = self.weight_quantizer[0].axis elif "weight_quantizer.1." in k: weight_quantizer_axis = self.weight_quantizer[1].axis - if isinstance(weight_quantizer_axis, tuple): + if isinstance(weight_quantizer_axis, tuple) or PINNED_PAYLOAD_PREFIX in k: shard_axis_dict[k] = 1 if k == "input_quantizer._pre_quant_scale": shard_axis_dict[k] = 0 @@ -1190,6 +1192,11 @@ def _is_param_grad_enabled_for_megatron(pname: str, model: torch.nn.Module) -> b def get_mcore_layerwise_calibration_layers( model: torch.nn.Module, ) -> list[torch.nn.Module] | torch.nn.ModuleList | None: + # VLM wrappers (e.g. GLM-5.3-Flash) keep the decoder only on their language model. + if not hasattr(model, "decoder") and isinstance( + getattr(model, "language_model", None), MegatronModule + ): + model = model.language_model if not hasattr(model, "decoder") or not hasattr(model.decoder, "layers"): return None decoder_layers = model.decoder.layers diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 56e7554f522..7d9609063b6 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -348,10 +348,18 @@ def _patched_forward(self, *args, **kwargs): ) bind_forward_method(layer, _patched_forward, "_original_forward") + # A VLM wrapper and its language model are both patched, and the wrapper's forward calls the + # language model's. Only the outermost call may turn the stop into ``None``: if the inner one + # did, the wrapper would go on running its own code on a ``None`` output. + depth = [0] + def _early_stop_forward(module_self, *args, **kwargs): + depth[0] += 1 try: return module_self._original_forward(*args, **kwargs) except _EarlyStopForwardError: + if depth[0] > 1: + raise return None except RuntimeError as e: if "meta" not in str(e).lower(): @@ -363,14 +371,32 @@ def _early_stop_forward(module_self, *args, **kwargs): "decoder blocks. Use non-layerwise calibration for this " "architecture." ) from e + finally: + depth[0] -= 1 - bind_forward_method(self.model, _early_stop_forward, "_original_forward") + for root in self._early_stop_roots(): + bind_forward_method(root, _early_stop_forward, "_original_forward") except Exception: self._cleanup_layers() raise self._patched = True + def _early_stop_roots(self) -> list[nn.Module]: + """Modules whose forward ends the pass early. + + The model, plus the language model of a VLM wrapper, which text calibration drives directly. + """ + roots = [self.model] + lm = getattr(self.model, "language_model", None) + if ( + isinstance(lm, nn.Module) + and self._decoder_layers + and any(m is self._decoder_layers[0] for m in lm.modules()) + ): + roots.append(lm) + return roots + def _cleanup_layers(self): """Best-effort cleanup of any patched layers and model forward.""" if self._decoder_layers is not None: @@ -378,8 +404,9 @@ def _cleanup_layers(self): if isinstance(layer, _SkipLayer): self._decoder_layers[idx] = layer._original - if hasattr(self.model, "_original_forward"): - unpatch_forward_method(self.model, "_original_forward") + for root in self._early_stop_roots(): + if hasattr(root, "_original_forward"): + unpatch_forward_method(root, "_original_forward") if self._decoder_layers is not None: for layer in self._decoder_layers: diff --git a/modelopt/torch/utils/plugins/megatron_generate.py b/modelopt/torch/utils/plugins/megatron_generate.py index e19932d60e4..bc43a9c5370 100644 --- a/modelopt/torch/utils/plugins/megatron_generate.py +++ b/modelopt/torch/utils/plugins/megatron_generate.py @@ -261,6 +261,11 @@ def megatron_prefill( if isinstance(output, (tuple, list)): output = output[0] + # Layerwise calibration ends the forward early at the layer being captured (PP=1 only: with + # PP>1 the other stages would still wait on this one). + if output is None and skip_return_logits and not is_pp: + return None + # For PP non-last stages, forward activations to the next stage and return early. if is_pp and not pp_last: pp_dtype = model.config.pipeline_dtype or ( diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index e6c25dcb4a2..6977061d282 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -26,6 +26,7 @@ from _test_utils.torch.transformers_models import ( create_tiny_qwen3_5_vl_dir, create_tiny_qwen3_dir, + create_tiny_qwen3_moe_dir, get_tiny_tokenizer, ) from transformers import AutoModelForImageTextToText @@ -122,9 +123,14 @@ def test_distill_llm_sft(tmp_path, num_gpus): @pytest.mark.skipif(not HAS_DIRECT_HF_SFT, reason="Megatron-Bridge lacks DirectHFSFTDatasetConfig") -def test_distill_llm_sft_hf_chat(tmp_path, num_gpus): - """--sft_hf_dataset distills on chat jsonl rendered by the model's chat template.""" - teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) +@pytest.mark.parametrize("packed", [False, True]) +def test_distill_llm_sft_hf_chat(tmp_path, num_gpus, packed): + """--sft_hf_dataset distills on chat jsonl rendered by the model's chat template. + + ``packed`` runs an MoE student with conversations packed per micro-batch and the aux loss off. + """ + create_dir = create_tiny_qwen3_moe_dir if packed else create_tiny_qwen3_dir + teacher_hf_path = create_dir(tmp_path, with_tokenizer=True) train_iters = 2 gbs = 4 # More records than train_iters * gbs so the sampler does not run dry; every 4th is multi-turn. @@ -146,8 +152,10 @@ def test_distill_llm_sft_hf_chat(tmp_path, num_gpus): data_files[split].write_text("\n".join(json.dumps(r) for r in records) + "\n") distill_output_dir = tmp_path / "distill_output" + packing = {"moe_aux_loss_coeff": 0} if packed else {} distill_cmd_parts = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--sft"], + ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--sft"] + + (["--sft_in_batch_packing"] if packed else []), student_hf_path=teacher_hf_path, teacher_hf_path=teacher_hf_path, sft_hf_dataset=data_files["train"], @@ -156,7 +164,8 @@ def test_distill_llm_sft_hf_chat(tmp_path, num_gpus): tp_size=num_gpus, pp_size=1, seq_length=64, - mbs=1, + mbs=2 if packed else 1, + **packing, gbs=gbs, train_iters=train_iters, lr_warmup_iters=1, diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 2950234e16e..043970ef42a 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -354,12 +354,13 @@ def test_megatron_ggml_export_rejects_tensor_parallelism(qformat): exporter.save_pretrained("unused", "unused") -def test_megatron_export_rejects_sequential_experts_at_ep(): - """SequentialMLP experts are numbered by local position, so EP>1 export must refuse them.""" +def test_megatron_export_rejects_packed_local_experts_at_ep(): + """Packed local-expert rules write a rank's experts as one tensor, so EP>1 export must refuse.""" experts = torch.nn.Module() experts.local_experts = torch.nn.ModuleList() exporter = object.__new__(GPTModelExporter) exporter.model = torch.nn.Sequential(experts) + exporter.rules = {"use_packed_local_experts": True} with ( patch.object(exporter, "_is_sidecar_writer_rank", return_value=False), @@ -369,11 +370,60 @@ def test_megatron_export_rejects_sequential_experts_at_ep(): patch.object(uem, "get_pipeline_model_parallel_rank", return_value=0), patch.object(uem, "get_pipeline_model_parallel_world_size", return_value=1), patch.object(uem, "get_tensor_model_parallel_rank", return_value=0), - pytest.raises(NotImplementedError, match="grouped-GEMM experts"), + pytest.raises(NotImplementedError, match="packed local experts"), ): exporter.save_pretrained("unused", "unused") +def _test_sequential_experts_export_at_ep(model_dir, tmp_path, rank, size): + config = transformers.AutoConfig.from_pretrained(model_dir) + model = get_mcore_gpt_model( + expert_model_parallel_size=size, + initialize_megatron=True, + num_layers=config.num_hidden_layers, + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_query_groups=config.num_key_value_heads, + ffn_hidden_size=config.intermediate_size, + max_sequence_length=config.max_position_embeddings, + vocab_size=config.vocab_size, + normalization="RMSNorm", + kv_channels=config.hidden_size // config.num_attention_heads, + qk_layernorm=True, + num_moe_experts=config.num_experts, + moe_ffn_hidden_size=config.moe_intermediate_size, + moe_grouped_gemm=False, + ).cuda() + export_dir = tmp_path / "export" + stale = export_dir / "model-00001-of-00002-ep007.json" # left by an export at a larger EP + if rank == 0: + export_dir.mkdir() + stale.write_text('{"metadata": {"total_size": 0}, "weight_map": {}}') + export_mcore_gpt_to_hf(model, model_dir, dtype=torch.bfloat16, export_dir=str(export_dir)) + assert not stale.exists() + + # Each rank's experts must land under their global ids, from whichever shard holds them. + weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"] + for layer_id, layer in enumerate(model.decoder.layers): + for expert_id, expert in zip( + layer.mlp.local_expert_indices, layer.mlp.experts.local_experts + ): + name = f"model.layers.{layer_id}.mlp.experts.{expert_id}.down_proj.weight" + with safe_open(str(export_dir / weight_map[name]), framework="pt") as f: + exported = f.get_tensor(name) + assert torch.equal(exported, expert.linear_fc2.weight.detach().bfloat16().cpu()), name + if rank == 0: + assert_exported_checkpoint_matches( + export_dir, model_dir, check_values=False, allow_unexpected=("mlp.gate.expert_bias",) + ) + + +def test_sequential_experts_export_at_ep(dist_workers_size_2, tmp_path): + """SequentialMLP experts exported at EP=2 keep every expert under its global id.""" + model_dir = create_tiny_qwen3_moe_dir(tmp_path) + dist_workers_size_2.run(partial(_test_sequential_experts_export_at_ep, model_dir, tmp_path)) + + def _test_mla_export_keeps_hf_head_dim(model_dir, rank, size): model = get_mcore_gpt_model( tensor_model_parallel_size=size, diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py index 20d02e7537a..8e95163a8ff 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_megatron_export.py @@ -182,6 +182,13 @@ def forward_loop(model): (source / "quant_recipe.yaml").write_text( yaml.safe_dump({stale_quantizer: {"_disabled": True}}) ) + # A pre-quantized source (e.g. FP8) describes weights the BF16 export no longer has. + config = json.loads((source / "config.json").read_text()) + config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128]} + (source / "config.json").write_text(json.dumps(config)) + (source / "hf_quant_config.json").write_text( + json.dumps({"quantization": {"quant_algo": "FP8"}}) + ) export_dir = tmp_path / "vllm_export" exporter = VllmFqGPTModelExporter(model, source, dtype=torch.bfloat16) @@ -241,6 +248,7 @@ def forward_loop(model): assert stale_quantizer + "._amax" not in state assert stale_quantizer not in recipe assert {"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"} <= weight_map.keys() + assert "quantization_config" not in json.loads((export_dir / "config.json").read_text()) model.share_embeddings_and_output_weights = True for quantizer in model.output_layer.modules(): diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 177ce61fc4b..9e2a7be8a0b 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -67,12 +67,14 @@ import modelopt import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( restore_sharded_modelopt_state, save_sharded_modelopt_state, ) from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher +from modelopt.torch.quantization.ggml.common import PINNED_PAYLOAD_PREFIX from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.plugins.megatron import ( @@ -99,6 +101,7 @@ from modelopt.torch.quantization.utils import is_quantized_linear from modelopt.torch.quantization.utils.calib_utils import GPTQHelper from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill try: from megatron.core.extensions.transformer_engine import TERowParallelLinear @@ -1688,6 +1691,76 @@ def test_mcore_layerwise_calibration_layers_do_not_mutate_decoder(dist_workers_s dist_workers_size_1.run(_test_mcore_layerwise_calibration_layers_do_not_mutate_decoder) +class _VLMWrapper(MegatronModule): + """A VLM root that, like Megatron-Bridge's GLM-5.3-Flash model, has no ``decoder`` of its own.""" + + def __init__(self, language_model): + super().__init__(config=language_model.config) + self.language_model = language_model + + def forward(self, *args, **kwargs): + # Post-processes the language model's output, as real wrappers do. + return self.language_model(*args, **kwargs).float() + + +def _test_layerwise_calibrate_vlm_text_forward_loop(drive_wrapper, rank, size): + initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED) + language_model = _gpt_model_provider(tp_size=1, hidden_size=64) + model = _VLMWrapper(language_model) + tokens = torch.randint(0, language_model.vocab_size, (2, 16), device="cuda") + + # Text calibration in examples/megatron_bridge/quantize.py drives the inner language model; image-text + # calibration drives the wrapper, so both patched forwards are on the stack when the stop is raised. + def forward_loop(_model): + megatron_prefill( + model if drive_wrapper else language_model, tokens, skip_return_logits=True + ) + + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + mtq.quantize(model, config, forward_loop) + + calibrated = [m for m in language_model.decoder.modules() if is_quantized_linear(m)] + assert calibrated + assert all(m.input_quantizer.amax is not None for m in calibrated) + + +@pytest.mark.parametrize("drive_wrapper", [False, True]) +def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper): + dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper)) + + +def _test_ggml_gptq_payload_sharded_state_dict(tmp_path, rank, size): + initialize_for_megatron(tensor_model_parallel_size=size, seed=SEED) + model_ref = _gpt_model_provider(tp_size=size) + model_test = _gpt_model_provider(tp_size=size) + + config = load_recipe("general/ptq/iq2_xs").quantize.model_dump() + gptq = next(algo for algo in config["algorithm"] if algo["method"] == "gptq") + config["algorithm"] = {**gptq, "layerwise": {"enable": False}} + mtq.quantize(model_ref, config, get_forward(model_ref)) + pinned = {k: v for k, v in model_ref.state_dict().items() if PINNED_PAYLOAD_PREFIX in k} + assert pinned, "GPTQ pinned no GGML payload" + + save_distributed_checkpoint(tmp_path, model_ref) + save_sharded_modelopt_state([model_ref], tmp_path) + restore_sharded_modelopt_state([model_test], tmp_path) + model_test = load_distributed_checkpoint(tmp_path, model_test) + + restored = model_test.state_dict() + matches = torch.tensor( + all(k in restored and torch.equal(v, restored[k]) for k, v in pinned.items()), device="cuda" + ) + # Agree across ranks before asserting, so a mismatch on one rank fails instead of hanging. + torch.distributed.all_reduce(matches, op=torch.distributed.ReduceOp.MIN) + assert matches.item(), "GPTQ's pinned GGML payloads did not survive the checkpoint" + + +def test_ggml_gptq_payload_sharded_state_dict(dist_workers, tmp_path): + """GPTQ's pinned GGML codes are saved, so a reload exports them instead of re-encoding.""" + dist_workers.run(partial(_test_ggml_gptq_payload_sharded_state_dict, tmp_path)) + + @pytest.mark.parametrize("ep_size", [1, 2]) @pytest.mark.parametrize("sync_weight_amax", [True, False]) def test_layer_sync_moe_local_experts_amax(dist_workers, ep_size, sync_weight_amax):