From 78fd487bbe13d53fd0119ec45eee4810bb4bf6b9 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:00:50 -0700 Subject: [PATCH 1/8] feat: VLM distillation and HF export for Megatron-Bridge examples - distill.py: distill only the language_model submodule of a VLM (vision tower / projector untouched); reuses the LLM training path for everything else. - Add export_distilled_megatron_to_hf.py to convert a distilled Megatron checkpoint (any iteration) to HF -- LLM/Puzzletron via export_ckpt, VLM via in-memory reassembly (vision from HF + distilled LM from the checkpoint). - Rename export.py -> export_quantized_megatron_to_hf.py; distill.py's final export reuses the shared export helpers. - mbridge: load_modelopt_megatron_checkpoint gains restore_modelopt_state (weights-only reload of a full-precision distilled student, no teacher). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 3 +- examples/megatron_bridge/README.md | 48 ++-- .../megatron_bridge/_distillation_provider.py | 97 ++++++++ examples/megatron_bridge/distill.py | 194 +++++---------- .../export_distilled_megatron_to_hf.py | 222 ++++++++++++++++++ ....py => export_quantized_megatron_to_hf.py} | 2 +- examples/megatron_bridge/quantize.py | 11 +- .../README.md | 8 +- examples/megatron_bridge/tutorials/README.md | 2 +- modelopt/torch/utils/plugins/mbridge.py | 53 +---- .../examples/megatron_bridge/test_distill.py | 93 +++++++- .../megatron_bridge/test_prune_minitron.py | 16 ++ tests/examples/megatron_bridge/test_qad.py | 96 ++++++-- .../megatron_bridge/test_quantize_export.py | 45 +--- 14 files changed, 622 insertions(+), 268 deletions(-) create mode 100644 examples/megatron_bridge/_distillation_provider.py create mode 100644 examples/megatron_bridge/export_distilled_megatron_to_hf.py rename examples/megatron_bridge/{export.py => export_quantized_megatron_to_hf.py} (99%) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..3d4506f2a83 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -34,6 +34,7 @@ Changelog - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned. - Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations. - Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported. +- Add Megatron-Bridge distillation support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. - Add NVFP4 Four-Over-Six (4/6) weight quantization (``mtq.NVFP4_FOUR_OVER_SIX_CFG``): MSE weight calibration picks, per block, between an M=6 and an M=4 dynamic range (the choice is folded into the FP8 per-block scales), with the ``four_over_six: true`` flag normalizing those scales by 256 (vs 448) for M=4 headroom. Supported via ``mtq.quantize`` and HF / Megatron export only -- **not** ``mtq.compress``, which does not preserve the per-block M=4/M=6 choice - Add dLLM (tied-weight PTQ and HF-checkpoint export) support for diffusion-based encoder-decoder LLMs (e.g. DiffusionGemma) whose encoder/decoder stacks share parameters via HF ``_tied_weights_keys``. @@ -81,7 +82,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* -- Add quantization examples for the Megatron-Bridge framework (``examples/megatron_bridge/``): post-training quantization (`quantize.py `_ calibrates an HF model via ``--quant_cfg`` alias / full config name or a ``--recipe`` YAML, with optional KV-cache quant, weight-only, compression, and MoE expert-ratio calibration, and saves a Megatron checkpoint with tensor / pipeline / expert parallelism), export to a deployable HuggingFace (unified) checkpoint for TensorRT-LLM / vLLM / SGLang (`export.py `_), and Quantization Aware Distillation (extend existing `distill.py `_). See `examples/megatron_bridge/README.md `_ for details. +- Add quantization examples for the Megatron-Bridge framework (``examples/megatron_bridge/``): post-training quantization (`quantize.py `_ calibrates an HF model via ``--quant_cfg`` alias / full config name or a ``--recipe`` YAML, with optional KV-cache quant, weight-only, compression, and MoE expert-ratio calibration, and saves a Megatron checkpoint with tensor / pipeline / expert parallelism), export to a deployable HuggingFace (unified) checkpoint for TensorRT-LLM / vLLM / SGLang (`export_quantized_megatron_to_hf.py `_), and Quantization Aware Distillation (extend existing `distill.py `_). See `examples/megatron_bridge/README.md `_ for details. - Add Megatron Core export/import mapping for Qwen3-VL (``Qwen3VLForConditionalGeneration``) vision-language models. The mapping handles the ``model.language_model.`` weight prefix used by Qwen3-VL. - Add shared Megatron-Core calibration forward loop: ``modelopt.torch.utils.plugins.megatron_calibration.get_megatron_calibration_forward_loop`` produces the ``forward_loop`` callable expected by ``mtq.quantize`` / ``mtp.prune``. Replaces the bespoke calibration loops in Megatron-LM and Megatron-Bridge for quantization and pruning with a single canonical implementation. - Support Megatron-Core checkpoint restore and export for MSE ``NVFP4StaticQuantizer``. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 60dfed7fc7d..7075cc9f5a2 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -47,6 +47,13 @@ docker run \ > [!WARNING] > Use `python -m pip` instead of `pip` to avoid conflicts with the system-wide installed packages in the NeMo containers. You may also refer to this [doc](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docker/common/README.md#installing-packages-inside-the-container) on how to correctly install packages in the NeMo containers without breaking existing torch installation. +> [!NOTE] +> **Working with MoE Vision-Language Models (e.g. Qwen3.5-VL-MoE)?** The `nemo:26.06` container's Megatron-Bridge lacks the MoE expert weight mappings these models need (dense VLMs such as Gemma3-VL and Qwen3-VL work as-is). Until the `nemo:26.08` container is released, mount the latest [Megatron-Bridge `main`](https://github.com/NVIDIA-NeMo/Megatron-Bridge) source over the pre-installed copy by adding this to the `docker run` command above: +> +> ```bash +> -v /path/to/Megatron-Bridge/src/megatron/bridge:/opt/Megatron-Bridge/src/megatron/bridge +> ``` + You also need to login with your HuggingFace token to download gated datasets / models. Note that the default dataset for pruning and quantization is [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2), which is gated. @@ -59,7 +66,7 @@ hf auth login --token This section shows how to quantize a HuggingFace model using ModelOpt in the Megatron-Bridge framework. Quantization is a two-step flow: 1. [quantize.py](quantize.py) applies post-training quantization (PTQ) with calibration and saves a **Megatron checkpoint** (with ModelOpt state). Tensor / pipeline / expert parallelism are all supported, and the checkpoint can be reloaded for further training (Quantization Aware Training / Quantization Aware Distillation). -2. [export.py](export.py) converts that Megatron checkpoint to a **HuggingFace (unified) checkpoint** that deploys directly with TensorRT-LLM, vLLM, or SGLang. +2. [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) converts that Megatron checkpoint to a **HuggingFace (unified) checkpoint** that deploys directly with TensorRT-LLM, vLLM, or SGLang. `quantize.py` supports the following formats via `--quant_cfg` (e.g. `fp8`, `nvfp4`, `int8_sq`, `int4_awq`, `w4a8_awq`, ...). You can also pass any full config name exposed by ModelOpt (e.g. `NVFP4_DEFAULT_CFG`) or a YAML `--recipe` (e.g. `general/ptq/nvfp4_default-kv_fp8`, authoritative for quant_cfg + algorithm + KV-cache). KV-cache quantization can be enabled on top via `--kv_cache_quant` (e.g. `fp8`, `nvfp4`). @@ -81,7 +88,7 @@ torchrun --nproc_per_node 2 quantize.py \ **Step 2 — export** the Megatron checkpoint to a deployable HuggingFace checkpoint: ```bash -torchrun --nproc_per_node 2 export.py \ +torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \ --hf_model_name_or_path Qwen/Qwen3-8B \ --megatron_path /tmp/Qwen3-8B-NVFP4-megatron \ --pp_size 2 \ @@ -89,12 +96,12 @@ torchrun --nproc_per_node 2 export.py \ ``` > [!NOTE] -> The HuggingFace unified exporter does not gather tensor-parallel-sharded weights. Use `--pp_size` on `export.py` to shard a large model with pipeline parallelism across GPUs for export. +> The HuggingFace unified exporter does not gather tensor-parallel-sharded weights. Use `--pp_size` on `export_quantized_megatron_to_hf.py` to shard a large model with pipeline parallelism across GPUs for export. > [!TIP] > To recover the accuracy lost during quantization, fine-tune the quantized Megatron checkpoint (from step 1) with [Quantization Aware Distillation (QAD)](#quantization-aware-distillation-qad) before running the step 2 export. -To see the full usage for advanced configurations, run `torchrun --nproc_per_node 1 quantize.py --help` (or `export.py --help`). +To see the full usage for advanced configurations, run `torchrun --nproc_per_node 1 quantize.py --help` (or `export_quantized_megatron_to_hf.py --help`). ### Vision-Language Models (VLMs) @@ -104,7 +111,7 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati - A **text** dataset runs text-only calibration of the language model (vision tower idle). > [!NOTE] -> HuggingFace unified export (`export.py`) of a quantized VLM is not yet supported; the quantized VLM is saved in Megatron checkpoint format only. +> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM is not yet supported; the quantized VLM is saved in Megatron checkpoint format only. ## Sanity-Check Generation @@ -204,7 +211,7 @@ torchrun --nproc_per_node 8 distill.py \ --output_dir /output/qwen3_8b_nvfp4_qad ``` -The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export.py](export.py) (point `--megatron_path` at `/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). +The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). ### Slurm Usage @@ -214,27 +221,36 @@ To run the distillation script on a Slurm cluster for multi-node training, you j The distilled checkpoint is saved in Megatron distributed format. If you need a HuggingFace checkpoint, there are two ways to convert it: -**Inline** -- add `--hf_export_path` and `--student_hf_model` to the `distill.py` command to automatically convert the final checkpoint after distillation: +**Inline** -- add `--hf_export_path` to the `distill.py` command to automatically convert the **final** checkpoint after distillation: ```bash torchrun --nnodes 1 --nproc_per_node 8 distill.py \ ... \ - --hf_export_path /path/to/save/distilled_hf_ckpt \ - --student_hf_model Qwen/Qwen3-4B + --hf_export_path /path/to/save/distilled_hf_ckpt ``` -`--student_hf_model` should match the base architecture of the student (used as a template for export). For non-Puzzletron (i.e. standard) models, it should be same as `--student_hf_path`. +`--student_hf_model` is also needed for heterogeneous (Puzzletron/NAS) students whose architecture differs from `--student_hf_path`; for standard (homogeneous) models it defaults to `--student_hf_path`, and it is unused for VLMs. -**Separate conversion** -- convert any saved iteration using the Megatron-Bridge conversion script: +**Separate conversion** -- convert **any** saved iteration (intermediate or final) with [export_distilled_megatron_to_hf.py](export_distilled_megatron_to_hf.py): ```bash -uv run python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ - --hf-model \ - --megatron-path /checkpoints/iter_ \ - --hf-path +torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ + --student_hf_path \ + --megatron_path /checkpoints/iter_ \ + --hf_export_path /path/to/save/distilled_hf_ckpt ``` -For more details, see the [Megatron-Bridge conversion README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion). +### Vision-Language Models (VLMs) + +For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `distill.py` automatically distills only the **language model** (on text data) and leaves the vision tower and vision-language projector untouched, then saves the full VLM back. This also aligns with pruning and quantization behavior where we leave the vision tower and projector untouched. This composes with pruning and QAD (`--student_megatron_path`) exactly as for LLMs, and the HF export reuses `--student_hf_path` (no `--student_hf_model` needed). + +```bash +torchrun --nproc_per_node 8 distill.py \ + --tp_size 8 \ + --teacher_hf_path Qwen/Qwen3-VL-2B-Thinking \ + --student_hf_path Qwen/Qwen3-VL-2B-Thinking \ + ... +``` ### Distillation Results diff --git a/examples/megatron_bridge/_distillation_provider.py b/examples/megatron_bridge/_distillation_provider.py new file mode 100644 index 00000000000..b78d256eda2 --- /dev/null +++ b/examples/megatron_bridge/_distillation_provider.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Minimal extension of Megatron-Bridge's ``DistillationProvider``. + +Adds two things over the stock provider: (1) the KD conversion runs in a post-weight-load pre-wrap hook +instead of in ``provide()``, and (2) a ``distill_submodule`` option to distill only a submodule (e.g. a +VLM ``language_model``, leaving the vision tower / projector untouched). This is the same behavior as the +upstream change in NVIDIA-NeMo/Megatron-Bridge and is implemented here as a small delta (via a dynamic +subclass, without mutating the stock class) so the example works on the current container. + +TODO: Remove this module and import ``convert_to_distillation_provider`` directly from +``megatron.bridge.models.distillation_provider`` once we require the nemo:26.08 container (Megatron-Bridge#4707). +""" + +from megatron.bridge.models.distillation_provider import ( + convert_to_distillation_provider as _base_convert_to_distillation_provider, +) +from megatron.core.utils import unwrap_model + +import modelopt.torch.distill as mtd +import modelopt.torch.distill.plugins.megatron as mtd_mcore + + +def _provide(self, pre_process=None, post_process=None, vp_stage=None): + """Build the un-converted student; the KD conversion is deferred to ``_convert_hook``.""" + if vp_stage is not None: + raise ValueError("ModelOpt KD currently does not support virtual-pipeline parallel.") + return self._super_class.provide(self, pre_process, post_process, vp_stage) + + +def _convert_hook(self, model_chunks): + """Pre-wrap hook (runs after weight-load): distill the whole model or ``distill_submodule``.""" + assert len(model_chunks) == 1, "ModelOpt KD does not support virtual pipeline (>1 model chunk)." + student = unwrap_model(model_chunks[0]) + # Hack to get teacher's pre-wrap hooks called to potentially load HF weights + teacher = unwrap_model( + self.teacher.provide_distributed_model(wrap_with_ddp=False, mixed_precision_wrapper=None)[0] + ) + if self.distill_submodule is not None: + # Retain the full model so the (in-place) distilled submodule can be exported back within it. + self.full_model = student + student = getattr(student, self.distill_submodule) + teacher = getattr(teacher, self.distill_submodule) + + kd_cfg = mtd_mcore.setup_distillation_config(self.kd_config, student.config, teacher.config) + modelopt_cfg = { + "teacher_model": teacher, + "criterion": kd_cfg.criterion, + "loss_balancer": kd_cfg.loss_balancer, + } + kd_model = mtd.convert(student, mode=[("kd_loss", modelopt_cfg)]) + mtd_mcore.adjust_distillation_model_for_mcore(kd_model, kd_cfg) + return [kd_model] + + +def convert_to_distillation_provider( + student_provider, teacher_provider, kd_config=None, *, distill_submodule=None +): + """Like ``megatron.bridge``'s ``convert_to_distillation_provider`` but defers the KD conversion to a + pre-wrap hook (so the student is weight-loaded first) and can target a submodule. See module docstring. + """ + provider = _base_convert_to_distillation_provider(student_provider, teacher_provider, kd_config) + # Dynamically subclass the (already rebased) provider class to add the deferred-convert behavior + # without mutating Megatron-Bridge's DistillationProvider. isinstance(provider, DistillationProvider) + # stays True, so megatron.bridge.training.distill.distill() still accepts it. + submodule_cls = type( + "SubmoduleDistillationProvider", + (type(provider),), + { + "provide": _provide, + "_convert_hook": _convert_hook, + "distill_submodule": distill_submodule, + }, + ) + # Use object.__setattr__ to bypass DistillationProvider.__setattr__, which mirrors every attribute + # set onto the teacher -- assigning ``__class__`` normally would also switch the teacher's class. + object.__setattr__(provider, "__class__", submodule_cls) + # Append the convert hook after the bridge's weight-load hook so the student is fully weight-loaded + # before conversion. Set _pre_wrap_hooks via object.__setattr__ (not register_pre_wrap_hook) to + # bypass the teacher-mirroring __setattr__: when the student starts with no hooks (QAD builds it + # with load_weights=False), the mirror would share the hook list with the teacher, so building the + # teacher inside _convert_hook would re-run _convert_hook -> infinite recursion. + hooks = [*getattr(provider, "_pre_wrap_hooks", []), provider._convert_hook] + object.__setattr__(provider, "_pre_wrap_hooks", hooks) + return provider diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 283504f1ec2..bbb8a26731a 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -23,14 +23,11 @@ import argparse import contextlib import os -from dataclasses import fields 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, ) @@ -49,116 +46,18 @@ 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 +from modelopt.torch.utils import print_args, print_rank_0, warn_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 -def _patched_to_cfg_dict(self): - """Patched DistillationProvider.to_cfg_dict method for heterogeneous teacher and student models. - - TODO: Remove once we drop nemo:26.02 container support - """ - from megatron.bridge.training.utils.config_utils import _ConfigContainerBase - - result = {"_target_": f"{self._super_class.__module__}.{self._super_class.__qualname__}"} - # Use fields from the actual student provider class, not DistillationProvider. - # DistillationProvider's __dataclass_fields__ only includes TransformerConfig fields - # (set at class definition time), missing GPTModelProvider-level fields like - # vocab_size, share_embeddings_and_output_weights, etc. - excluded_fields = {"teacher", "kd_config"} - for field in fields(self._super_class): - if field.name.startswith("_") or field.name in excluded_fields: - continue - if hasattr(self, field.name): - result[field.name] = _ConfigContainerBase._convert_value_to_dict( - getattr(self, field.name) - ) - for field in fields(self): - if field.name.startswith("_") or field.name in excluded_fields: - continue - if field.name not in result: - result[field.name] = _ConfigContainerBase._convert_value_to_dict( - getattr(self, field.name) - ) - return result - - -DistillationProvider.to_cfg_dict = _patched_to_cfg_dict - - -# 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.") @@ -289,8 +188,9 @@ def get_args(): type=str, required=False, default=None, - 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.", + help="Export template: HF model whose architecture matches the exported student. Only needed " + "for heterogeneous (Puzzletron/NAS) students whose architecture differs from --student_hf_path. " + "Defaults to --student_hf_path; unused for VLMs.", ) args = parser.parse_args() @@ -298,8 +198,8 @@ def get_args(): if not args.use_mock_data and not args.data_paths: raise ValueError("Must provide either --data_paths or set --use_mock_data.") - if args.hf_export_path and not args.student_hf_model: - raise ValueError("Must provide --student_hf_model if --hf_export_path is provided.") + if args.student_hf_model is None: + args.student_hf_model = args.student_hf_path print_args(args) @@ -310,6 +210,11 @@ def main(args: argparse.Namespace): checkpoint_dir = os.path.join(args.output_dir, "checkpoints") tensorboard_dir = os.path.join(args.output_dir, "tb_logs") + is_vlm = hasattr( + AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), + "vision_config", + ) + # Build student and teacher model providers def _build_model_provider(hf_path, load_weights=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) @@ -317,7 +222,8 @@ def _build_model_provider(hf_path, load_weights=True): # Override parallelism / training settings provider.tensor_model_parallel_size = args.tp_size - provider.sequence_parallel = args.tp_size > 1 + # TODO: re-enable VLM SP from the nemo:26.08 container (fixed upstream by Megatron-LM#5628) + provider.sequence_parallel = args.tp_size > 1 and not is_vlm provider.pipeline_model_parallel_size = args.pp_size provider.pipeline_dtype = torch.bfloat16 provider.context_parallel_size = args.cp_size @@ -347,23 +253,41 @@ def _build_model_provider(hf_path, load_weights=True): student_provider.gradient_accumulation_fusion = False teacher_provider = _build_model_provider(args.teacher_hf_path) - # Wrap into DistillationProvider kd_config = ModelOptDistillConfig( skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) - distill_provider = convert_to_distillation_provider( - student_provider, teacher_provider, kd_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" + ) + run_provider = convert_to_distillation_provider( + student_provider, + teacher_provider, + kd_config, + distill_submodule="language_model" if is_vlm else None, ) 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)." ) - # 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 + + 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 + + run_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) # Build optimizer and scheduler optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing( @@ -395,7 +319,7 @@ def _build_model_provider(hf_path, load_weights=True): # Assemble ConfigContainer and run distillation config = ConfigContainer( - model=distill_provider, + model=run_provider, train=TrainingConfig( train_iters=args.train_iters, eval_interval=args.eval_interval, @@ -428,7 +352,7 @@ def _build_model_provider(hf_path, load_weights=True): wandb_exp_name=args.wandb_exp_name, ), tokenizer=TokenizerConfig( - tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + tokenizer_type="NullTokenizer", vocab_size=run_provider.vocab_size ), checkpoint=CheckpointConfig( save_interval=args.eval_interval, @@ -450,7 +374,22 @@ def _build_model_provider(hf_path, load_weights=True): " in megatron distributed checkpoint format.\n" ) - if args.hf_export_path: + 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. + run_provider.initialize_model_parallel(seed=args.seed) + full_student = run_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_model, + 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: 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 @@ -460,20 +399,13 @@ def _build_model_provider(hf_path, load_weights=True): dist.cleanup() if is_rank_0: - 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( + export_llm_to_hf( megatron_path=f"{checkpoint_dir}/iter_{args.train_iters:07d}", - hf_path=args.hf_export_path, - show_progress=True, - strict=True, + 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, ) - # 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_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py new file mode 100644 index 00000000000..ffa41825463 --- /dev/null +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert a full-precision distilled Megatron checkpoint (produced by distill.py) to HuggingFace. + +Two mechanisms, dispatched on the model type: + + - LLM (Homogeneous or Puzzletron Heterogeneous): the full model is on disk, so it is exported directly with + ``AutoBridge.export_ckpt`` (which reads the checkpoint's actual per-layer shapes and therefore + handles both homogeneous and heterogeneous students). + - VLM: only the ``language_model`` submodule is distilled and checkpointed, so the full VLM is + reassembled in memory -- vision tower + projector from the original HF model (--student_hf_path), + the distilled language model from the checkpoint -- and written with ``AutoBridge.save_hf_weights``. + +These two helpers (``export_llm_to_hf`` / ``save_vlm_to_hf``) are also reused by distill.py for its +final-checkpoint export. + +Example LLM (Homogeneous or Puzzletron Heterogeneous): + + torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ + --student_hf_path Qwen/Qwen3-0.6B \ + --megatron_path /tmp/distill-out/checkpoints/iter_0000500 \ + --hf_export_path /tmp/distilled-hf_iter_0000500 + +Example VLM (checkpoint reshards on load, so TP/PP/EP need not match training): + + torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ + --student_hf_path Qwen/Qwen3-VL-4B-Instruct \ + --megatron_path /tmp/distill-out/checkpoints/iter_0000500 \ + --hf_export_path /tmp/distilled-vlm-hf_iter_0000500 + +See `README.md` in this directory for more details. +""" + +import argparse + +import torch +from megatron.bridge import AutoBridge +from transformers import AutoConfig + +import modelopt.torch.utils.distributed as dist +from modelopt.torch.export import copy_hf_ckpt_remote_code +from modelopt.torch.utils import print_args, print_rank_0 +from modelopt.torch.utils.plugins.mbridge import ( + load_mbridge_model_from_hf, + load_modelopt_megatron_checkpoint, +) + + +def export_llm_to_hf( + megatron_path: str, + hf_export_path: str, + student_hf_path: str, + template_hf: str | None = None, + trust_remote_code: bool = False, +) -> None: + """Export a LLM (Homogeneous or Puzzletron Heterogeneous) Megatron checkpoint to HF. + + Args: + megatron_path: Megatron checkpoint directory (an ``iter_*`` dir or its parent). + hf_export_path: Directory to write the HuggingFace checkpoint to. + student_hf_path: Student HF model used for the exported config / tokenizer. + template_hf: HF model whose architecture matches the exported student. Defaults to + ``student_hf_path``; pass a distinct value only for heterogeneous (Puzzletron/NAS) + students whose architecture differs from ``student_hf_path``. + trust_remote_code: Whether to trust remote code when loading the HF model. + """ + export_bridge = AutoBridge.from_hf_pretrained( + template_hf or student_hf_path, trust_remote_code=trust_remote_code + ) + export_bridge.export_ckpt( + megatron_path=megatron_path, hf_path=hf_export_path, show_progress=True, strict=True + ) + # Config / tokenizer come from the student definition (handles local paths and HF model IDs). + AutoConfig.from_pretrained( + student_hf_path, trust_remote_code=trust_remote_code + ).save_pretrained(hf_export_path) + + +def save_vlm_to_hf( + full_model, + hf_export_path: str, + student_hf_path: str, + trust_remote_code: bool = False, +) -> None: + """Write an in-memory full VLM (distilled LM already in place) to HF format. + + Only the language model is distilled; the vision tower / projector are the original weights, so + the original VLM config / tokenizer / remote code are reused and only the weights are written. + ``full_model.language_model`` must already be a plain module (any KD wrapper stripped by the + caller). Requires the model-parallel groups to be initialized (for the weight gather). + + Args: + full_model: The in-memory full VLM with the distilled language model in place. + hf_export_path: Directory to write the HuggingFace checkpoint to. + student_hf_path: Original VLM HF model providing config / tokenizer / remote code. + trust_remote_code: Whether to trust remote code when loading the HF model. + """ + export_bridge = AutoBridge.from_hf_pretrained( + student_hf_path, trust_remote_code=trust_remote_code + ) + export_bridge.hf_pretrained.save_artifacts(hf_export_path) + export_bridge.save_hf_weights([full_model], hf_export_path) + copy_hf_ckpt_remote_code(student_hf_path, hf_export_path) + + +def get_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument( + "--student_hf_path", + type=str, + required=True, + help="Student HF model (used for the exported config / tokenizer, and, for VLMs, the vision " + "tower / projector weights). Must match the model distilled by distill.py.", + ) + parser.add_argument( + "--megatron_path", + type=str, + required=True, + help="Distilled Megatron checkpoint to convert (an iter_* directory or its parent).", + ) + parser.add_argument( + "--hf_export_path", + type=str, + required=True, + help="Directory to write the exported HuggingFace checkpoint to.", + ) + parser.add_argument( + "--student_hf_model", + type=str, + default=None, + help="Export template: HF model whose architecture matches the exported student. Only needed " + "for heterogeneous (Puzzletron/NAS) students whose architecture differs from --student_hf_path. " + "Defaults to --student_hf_path; unused for VLMs.", + ) + parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code") + parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size") + 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") + + args = parser.parse_args() + print_args(args) + + return args + + +def main(args: argparse.Namespace): + is_vlm = hasattr( + AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), + "vision_config", + ) + + if is_vlm: + # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM + # with the distilled checkpoint weights, then export the assembled VLM. + print_rank_0( + f"Reassembling distilled VLM and exporting to HF format at {args.hf_export_path}" + ) + _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf( + hf_model_name_or_path=args.student_hf_path, + trust_remote_code=args.trust_remote_code, + provider_overrides={ + "tensor_model_parallel_size": args.tp_size, + "pipeline_model_parallel_size": args.pp_size, + "expert_model_parallel_size": args.ep_size, + "context_parallel_size": args.cp_size, + # VLMs run with sequence parallelism off (see distill.py). + "sequence_parallel": False, + "pipeline_dtype": torch.bfloat16, + }, + init_model_parallel=True, + load_weights=True, # vision tower / projector + original LM; the LM is overwritten below + ) + # Load only the distilled language-model weights (skip ModelOpt-state restore -- the kd_loss + # mode / teacher are irrelevant for export and would otherwise require a teacher model). + load_modelopt_megatron_checkpoint( + [full_model.language_model], args.megatron_path, restore_modelopt_state=False + ) + save_vlm_to_hf( + full_model, + 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") + else: + print_rank_0(f"Exporting distilled checkpoint to HF format at {args.hf_export_path}") + # Save rank before destroying process group (dist.rank() won't work after destruction). + is_rank_0 = dist.rank() == 0 + # export_ckpt creates its own temporary process group; destroy this one first so cleanup + # does not hang on a barrier once rank 0 has left. + dist.cleanup() + if is_rank_0: + export_llm_to_hf( + megatron_path=args.megatron_path, + 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, + ) + print_rank_0(f"Exported HuggingFace checkpoint to {args.hf_export_path}") + + +if __name__ == "__main__": + dist.setup() + args = get_args() + try: + main(args) + finally: + dist.cleanup() diff --git a/examples/megatron_bridge/export.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py similarity index 99% rename from examples/megatron_bridge/export.py rename to examples/megatron_bridge/export_quantized_megatron_to_hf.py index 926b47755e5..17db5e6da34 100644 --- a/examples/megatron_bridge/export.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -26,7 +26,7 @@ Example usage to export an FP8 checkpoint produced by quantize.py: - torchrun --nproc_per_node 2 export.py \ + torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \ --hf_model_name_or_path Qwen/Qwen3-8B \ --megatron_path /tmp/Qwen3-8B-FP8-megatron \ --pp_size 2 \ diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index d57e74277c6..ecd0f17e688 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -22,10 +22,10 @@ 3. (Optional) Compress weights to a real low-bit representation. 4. Save the quantized model as a Megatron checkpoint (with ModelOpt state). The checkpoint can be reloaded for further training (QAT / distillation) or converted to a HuggingFace (unified) - checkpoint for deployment with `export.py` (see that script for TensorRT-LLM / vLLM / SGLang). + checkpoint for deployment with `export_quantized_megatron_to_hf.py` (for TensorRT-LLM / vLLM / SGLang). Tensor / pipeline / expert parallelism are all supported here — the Megatron checkpoint is saved -sharded and can be re-sharded on load (e.g. `export.py` reloads it at TP=1 for the HF export). +sharded and can be re-sharded on load (e.g. `export_quantized_megatron_to_hf.py` reloads it at TP=1 for the HF export). Example usage to quantize Qwen3-8B to NVFP4 on 2 GPUs (Tensor Parallelism = 2): 1024 samples from default dataset are used for calibration (sequence length = 4096). @@ -48,7 +48,8 @@ --seq_length 4096 \ --export_megatron_path /tmp/Qwen3-8B-NVFP4-megatron -To convert the saved Megatron checkpoint to a deployable HuggingFace checkpoint, run `export.py`. +To convert the saved Megatron checkpoint to a deployable HuggingFace checkpoint, use +`export_quantized_megatron_to_hf.py`. To see the full usage for advanced configurations, run: torchrun --nproc_per_node 1 quantize.py --help @@ -414,8 +415,8 @@ def forward_loop(_model=None): hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, ) print_rank_0( - f"\nSaved quantized model to {args.export_megatron_path} in Megatron format. " - "To deploy this model (TensorRT-LLM / vLLM / SGLang), convert it to a Unified HF ckpt with export.py" + f"\nSaved quantized model to {args.export_megatron_path} in Megatron format. To deploy this model " + "(TensorRT-LLM / vLLM / SGLang), convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py" ) # Sanity-check generation with the fake-quantized model. Skipped when --compress is set: the diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md index cb4e5f9b650..b2cdb2b3acb 100644 --- a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md +++ b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md @@ -323,7 +323,7 @@ Phase 2 starts as a separate run from a fresh HuggingFace student checkpoint, so
Checkpoint conversion command (click to expand) -> NOTE: Below command only works for non-quantized checkpoints. For quantized checkpoints, we use the `export.py` script in Section 5 to directly export the quantized checkpoint to Unified HF format for deployment. +> NOTE: Below command only works for non-quantized checkpoints. For quantized checkpoints, we use the `export_quantized_megatron_to_hf.py` script in Section 5 to directly export the quantized checkpoint to Unified HF format for deployment. ```bash python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ @@ -394,7 +394,7 @@ We use the same conversion script to convert the Phase 2 final checkpoint to Hug
Checkpoint conversion command (click to expand) -> NOTE: Below command only works for non-quantized checkpoints. For quantized checkpoints, we use the `export.py` script in Section 5 to directly export the quantized checkpoint to Unified HF format for deployment. +> NOTE: Below command only works for non-quantized checkpoints. For quantized checkpoints, we use the `export_quantized_megatron_to_hf.py` script in Section 5 to directly export the quantized checkpoint to Unified HF format for deployment. ```bash python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ @@ -474,7 +474,7 @@ This is done with the `MAMBA_MOE_FP8_CONSERVATIVE_CFG` config defined in [`model > [!NOTE] > You can also quantize to NVFP4 using `--quant_cfg MAMBA_MOE_NVFP4_CONSERVATIVE_CFG` or `MAMBA_MOE_NVFP4_AGGRESSIVE_CFG` (faster, more accuracy drop). NVFP4 typically needs further [Quantization Aware Distillation (QAD)](../../README.md#quantization-aware-distillation-qad) to recover accuracy, plus a Blackwell GPU for deployment. -Quantization is a two-step flow: `quantize.py` calibrates and saves a Megatron checkpoint, then `export.py` converts it to a deployable HuggingFace checkpoint (the unified HF exporter loads at TP=1, so pipeline parallelism is used to shard across GPUs). Both steps take a few minutes on 8x H100. +Quantization is a two-step flow: `quantize.py` calibrates and saves a Megatron checkpoint, then `export_quantized_megatron_to_hf.py` converts it to a deployable HuggingFace checkpoint (the unified HF exporter loads at TP=1, so pipeline parallelism is used to shard across GPUs). Both steps take a few minutes on 8x H100. **Step 1 — calibrate and save the quantized Megatron checkpoint:** @@ -501,7 +501,7 @@ torchrun --nproc_per_node 8 /opt/Model-Optimizer/examples/megatron_bridge/quanti Export command (click to expand) ```bash -torchrun --nproc_per_node 1 /opt/Model-Optimizer/examples/megatron_bridge/export.py \ +torchrun --nproc_per_node 1 /opt/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py \ --hf_model_name_or_path /path/to/distill_output_phase2_32k/checkpoints/hf_iter_0000800 \ --megatron_path /path/to/distill_output_phase2_32k/checkpoints/iter_0000800_fp8_megatron \ --trust_remote_code \ diff --git a/examples/megatron_bridge/tutorials/README.md b/examples/megatron_bridge/tutorials/README.md index 9a27c6914d9..9d6499c8c90 100644 --- a/examples/megatron_bridge/tutorials/README.md +++ b/examples/megatron_bridge/tutorials/README.md @@ -1,7 +1,7 @@ # Megatron-Bridge Tutorials End-to-end tutorials that combine ModelOpt optimization techniques on [NVIDIA Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) models. -Each one walks through a complete workflow using the scripts in [examples/megatron_bridge](../README.md) (`prune_minitron.py`, `distill.py`, `quantize.py`, `export.py`). +Each one walks through a complete workflow using the scripts in [examples/megatron_bridge](../README.md) (`prune_minitron.py`, `distill.py`, `quantize.py`, `export_quantized_megatron_to_hf.py`, `export_distilled_megatron_to_hf.py`). ## Available tutorials diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index e53e6516992..d44bf50ab0f 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -17,7 +17,6 @@ from typing import Any from megatron.bridge import AutoBridge -from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.utils import is_safe_repo from megatron.bridge.models.mamba.mamba_provider import MambaModelProvider @@ -39,48 +38,6 @@ __all__ = ["load_mbridge_model_from_hf", "load_modelopt_megatron_checkpoint"] -def _patch_qwen35_moe_sequential_expert_mappings() -> None: - """WAR: Add sequential (non-grouped) expert mappings to Megatron-Bridge's Qwen3.5 MoE bridge. - - The shipped bridge only maps grouped experts (``experts.gate_up_proj``), but pruning disables - grouped GEMM and needs the sequential ``experts.local_experts.*`` layout. This also covers - Qwen3.5-VL MoE, whose bridge delegates to the same ``_get_moe_lm_mappings`` helper. - - TODO: Remove once Megatron-Bridge maps sequential Qwen3.5 MoE experts natively (patched in 26.06.01). - """ - try: - from megatron.bridge.models.qwen.qwen35_bridge import Qwen35MoEBridge - except ImportError: - return - - orig = Qwen35MoEBridge._get_moe_lm_mappings - if getattr(orig, "_modelopt_sequential_experts", False): - return - # No-op if the installed bridge already maps sequential experts. - if any( - "local_experts" in str(getattr(m, "megatron_param", "")) - for m in orig(hf_prefix="model.", megatron_prefix="") - ): - return - - def _get_moe_lm_mappings(hf_prefix="model.", megatron_prefix=""): - return [ - *orig(hf_prefix=hf_prefix, megatron_prefix=megatron_prefix), - GatedMLPMapping( - megatron_param=f"{megatron_prefix}decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", - gate=f"{hf_prefix}layers.*.mlp.experts.*.gate_proj.weight", - up=f"{hf_prefix}layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param=f"{megatron_prefix}decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", - hf_param=f"{hf_prefix}layers.*.mlp.experts.*.down_proj.weight", - ), - ] - - _get_moe_lm_mappings._modelopt_sequential_experts = True # type: ignore[attr-defined] - Qwen35MoEBridge._get_moe_lm_mappings = staticmethod(_get_moe_lm_mappings) - - def load_mbridge_model_from_hf( *, hf_model_name_or_path: str, @@ -113,7 +70,6 @@ def load_mbridge_model_from_hf( A tuple of (bridge, provider, model, unwrapped_model, tokenizer). """ print_rank_0(f"Loading Megatron-Bridge model from HF: {hf_model_name_or_path}") - _patch_qwen35_moe_sequential_expert_mappings() trust_remote_code = is_safe_repo( trust_remote_code=trust_remote_code, hf_path=hf_model_name_or_path, @@ -162,16 +118,21 @@ def load_mbridge_model_from_hf( return bridge, provider, model, unwrapped_model, tokenizer -def load_modelopt_megatron_checkpoint(model: list[MegatronModule], megatron_path: str) -> None: +def load_modelopt_megatron_checkpoint( + model: list[MegatronModule], megatron_path: str, restore_modelopt_state: bool = True +) -> None: """Load Megatron checkpoint weights (with modelopt_state). Args: model: The (pre-built) Megatron model to load the checkpoint into. megatron_path: Path to the quantized Megatron checkpoint (produced by ``quantize.py``) + restore_modelopt_state: Whether to restore the ModelOpt state (e.g. quantizers) before loading + weights. Set ``False`` to load weights only -- e.g. to reload a full-precision distilled + student without reconstructing the ``kd_loss`` mode (which would require a teacher model). """ # Restore the ModelOpt state before loading weights. # has_modelopt_state / load_modelopt_state resolves the latest iter_* directory - if has_modelopt_state(megatron_path): + if restore_modelopt_state and has_modelopt_state(megatron_path): load_modelopt_state(model, megatron_path) # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly _load_model_weights_from_checkpoint(_get_modelopt_checkpoint_path(megatron_path), model) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 8b5e206567d..10649ba2195 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -16,14 +16,20 @@ from pathlib import Path +import torch from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.puzzletron.utils import create_and_save_small_hf_model -from _test_utils.torch.transformers_models import create_tiny_qwen3_dir, get_tiny_tokenizer +from _test_utils.torch.transformers_models import ( + create_tiny_qwen3_5_vl_dir, + create_tiny_qwen3_dir, + get_tiny_tokenizer, +) +from transformers import AutoModelForImageTextToText from modelopt.torch.puzzletron.anymodel import convert_model -def test_distill_and_convert(tmp_path: Path, num_gpus): +def test_distill_llm(tmp_path, num_gpus): teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) train_iters = 2 distill_output_dir = tmp_path / "distill_output" @@ -44,7 +50,6 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): eval_iters=1, log_interval=1, hf_export_path=distilled_hf_path, - student_hf_model=teacher_hf_path, ) run_example_command(distill_cmd_parts, example_path="megatron_bridge") @@ -52,6 +57,88 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): assert (distilled_hf_path / "config.json").exists() +# NOTE: Qwen3.5-VL-MoE covered by test_qad.py +def test_distill_vlm(tmp_path, num_gpus): + # Self-distillation of a tiny VLM: only the language model is distilled; the vision tower and the + # vision->language projector must be left byte-for-byte untouched. + # + # Distillation runs under tensor parallelism (sequence parallelism is disabled for VLMs -- see + # distill.py -- pending the standalone-LM SP fix in the nemo:26.08 container). + vlm_hf_path, teacher_model = create_tiny_qwen3_5_vl_dir( + tmp_path, + with_tokenizer=True, + return_model=True, + num_hidden_layers=2, + intermediate_size=128, + ) + + # The language model spans ``model.language_model.*`` plus the (possibly untied) output head + # ``lm_head`` -- both are distilled and may change. Everything else (vision tower + projector) + # must stay byte-for-byte identical. + def _is_language_model(name: str) -> bool: + return name.startswith(("model.language_model.", "lm_head")) + + teacher_non_lm = { + name: param.detach().clone() + for name, param in teacher_model.named_parameters() + if not _is_language_model(name) + } + assert teacher_non_lm, ( + "Expected non-language-model params (vision tower / projector) in the VLM." + ) + + train_iters = 2 + distill_output_dir = tmp_path / "distill_output" + distilled_hf_path = tmp_path / "distilled_hf" + distill_cmd_parts = extend_cmd_parts( + ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"], + student_hf_path=vlm_hf_path, + teacher_hf_path=vlm_hf_path, + output_dir=distill_output_dir, + tp_size=num_gpus, + pp_size=1, + seq_length=16, + mbs=1, + gbs=4, + train_iters=train_iters, + lr_warmup_iters=1, + eval_interval=train_iters, + eval_iters=1, + log_interval=1, + ) + run_example_command(distill_cmd_parts, example_path="megatron_bridge") + + megatron_ckpt = distill_output_dir / f"checkpoints/iter_{train_iters:07d}" + assert megatron_ckpt.exists() + + # Separately convert the distilled Megatron checkpoint to HF with the standalone export script + # TODO: VLMs disable sequence parallelism, so tensor parallelism can't be used here. + # Flip to tp_size=num_gpus in nemo:26.08 container + export_cmd_parts = extend_cmd_parts( + ["torchrun", f"--nproc_per_node={num_gpus}", "export_distilled_megatron_to_hf.py"], + student_hf_path=vlm_hf_path, + megatron_path=megatron_ckpt, + hf_export_path=distilled_hf_path, + pp_size=num_gpus, + ) + run_example_command(export_cmd_parts, example_path="megatron_bridge") + + assert (distilled_hf_path / "config.json").exists() + + # from_pretrained (default strict load) verifies the saved weights match the (unchanged) config. + distilled_model = AutoModelForImageTextToText.from_pretrained(distilled_hf_path) + assert hasattr(distilled_model.config, "vision_config") + distilled_params = dict(distilled_model.named_parameters()) + # Everything outside the language model is identical to the teacher (vision tower untouched). + for name, teacher_param in teacher_non_lm.items(): + assert name in distilled_params, ( + f"Missing non-language-model param after distillation: {name}" + ) + assert torch.equal(distilled_params[name].float(), teacher_param.float()), ( + f"Non-language-model param changed during LM-only distillation: {name}" + ) + + def test_distill_puzzletron_anymodel(tmp_path: Path, num_gpus): """Integration test for distill.py with Puzzletron AnyModel (heterogeneous) checkpoints. diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 8ca09a77967..f3fa786176f 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -23,9 +23,21 @@ create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_dir, ) +from megatron.bridge.models.conversion import model_bridge from transformers import AutoModelForCausalLM, AutoModelForImageTextToText +def _qwen35_moe_bridge_supported() -> bool: + """Whether MBridge supports Qwen3.5-MoE per-expert weight assembly, i.e. nemo:26.08+. + + Mount an updated MBridge to run them on 26.06. + """ + try: + return hasattr(model_bridge, "_fuse_per_expert_hf_weight") + except Exception: + return False + + @pytest.mark.parametrize( ("create_teacher", "megatron_format"), [ @@ -109,6 +121,10 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): max_position_embeddings=1024, ), id="qwen3_5_moe_vl", + marks=pytest.mark.skipif( + not _qwen35_moe_bridge_supported(), + reason="Qwen3.5-MoE needs Megatron-Bridge native MoE support (nemo:26.08+)", + ), ), ], ) diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index c1dfc26400e..5982bad7b53 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -18,27 +18,81 @@ import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.transformers_models import create_tiny_qwen3_dir +from _test_utils.torch.transformers_models import ( + create_tiny_gemma3vl_dir, + create_tiny_qwen3_5_moe_vl_dir, + create_tiny_qwen3_dir, +) +from megatron.bridge.models.conversion import model_bridge -@pytest.mark.timeout(720) # Multiple steps in one test hence takes longer than default timeout -def test_qad(tmp_path: Path, num_gpus): - """Quantize a tiny Qwen3, run QAD from the quantized student, and export to a unified HF ckpt.""" - hf_model_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) - quantized_megatron_path = tmp_path / "qwen3_fp8_megatron" +def _qwen35_moe_bridge_supported() -> bool: + """Whether MBridge supports Qwen3.5-MoE per-expert weight assembly, i.e. nemo:26.08+. + + Mount an updated MBridge to run them on 26.06. + """ + try: + return hasattr(model_bridge, "_fuse_per_expert_hf_weight") + except Exception: + return False + + +@pytest.mark.timeout(720) # Multiple steps in one test hence takes longer than the default timeout +@pytest.mark.parametrize( + ("create_student", "is_vlm", "is_moe"), + [ + (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), False, False), + ( + # TODO: remove once qwen35_vl_moe is supported + lambda tmp_path: create_tiny_gemma3vl_dir( + tmp_path, + with_processor=True, + num_hidden_layers=2, + intermediate_size=128, + max_position_embeddings=512, + ), + True, + False, + ), + pytest.param( + lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir( + tmp_path, with_processor=True, num_hidden_layers=2 + ), + True, + True, + marks=pytest.mark.skipif( + not _qwen35_moe_bridge_supported(), + reason="Qwen3.5-MoE needs Megatron-Bridge native MoE support (nemo:26.08+)", + ), + ), + ], + ids=["qwen3", "gemma3vl", "qwen3_5_moe_vl"], +) +def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): + """Quantize a tiny model, run QAD from the quantized student, and export the result. + + For VLMs only the language model is quantized and distilled (vision tower / projector untouched), + and a text calibration dataset infers text-only LM calibration. VLM quantized-HF export is + unsupported, so the VLM case stops at the distilled Megatron checkpoint and verifies the ModelOpt + (quantize) state survived distillation; the LLM case additionally exports a unified HF checkpoint. + """ + hf_model_path = create_student(tmp_path) + quantized_megatron_path = tmp_path / "quantized_megatron" distill_output_dir = tmp_path / "qad_output" - hf_export_path = tmp_path / "qwen3_qad_fp8_hf" train_iters = 2 + # TODO: MoE VLMs run single-GPU: MoE layers require sequence parallelism when TP>1, but VLM distillation + # disables SP (the standalone-LM SP fix lands in the nemo:26.08 container). Others use all GPUs. + tp_size = 1 if is_moe else num_gpus - # Step 1: PTQ the model to FP8 and save a Megatron checkpoint carrying the ModelOpt state. + # Step 1: PTQ the (language) model to FP8 and save a Megatron checkpoint carrying the ModelOpt state. quantize_cmd = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"], + ["torchrun", f"--nproc_per_node={tp_size}", "quantize.py", "--skip_generate"], hf_model_name_or_path=hf_model_path, recipe="general/ptq/fp8_default-kv_fp8", - tp_size=num_gpus, - calib_dataset_name="cnn_dailymail", - calib_num_samples=16, - calib_batch_size=4, + tp_size=tp_size, + calib_dataset_name="cnn_dailymail", # text dataset -> (for VLMs) text-only LM calibration + calib_num_samples=8, + calib_batch_size=2, seq_length=16, export_megatron_path=quantized_megatron_path, ) @@ -47,16 +101,16 @@ def test_qad(tmp_path: Path, num_gpus): "Expected modelopt_state in the quantized Megatron checkpoint" ) - # Step 2: QAD -- load the quantized student from the Megatron checkpoint (restoring the - # ModelOpt quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint - # must keep the ModelOpt state so it can later be exported as a quantized HF checkpoint. + # Step 2: QAD -- load the quantized student from the Megatron checkpoint (restoring the ModelOpt + # quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint must keep the + # ModelOpt state so the quantizers survive distillation. distill_cmd = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"], + ["torchrun", f"--nproc_per_node={tp_size}", "distill.py", "--use_mock_data"], student_hf_path=hf_model_path, student_megatron_path=quantized_megatron_path, teacher_hf_path=hf_model_path, output_dir=distill_output_dir, - tp_size=num_gpus, + tp_size=tp_size, seq_length=16, mbs=1, gbs=4, @@ -73,10 +127,14 @@ def test_qad(tmp_path: Path, num_gpus): "Expected modelopt_state to be preserved in the distilled (QAD) checkpoint" ) + if is_vlm: + return # VLM quantized-HF export is unsupported; stop at the distilled Megatron checkpoint + # Step 3: export the distilled quantized checkpoint to a unified HF checkpoint. hf_quant_config.json # is only written for a quantized model, so its presence confirms the quantizers survived QAD. + hf_export_path = tmp_path / "qad_fp8_hf" export_cmd = extend_cmd_parts( - ["torchrun", "--nproc_per_node=1", "export.py"], + ["torchrun", "--nproc_per_node=1", "export_quantized_megatron_to_hf.py"], hf_model_name_or_path=hf_model_path, megatron_path=distilled_megatron_path, export_unified_hf_path=hf_export_path, diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 8e9402337f7..55a97b49d00 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -12,14 +12,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for quantize.py and export.py scripts.""" +"""Tests for quantize.py and export_quantized_megatron_to_hf.py scripts.""" from pathlib import Path from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.transformers_models import create_tiny_qwen3_5_vl_dir, create_tiny_qwen3_dir +from _test_utils.torch.transformers_models import create_tiny_qwen3_dir +# NOTE: Qwen3.5-VL covered by test_qad.py def test_quantize_and_export(tmp_path: Path, num_gpus): """Quantize a tiny Qwen3 via a YAML recipe and export it to a unified HF checkpoint.""" # Use a vLLM-friendly head_dim (64) since the default tiny config (head_dim=2) is unsupported. @@ -56,7 +57,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus): # Step 2: export to HF export_cmd = extend_cmd_parts( - ["torchrun", "--nproc_per_node=1", "export.py"], + ["torchrun", "--nproc_per_node=1", "export_quantized_megatron_to_hf.py"], hf_model_name_or_path=hf_model_path, megatron_path=megatron_path, export_unified_hf_path=hf_export_path, @@ -80,41 +81,3 @@ def test_quantize_and_export(tmp_path: Path, num_gpus): # ) # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4)) # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text - - -def test_quantize_vlm(tmp_path: Path, num_gpus): - """Quantize a tiny Qwen3.5-VL's language model and save a Megatron checkpoint. - - Only the language model is quantized (vision quantizers disabled, ModelOpt state on the root); a - text calibration dataset infers text-only calibration. Saves Megatron format only -- HF unified - export of a VLM is unsupported, matching Megatron-Bridge's quantize_vlm.py. - """ - hf_model_path = create_tiny_qwen3_5_vl_dir( - tmp_path, - with_processor=True, - hidden_size=128, - num_attention_heads=2, - num_key_value_heads=2, - num_hidden_layers=2, - intermediate_size=256, - max_position_embeddings=512, - ) - megatron_path = tmp_path / "qwen3_5_vl_fp8_megatron" - - quantize_cmd = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"], - hf_model_name_or_path=hf_model_path, - recipe="general/ptq/fp8_default-kv_fp8", - tp_size=num_gpus, - # A text dataset on a VLM infers text-only calibration of its language model (ablation path). - calib_dataset_name="cnn_dailymail", - calib_num_samples=4, - calib_batch_size=2, - seq_length=16, - export_megatron_path=megatron_path, - ) - run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) - assert (megatron_path / "latest_checkpointed_iteration.txt").exists() - assert list(megatron_path.rglob("modelopt_state")), ( - "Expected modelopt_state in the Megatron checkpoint" - ) From bdb745c7e19c51a3e4277619896df49ce0a32b44 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:11:19 -0700 Subject: [PATCH 2/8] address review: VLM export template arg, quantize hint scope, dedupe helper - distill.py: VLM HF export uses --student_hf_path (the VLM config/tokenizer source), not --student_hf_model (which is unused for VLMs and only coincides via its default). - quantize.py: only point non-VLM checkpoints at export_quantized_megatron_to_hf.py; quantized-VLM HF export is unsupported, so say so for VLMs. - Move the duplicated _qwen35_moe_bridge_supported() helper into tests/_test_utils/examples/megatron_bridge.py. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/distill.py | 2 +- examples/megatron_bridge/quantize.py | 14 +++++++--- tests/_test_utils/examples/megatron_bridge.py | 28 +++++++++++++++++++ .../megatron_bridge/test_prune_minitron.py | 15 ++-------- tests/examples/megatron_bridge/test_qad.py | 15 ++-------- 5 files changed, 43 insertions(+), 31 deletions(-) create mode 100644 tests/_test_utils/examples/megatron_bridge.py diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index bbb8a26731a..767ddfc3a06 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -385,7 +385,7 @@ def _restore_student_hook(model_chunks): save_vlm_to_hf( full_student, args.hf_export_path, - args.student_hf_model, + 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") diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index ecd0f17e688..3454da00441 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -414,10 +414,16 @@ def forward_loop(_model=None): hf_tokenizer_path=args.hf_model_name_or_path, hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, ) - print_rank_0( - f"\nSaved quantized model to {args.export_megatron_path} in Megatron format. To deploy this model " - "(TensorRT-LLM / vLLM / SGLang), convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py" - ) + if is_vlm: + print_rank_0( + f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format " + "(HuggingFace unified export of a quantized VLM is not yet supported)." + ) + else: + print_rank_0( + f"\nSaved quantized model to {args.export_megatron_path} in Megatron format. To deploy this model " + "(TensorRT-LLM / vLLM / SGLang), convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py" + ) # Sanity-check generation with the fake-quantized model. Skipped when --compress is set: the # weights are now real low-bit and megatron_generate may not support compressed forward for diff --git a/tests/_test_utils/examples/megatron_bridge.py b/tests/_test_utils/examples/megatron_bridge.py new file mode 100644 index 00000000000..0989578c3c6 --- /dev/null +++ b/tests/_test_utils/examples/megatron_bridge.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared helpers for Megatron-Bridge example tests.""" + + +def qwen35_moe_bridge_supported() -> bool: + """Whether MBridge supports Qwen3.5-MoE per-expert weight assembly, i.e. nemo:26.08+. + + Mount an updated MBridge to run these on 26.06. + """ + try: + from megatron.bridge.models.conversion import model_bridge + + return hasattr(model_bridge, "_fuse_per_expert_hf_weight") + except Exception: + return False diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index f3fa786176f..f21f25d5209 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -16,6 +16,7 @@ import pytest import torch +from _test_utils.examples.megatron_bridge import qwen35_moe_bridge_supported from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.transformers_models import ( create_tiny_gemma3vl_dir, @@ -23,21 +24,9 @@ create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_dir, ) -from megatron.bridge.models.conversion import model_bridge from transformers import AutoModelForCausalLM, AutoModelForImageTextToText -def _qwen35_moe_bridge_supported() -> bool: - """Whether MBridge supports Qwen3.5-MoE per-expert weight assembly, i.e. nemo:26.08+. - - Mount an updated MBridge to run them on 26.06. - """ - try: - return hasattr(model_bridge, "_fuse_per_expert_hf_weight") - except Exception: - return False - - @pytest.mark.parametrize( ("create_teacher", "megatron_format"), [ @@ -122,7 +111,7 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): ), id="qwen3_5_moe_vl", marks=pytest.mark.skipif( - not _qwen35_moe_bridge_supported(), + not qwen35_moe_bridge_supported(), reason="Qwen3.5-MoE needs Megatron-Bridge native MoE support (nemo:26.08+)", ), ), diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index 5982bad7b53..6ccf97a52a2 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -17,24 +17,13 @@ from pathlib import Path import pytest +from _test_utils.examples.megatron_bridge import qwen35_moe_bridge_supported from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.transformers_models import ( create_tiny_gemma3vl_dir, create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_dir, ) -from megatron.bridge.models.conversion import model_bridge - - -def _qwen35_moe_bridge_supported() -> bool: - """Whether MBridge supports Qwen3.5-MoE per-expert weight assembly, i.e. nemo:26.08+. - - Mount an updated MBridge to run them on 26.06. - """ - try: - return hasattr(model_bridge, "_fuse_per_expert_hf_weight") - except Exception: - return False @pytest.mark.timeout(720) # Multiple steps in one test hence takes longer than the default timeout @@ -61,7 +50,7 @@ def _qwen35_moe_bridge_supported() -> bool: True, True, marks=pytest.mark.skipif( - not _qwen35_moe_bridge_supported(), + not qwen35_moe_bridge_supported(), reason="Qwen3.5-MoE needs Megatron-Bridge native MoE support (nemo:26.08+)", ), ), From 0ba8471d16c1548ba05721172a31197fbf727b99 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 9 Jul 2026 08:43:47 -0700 Subject: [PATCH 3/8] minor Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/README.md | 2 +- .../megatron_bridge/_distillation_provider.py | 15 +++++++++++++-- examples/megatron_bridge/distill.py | 13 ++++++------- examples/megatron_bridge/prune_minitron.py | 5 +++++ modelopt/torch/nas/plugins/__init__.py | 6 ++++-- 5 files changed, 29 insertions(+), 12 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 7075cc9f5a2..6886018dfa6 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -51,7 +51,7 @@ docker run \ > **Working with MoE Vision-Language Models (e.g. Qwen3.5-VL-MoE)?** The `nemo:26.06` container's Megatron-Bridge lacks the MoE expert weight mappings these models need (dense VLMs such as Gemma3-VL and Qwen3-VL work as-is). Until the `nemo:26.08` container is released, mount the latest [Megatron-Bridge `main`](https://github.com/NVIDIA-NeMo/Megatron-Bridge) source over the pre-installed copy by adding this to the `docker run` command above: > > ```bash -> -v /path/to/Megatron-Bridge/src/megatron/bridge:/opt/Megatron-Bridge/src/megatron/bridge +> -v ${MEGATRON_BRIDGE_DIR}:/opt/Megatron-Bridge > ``` You also need to login with your HuggingFace token to download gated datasets / models. diff --git a/examples/megatron_bridge/_distillation_provider.py b/examples/megatron_bridge/_distillation_provider.py index b78d256eda2..937d3254ae0 100644 --- a/examples/megatron_bridge/_distillation_provider.py +++ b/examples/megatron_bridge/_distillation_provider.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Minimal extension of Megatron-Bridge's ``DistillationProvider``. +"""Minimal extension of Megatron-Bridge's ``DistillationProvider`` for nemo:26.06 and older containers. Adds two things over the stock provider: (1) the KD conversion runs in a post-weight-load pre-wrap hook instead of in ``provide()``, and (2) a ``distill_submodule`` option to distill only a submodule (e.g. a @@ -24,6 +24,8 @@ ``megatron.bridge.models.distillation_provider`` once we require the nemo:26.08 container (Megatron-Bridge#4707). """ +import inspect + from megatron.bridge.models.distillation_provider import ( convert_to_distillation_provider as _base_convert_to_distillation_provider, ) @@ -65,7 +67,7 @@ def _convert_hook(self, model_chunks): return [kd_model] -def convert_to_distillation_provider( +def _shim_convert_to_distillation_provider( student_provider, teacher_provider, kd_config=None, *, distill_submodule=None ): """Like ``megatron.bridge``'s ``convert_to_distillation_provider`` but defers the KD conversion to a @@ -95,3 +97,12 @@ def convert_to_distillation_provider( hooks = [*getattr(provider, "_pre_wrap_hooks", []), provider._convert_hook] object.__setattr__(provider, "_pre_wrap_hooks", hooks) return provider + + +# Prefer Megatron-Bridge's native implementation when it supports submodule distillation; otherwise +# fall back to the local back-port for older containers. +convert_to_distillation_provider = ( + _base_convert_to_distillation_provider + if "distill_submodule" in inspect.signature(_base_convert_to_distillation_provider).parameters + else _shim_convert_to_distillation_provider +) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 767ddfc3a06..947a1387e57 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -210,11 +210,6 @@ def main(args: argparse.Namespace): checkpoint_dir = os.path.join(args.output_dir, "checkpoints") tensorboard_dir = os.path.join(args.output_dir, "tb_logs") - is_vlm = hasattr( - AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), - "vision_config", - ) - # Build student and teacher model providers def _build_model_provider(hf_path, load_weights=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) @@ -222,8 +217,7 @@ def _build_model_provider(hf_path, load_weights=True): # Override parallelism / training settings provider.tensor_model_parallel_size = args.tp_size - # TODO: re-enable VLM SP from the nemo:26.08 container (fixed upstream by Megatron-LM#5628) - provider.sequence_parallel = args.tp_size > 1 and not is_vlm + provider.sequence_parallel = args.tp_size > 1 provider.pipeline_model_parallel_size = args.pp_size provider.pipeline_dtype = torch.bfloat16 provider.context_parallel_size = args.cp_size @@ -257,6 +251,11 @@ def _build_model_provider(hf_path, load_weights=True): skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) + 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). " diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index bdd40b3ad8f..1d32d0df09b 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -72,6 +72,11 @@ from modelopt.torch.utils.plugins.megatron_mmlu import megatron_mmlu from modelopt.torch.utils.vlm_dataset_utils import get_supported_vlm_datasets +# isort: off +# Register Megatron-Bridge model-specific NAS/pruning plugins here to avoid a circular import +import modelopt.torch.nas.plugins.mbridge # noqa: F401 +# isort: on + # Default calibration datasets when --calib_dataset_name is not set DEFAULT_TEXT_CALIB_DATASET = "nemotron-post-training-dataset-v2" DEFAULT_VLM_CALIB_DATASET = "nemotron_vlm_dataset_v2" diff --git a/modelopt/torch/nas/plugins/__init__.py b/modelopt/torch/nas/plugins/__init__.py index 3abcd799199..b9bd6b8a11c 100644 --- a/modelopt/torch/nas/plugins/__init__.py +++ b/modelopt/torch/nas/plugins/__init__.py @@ -23,5 +23,7 @@ from .megatron import * from .megatron_model_stats import * -with import_plugin("megatron.bridge"): - from .mbridge import * +# NOTE: Megatron-Bridge plugin is intentionally NOT auto-imported here to avoid a circular import. +# It is imported from the pruning entrypoint (examples/megatron_bridge/prune_minitron.py) +# with import_plugin("megatron.bridge"): +# from .mbridge import * From cbddfbe8cb27827e5fa352778256ef7424c0d45c Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 14 Jul 2026 12:53:51 -0700 Subject: [PATCH 4/8] VLM pruning improvements Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/megatron_bridge/README.md | 6 +- examples/megatron_bridge/prune_minitron.py | 21 ++- .../torch/nas/plugins/megatron_model_stats.py | 13 +- .../torch/prune/plugins/mcore_minitron.py | 142 +++++++++++------- .../nas/plugins/test_megatron_model_stats.py | 10 ++ .../test_mcore_gpt_minitron_pruning.py | 16 +- 7 files changed, 138 insertions(+), 72 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3d4506f2a83..7dd96e6ac69 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -34,7 +34,7 @@ Changelog - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned. - Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations. - Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported. -- Add Megatron-Bridge distillation support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. +- Add Megatron-Bridge distillation and QAD support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. - Add NVFP4 Four-Over-Six (4/6) weight quantization (``mtq.NVFP4_FOUR_OVER_SIX_CFG``): MSE weight calibration picks, per block, between an M=6 and an M=4 dynamic range (the choice is folded into the FP8 per-block scales), with the ``four_over_six: true`` flag normalizing those scales by 256 (vs 448) for M=4 headroom. Supported via ``mtq.quantize`` and HF / Megatron export only -- **not** ``mtq.compress``, which does not preserve the per-block M=4/M=6 choice - Add dLLM (tied-weight PTQ and HF-checkpoint export) support for diffusion-based encoder-decoder LLMs (e.g. DiffusionGemma) whose encoder/decoder stacks share parameters via HF ``_tied_weights_keys``. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 6886018dfa6..2a3ff9558bf 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -115,10 +115,14 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati ## Sanity-Check Generation -[generate_vllm.py](generate_vllm.py) runs a quick generation check on a unified HuggingFace checkpoint using vLLM. vLLM auto-detects the ModelOpt quantization from the exported `hf_quant_config.json`, so no extra quant flags are needed: +[generate_vllm.py](generate_vllm.py) runs a quick generation check on an exported HuggingFace checkpoint using vLLM — a useful smoke test for both a **quantized** and a **pruned** model to confirm it still produces coherent text before investing in [distillation](#distillation) to recover accuracy. For quantized checkpoints, vLLM auto-detects the ModelOpt quantization from the exported `hf_quant_config.json`, so no extra flags are needed: ```bash +# Quantized model python generate_vllm.py --model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 --trust_remote_code + +# Pruned model +python generate_vllm.py --model /tmp/Qwen3-8B-Pruned-6B ``` ## Distillation diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 1d32d0df09b..5bead2491ba 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -549,22 +549,17 @@ def score_func(m): ) setattr(provider, hybrid_key, getattr(unwrapped_model, hybrid_key)) - if args.output_megatron_path is not None: - print_rank_0( - f"Saved pruned model to {args.output_megatron_path} in Megatron checkpoint format" - ) + # NOTE: Issue with NemotronH tokenizer's len() hence using use_fast=True as a WAR. + architectures = getattr(bridge.hf_pretrained.config, "architectures", None) or [] + use_fast_tokenizer = "NemotronHForCausalLM" in architectures + tokenizer_kwargs = {"trust_remote_code": args.trust_remote_code, "use_fast": use_fast_tokenizer} - # NOTE: Issue with NemotronH tokenizer's len() hence using use_fast=True as a WAR. - architectures = getattr(bridge.hf_pretrained.config, "architectures", None) or [] - use_fast_tokenizer = "NemotronHForCausalLM" in architectures + if args.output_megatron_path is not None: bridge.save_megatron_model( model, args.output_megatron_path, hf_tokenizer_path=args.hf_model_name_or_path, - hf_tokenizer_kwargs={ - "trust_remote_code": args.trust_remote_code, - "use_fast": use_fast_tokenizer, - }, + hf_tokenizer_kwargs=tokenizer_kwargs, ) print_rank_0( f"Saved pruned model to {args.output_megatron_path} in Megatron checkpoint format" @@ -572,7 +567,9 @@ def score_func(m): else: print_rank_0(f"Saving pruned model to {args.output_hf_path} in HF checkpoint format") - # [WAR] Hacky way to save pruned HF model until Megatron-Bridge natively supports it + # [WAR] Save the pruned HF model by hand until Megatron-Bridge natively supports it. + # TODO: Replace this whole block with ``AutoBridge.from_auto_config(...).save_hf_weights(...)`` + # once the Megatron-Bridge fix ships (nemo:26.08). bridge.hf_pretrained.save_artifacts(args.output_hf_path) hf_cfg = AutoConfig.from_pretrained( args.output_hf_path, trust_remote_code=args.trust_remote_code diff --git a/modelopt/torch/nas/plugins/megatron_model_stats.py b/modelopt/torch/nas/plugins/megatron_model_stats.py index 223f7201fa6..b3a8749962a 100644 --- a/modelopt/torch/nas/plugins/megatron_model_stats.py +++ b/modelopt/torch/nas/plugins/megatron_model_stats.py @@ -394,14 +394,23 @@ def _get(attr: str, default: Any = None) -> Any: # GatedDeltaNet (linear-attention) hybrid, e.g. Qwen3-Next / Qwen3.5: every ``linear_attention_freq``-th # layer keeps full attention, the rest use GatedDeltaNet linear attention. experimental_attention_variant: str | None = _get("experimental_attention_variant", None) - linear_attention_freq: int | None = _get("linear_attention_freq", None) + # ``linear_attention_freq`` is either an int interval (every N-th layer is full attention) or an + # explicit per-layer pattern (list; 1 = linear_attention / GatedDeltaNet, 0 = full_attention). + linear_attention_freq: int | list[int] | None = _get("linear_attention_freq", None) is_gdn = experimental_attention_variant == "gated_delta_net" and bool(linear_attention_freq) + + def _is_linear_attention_layer(i: int) -> bool: + """Whether layer ``i`` uses GatedDeltaNet linear attention (vs full attention).""" + if isinstance(linear_attention_freq, (list, tuple)): + return bool(linear_attention_freq[i]) + return (i + 1) % linear_attention_freq != 0 + # Multi-Latent Attention (MLA), e.g. DeepSeek / Kimi: low-rank compressed q/kv projections. multi_latent_attention: bool = _get("multi_latent_attention", False) def _attn_params(i: int) -> int: """Params for the attention sublayer of layer ``i`` (GatedDeltaNet / MLA / standard).""" - if is_gdn and (i + 1) % linear_attention_freq != 0: + if is_gdn and _is_linear_attention_layer(i): return _gated_delta_net_layer_params( hidden_size, _get("linear_num_key_heads"), diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index cee43d85807..df61713b480 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -27,6 +27,7 @@ import io import sys from collections.abc import Callable +from contextlib import contextmanager from dataclasses import dataclass from functools import partial from itertools import product @@ -110,6 +111,10 @@ "num_layers", } +# Explicit per-layer config patterns (indexed by original 1-indexed layer number) that must be sliced +# to the surviving layers on depth pruning. Integer/None cadence values are ignored (need no update). +_PER_LAYER_CONFIG_PATTERNS = ("linear_attention_freq", "moe_layer_freq") + __all__ = [ "SUPPORTED_HPARAMS", "MCoreMinitronConfig", @@ -120,6 +125,31 @@ ] +def _get_hybrid_pattern_key(model: nn.Module) -> str | None: + """Return the attribute name carrying the hybrid block pattern for hybrid models, else None. + + Handles both ``MambaModel`` (which still uses ``hybrid_override_pattern``) and plain + ``HybridModel`` (the parent class introduced in modern Megatron-LM, which carries + ``hybrid_layer_pattern``). Detecting by attribute presence avoids fragile isinstance + checks against a class hierarchy that may shift across MCore versions. + """ + for attr in ("hybrid_override_pattern", "hybrid_layer_pattern"): + if getattr(model, attr, None): + return attr + return None + + +def _slice_per_layer_pattern(pattern: list | tuple | str, dropped_layers: set[int]): + """Slice a per-layer pattern to the surviving layers, preserving its type. + + ``pattern`` is indexed by the original 1-indexed layer number (a list/tuple such as + ``linear_attention_freq`` or a hybrid block-type string such as ``"M*M-"``). Entries whose + layer number is in ``dropped_layers`` are removed. + """ + kept = [p for i, p in enumerate(pattern) if (i + 1) not in dropped_layers] + return "".join(kept) if isinstance(pattern, str) else type(pattern)(kept) + + def drop_mcore_language_model_layers(model: nn.Module, *, layers_to_drop: list[int]) -> None: """Remove given layers (1-indexed) of the model (works with TP and/or PP). @@ -173,19 +203,18 @@ def drop_mcore_language_model_layers(model: nn.Module, *, layers_to_drop: list[i model.config.num_layers = new_num_layers - -def _get_hybrid_pattern_key(model: nn.Module) -> str | None: - """Return the attribute name carrying the hybrid block pattern for hybrid models, else None. - - Handles both ``MambaModel`` (which still uses ``hybrid_override_pattern``) and plain - ``HybridModel`` (the parent class introduced in modern Megatron-LM, which carries - ``hybrid_layer_pattern``). Detecting by attribute presence avoids fragile isinstance - checks against a class hierarchy that may shift across MCore versions. - """ - for attr in ("hybrid_override_pattern", "hybrid_layer_pattern"): - if getattr(model, attr, None): - return attr - return None + # Slice per-layer patterns to the surviving layers, else their length no longer matches + # ``num_layers`` and building/exporting the pruned model fails. + dropped = set(layers_to_drop) + # Explicit config lists (e.g. linear_attention_freq, moe_layer_freq); integer/None values are left untouched + for pattern_attr in _PER_LAYER_CONFIG_PATTERNS: + pattern = getattr(model.config, pattern_attr, None) + if isinstance(pattern, (list, tuple)): + setattr(model.config, pattern_attr, _slice_per_layer_pattern(pattern, dropped)) + # Hybrid block-type string (Mamba ``hybrid_override_pattern`` / ``hybrid_layer_pattern``). + hybrid_key = _get_hybrid_pattern_key(model) + if hybrid_key is not None: + setattr(model, hybrid_key, _slice_per_layer_pattern(getattr(model, hybrid_key), dropped)) def _rprint(*renderables: Any) -> None: @@ -385,36 +414,52 @@ def run_search(self) -> None: export_config = self.constraints["export_config"] # Prune homogeneously - self._prune(export_config, prune_depth=True) - - # Update the hybrid block-type pattern if pruning a hybrid model. - hybrid_key = _get_hybrid_pattern_key(self.model) - if hybrid_key is not None: - print_rank_0(f"Original {hybrid_key}: {getattr(self.model, hybrid_key)}") - new_num_layers = self.model.config.num_layers - assert self.sorted_layers is not None - kept_layers_numbers = self.sorted_layers[:new_num_layers] - setattr( - self.model, - hybrid_key, - "".join( - c - for i, c in enumerate(getattr(self.model, hybrid_key)) - if i + 1 in kept_layers_numbers - ), - ) - print_rank_0(f"Pruned {hybrid_key}: {getattr(self.model, hybrid_key)}") + self._prune(export_config) print_mcore_model_stats( self.model, "Pruned Model", self.config["seq_length"], self.config["batch_size"] ) - def _prune(self, export_config: dict, prune_depth: bool = True) -> None: + @contextmanager + def _temporarily_pruned(self, export_config: dict): + """Prune to ``export_config`` for the duration of the block, then restore the max subnet. + + Used to score a candidate subnet without permanently mutating the model. ``_prune`` drops + layers and slices per-layer patterns (config lists + hybrid string) in place; ``sample(max)`` + only restores hparam-controlled state, so those must be snapshotted and restored here — else + the mutations compound across candidates and corrupt the final export. + """ + model = self.model + all_layers = model.decoder.layers + start_layer_number = all_layers[0].layer_number + hybrid_key = _get_hybrid_pattern_key(model) + saved_patterns = { + attr: getattr(model.config, attr) + for attr in _PER_LAYER_CONFIG_PATTERNS + if isinstance(getattr(model.config, attr, None), (list, tuple)) + } + saved_hybrid = getattr(model, hybrid_key) if hybrid_key else None + try: + self._prune(export_config) + yield + finally: + # reset to max subnet, revert dropped layers, and restore in-place-sliced patterns + sample(model, sample_func=max) + for offset, layer in enumerate(all_layers): + layer.layer_number = start_layer_number + offset + model.decoder.layers = all_layers + for attr, pattern in saved_patterns.items(): + setattr(model.config, attr, pattern) + if hybrid_key: + setattr(model, hybrid_key, saved_hybrid) + + def _prune(self, export_config: dict) -> None: """Prune the model homogeneously based on the export_config by setting active choices for configurable hparams. + Also drops layers and slices all per-layer patterns. + Args: export_config: Dictionary mapping hyperparameter names to their pruned values. - prune_depth: Whether to drop layers based on sorted_layers (default: True). """ # Prune homogeneously for n, hp in named_hparams(self.model, configurable=True): @@ -422,13 +467,12 @@ def _prune(self, export_config: dict, prune_depth: bool = True) -> None: if hp_name in export_config: hp.active = export_config[hp_name] - # Drop layers if depth pruning is enabled - if prune_depth: - num_layers_hp = self.model.get_hparam("num_layers") - if num_layers_hp.active != num_layers_hp.max: - assert self.sorted_layers is not None - layers_to_drop = self.sorted_layers[num_layers_hp.active :] - drop_mcore_language_model_layers(self.model, layers_to_drop=layers_to_drop) + # Drop layers based on sorted_layers + num_layers_hp = self.model.get_hparam("num_layers") + if num_layers_hp.active != num_layers_hp.max: + assert self.sorted_layers is not None + layers_to_drop = self.sorted_layers[num_layers_hp.active :] + drop_mcore_language_model_layers(self.model, layers_to_drop=layers_to_drop) # Update model config with pruned architecture # kv_channels can be None so we need to save from original hidden_size and num_attention_heads @@ -565,19 +609,9 @@ def search_best_arch_by_metrics(self) -> dict: smoothing=0.7, ): if candidate.score is None: # not restored from checkpoint - all_layers = self.model.decoder.layers - start_layer_number = all_layers[0].layer_number - - self._prune(candidate.ss_config, prune_depth=True) - candidate.score = self.eval_score(silent=False) - self.save_search_checkpoint(verbose=False) - - # reset to max subnet and revert dropped layers - sample(self.model, sample_func=max) - for layer in all_layers: - layer.layer_number = start_layer_number - start_layer_number += 1 - self.model.decoder.layers = all_layers + with self._temporarily_pruned(candidate.ss_config): + candidate.score = self.eval_score(silent=False) + self.save_search_checkpoint(verbose=False) metrics_str = ", ".join( f"{self._fmt_metric(v, k)} {k}" for k, v in candidate.metrics.items() ) diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_model_stats.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_model_stats.py index d5fe807f8dd..5277ae79420 100644 --- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_model_stats.py +++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_model_stats.py @@ -287,6 +287,16 @@ def test_gated_delta_net_layers(self): _BASE_CFG, _V, num_layers=4, num_moe_experts=None, **_GDN_OVERRIDES ) assert total == _BASE_UNTIED + 3 * (_GDN + _DENSE_MLP) + (_ATTN + _DENSE_MLP) + # Real Qwen3.5 configs store the pattern as an explicit per-layer list (1=linear/GDN, 0=full) + # instead of an int cadence; the equivalent list ([1,1,1,0] == freq 4 here) must match. + total_list, _ = mcore_param_count( + _BASE_CFG, + _V, + num_layers=4, + num_moe_experts=None, + **{**_GDN_OVERRIDES, "linear_attention_freq": [1, 1, 1, 0]}, + ) + assert total_list == total def test_gated_delta_net_differs_from_plain_attention(self): # Without the variant flag, the same layers are counted as plain attention (no GDN dispatch). diff --git a/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py b/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py index cb66b11683b..3b3425f6717 100644 --- a/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py +++ b/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py @@ -557,8 +557,14 @@ def _build_and_prune_variant(size, export_config, *, num_attention_heads=4, **mo def _test_mcore_qwen35_gdn_moe_pruning(rank, size): - # Qwen3.5-like hybrid MoE LM: GatedDeltaNet + gated full-attention layers. Only hidden_size and - # MoE dims are pruned; attention/GDN internal heads and the fused output gate are kept. + # Qwen3.5-like hybrid MoE LM: GatedDeltaNet + gated full-attention layers. hidden_size and MoE + # dims are width-pruned and num_layers is depth-pruned; attention/GDN internal heads and the fused + # output gate are kept. Depth pruning also slices the per-layer ``linear_attention_freq`` list. + num_layers = min(size * 2, 8) + pruned_num_layers = num_layers // 2 + # Store the linear-attention pattern in list form (as real Qwen3.5 configs do, unlike the int + # cadence default) so depth pruning must slice it: 1 = linear (GatedDeltaNet), 0 = full attention. + linear_attention_freq = [0 if (i + 1) % 2 == 0 else 1 for i in range(num_layers)] pruned_hidden, pruned_moe_ffn, pruned_moe_shared, pruned_experts = 32, 16, 32, 2 model = _build_and_prune_variant( size, @@ -567,12 +573,14 @@ def _test_mcore_qwen35_gdn_moe_pruning(rank, size): "moe_ffn_hidden_size": pruned_moe_ffn, "moe_shared_expert_intermediate_size": pruned_moe_shared, "num_moe_experts": pruned_experts, + "num_layers": pruned_num_layers, }, num_query_groups=2, experimental_attention_variant="gated_delta_net", num_moe_experts=4, moe_ffn_hidden_size=32, moe_shared_expert_intermediate_size=64, + linear_attention_freq=linear_attention_freq, ) for layer in model.decoder.layers: @@ -593,6 +601,10 @@ def _test_mcore_qwen35_gdn_moe_pruning(rank, size): assert model.config.moe_ffn_hidden_size == pruned_moe_ffn assert model.config.moe_shared_expert_intermediate_size == pruned_moe_shared assert model.config.num_moe_experts == pruned_experts + # Depth pruning must slice the per-layer ``linear_attention_freq`` list to the surviving layers + assert model.config.num_layers == pruned_num_layers + assert isinstance(model.config.linear_attention_freq, list) + assert len(model.config.linear_attention_freq) == pruned_num_layers def test_mcore_qwen35_gdn_moe_pruning(dist_workers): From d5e97d8446ae112d20adc89a1dd3ffcc35c87333 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 15 Jul 2026 13:07:59 -0700 Subject: [PATCH 5/8] docs(megatron_bridge): address README review comments Restructure the Distillation section: non-quantized methods (real-data / VLM) followed by their HF conversion, then QAD (quantized, exported via export_quantized) and Slurm; move Sanity-Check Generation last; add a workflows intro; rename MEGATRON_BRIDGE_SRC_DIR; and clarify the --student_hf_model, --trust_remote_code, and export --pp_size notes. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/README.md | 87 +++++++++++++++--------------- 1 file changed, 45 insertions(+), 42 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 2a3ff9558bf..f1a4cd3f126 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -1,6 +1,6 @@ # Megatron Bridge -This directory contains examples of using Model Optimizer with [NeMo Megatron-Bridge](https://github.com/NVIDIA-Nemo/Megatron-Bridge) framework for quantization, distillation, pruning, etc. +This directory contains examples of using Model Optimizer with the [NeMo Megatron-Bridge](https://github.com/NVIDIA-Nemo/Megatron-Bridge) framework for quantization, pruning, and distillation. These workflows can be used on their own or combined.
@@ -8,9 +8,9 @@ This directory contains examples of using Model Optimizer with [NeMo Megatron-Br | :------------: | :------------: | :------------: | | Pre-Requisites | Development environment setup | \[[Link](#pre-requisites)\] | | Post-Training Quantization | Quantizing a model | \[[Link](#post-training-quantization)\] | -| Sanity-Check Generation | Quick generation check with vLLM | \[[Link](#sanity-check-generation)\] | | Distillation | Distilling a pruned or quantized model | \[[Link](#distillation)\] | | Pruning | Pruning a model using Minitron algorithm | \[[Link](#pruning)\] | +| Sanity-Check Generation | Quick generation check with vLLM | \[[Link](#sanity-check-generation)\] | | Resources | Extra links to relevant resources | \[[Link](#resources)\] |
@@ -51,7 +51,7 @@ docker run \ > **Working with MoE Vision-Language Models (e.g. Qwen3.5-VL-MoE)?** The `nemo:26.06` container's Megatron-Bridge lacks the MoE expert weight mappings these models need (dense VLMs such as Gemma3-VL and Qwen3-VL work as-is). Until the `nemo:26.08` container is released, mount the latest [Megatron-Bridge `main`](https://github.com/NVIDIA-NeMo/Megatron-Bridge) source over the pre-installed copy by adding this to the `docker run` command above: > > ```bash -> -v ${MEGATRON_BRIDGE_DIR}:/opt/Megatron-Bridge +> -v ${MEGATRON_BRIDGE_SRC_DIR}:/opt/Megatron-Bridge > ``` You also need to login with your HuggingFace token to download gated datasets / models. @@ -96,7 +96,7 @@ torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \ ``` > [!NOTE] -> The HuggingFace unified exporter does not gather tensor-parallel-sharded weights. Use `--pp_size` on `export_quantized_megatron_to_hf.py` to shard a large model with pipeline parallelism across GPUs for export. +> The HuggingFace unified exporter can't split weights across GPUs with tensor parallelism. For large models, use `--pp_size` on `export_quantized_megatron_to_hf.py` to shard the export across GPUs with pipeline parallelism instead. > [!TIP] > To recover the accuracy lost during quantization, fine-tune the quantized Megatron checkpoint (from step 1) with [Quantization Aware Distillation (QAD)](#quantization-aware-distillation-qad) before running the step 2 export. @@ -113,18 +113,6 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati > [!NOTE] > HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM is not yet supported; the quantized VLM is saved in Megatron checkpoint format only. -## Sanity-Check Generation - -[generate_vllm.py](generate_vllm.py) runs a quick generation check on an exported HuggingFace checkpoint using vLLM — a useful smoke test for both a **quantized** and a **pruned** model to confirm it still produces coherent text before investing in [distillation](#distillation) to recover accuracy. For quantized checkpoints, vLLM auto-detects the ModelOpt quantization from the exported `hf_quant_config.json`, so no extra flags are needed: - -```bash -# Quantized model -python generate_vllm.py --model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 --trust_remote_code - -# Pruned model -python generate_vllm.py --model /tmp/Qwen3-8B-Pruned-6B -``` - ## Distillation This section shows how to distill a student model from a teacher model in the Megatron-Bridge framework. @@ -195,35 +183,21 @@ torchrun --nproc_per_node 8 distill.py \ --output_dir /tmp/test_distill ``` -### Quantization Aware Distillation (QAD) +### Vision-Language Models (VLMs) -To recover the accuracy lost during [Post-Training Quantization](#post-training-quantization), distill the quantized model (student) from the original, unquantized model (teacher). Pass the quantized **Megatron checkpoint** produced by `quantize.py` via `--student_megatron_path` (the ModelOpt quantizers are restored automatically, so distillation trains the fake-quantized student), while `--student_hf_path` provides the student architecture and `--teacher_hf_path` points to the original unquantized model. We also use a smaller learning rate for QAD: +For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `distill.py` distills only the **language model** (on text data) and leaves the vision tower and projector untouched — matching the pruning and quantization behavior. It composes with pruning and QAD (`--student_megatron_path`) exactly as for LLMs, and the HF export reuses `--student_hf_path` (no `--student_hf_model` needed). ```bash torchrun --nproc_per_node 8 distill.py \ --tp_size 8 \ - --teacher_hf_path Qwen/Qwen3-8B \ - --student_hf_path Qwen/Qwen3-8B \ - --student_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \ - --data_paths 1.0 tokenized_qwen3/data1_text_document 1.0 tokenized_qwen3/data2_text_document \ - --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \ - --seq_length 8192 \ - --gbs 768 \ - --train_iters 1000 \ - --lr 1e-5 \ - --min_lr 5e-6 \ - --output_dir /output/qwen3_8b_nvfp4_qad + --teacher_hf_path Qwen/Qwen3-VL-2B-Thinking \ + --student_hf_path Qwen/Qwen3-VL-2B-Thinking \ + ... ``` -The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). - -### Slurm Usage - -To run the distillation script on a Slurm cluster for multi-node training, you just need use `python` instead of `torchrun` and set the number of nodes using `#SBATCH --nodes=` clause in your Slurm script. - ### Converting to Hugging Face format (optional) -The distilled checkpoint is saved in Megatron distributed format. If you need a HuggingFace checkpoint, there are two ways to convert it: +A **non-quantized** distilled checkpoint (LLM or VLM) is saved in Megatron distributed format. If you need a HuggingFace checkpoint, there are two ways to convert it (for a **QAD** checkpoint, which retains quantization state, use [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) instead — see [QAD](#quantization-aware-distillation-qad)): **Inline** -- add `--hf_export_path` to the `distill.py` command to automatically convert the **final** checkpoint after distillation: @@ -233,7 +207,7 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \ --hf_export_path /path/to/save/distilled_hf_ckpt ``` -`--student_hf_model` is also needed for heterogeneous (Puzzletron/NAS) students whose architecture differs from `--student_hf_path`; for standard (homogeneous) models it defaults to `--student_hf_path`, and it is unused for VLMs. +`--student_hf_path` builds the student and provides the exported config / tokenizer. `--student_hf_model` is only needed when the student is **heterogeneous** (from Puzzletron/NAS pruning) and so differs from `--student_hf_path`; it points to a template matching that pruned architecture. For homogeneous models and VLMs, omit it -- it defaults to `--student_hf_path`. **Separate conversion** -- convert **any** saved iteration (intermediate or final) with [export_distilled_megatron_to_hf.py](export_distilled_megatron_to_hf.py): @@ -244,18 +218,32 @@ torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ --hf_export_path /path/to/save/distilled_hf_ckpt ``` -### Vision-Language Models (VLMs) +### Quantization Aware Distillation (QAD) -For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `distill.py` automatically distills only the **language model** (on text data) and leaves the vision tower and vision-language projector untouched, then saves the full VLM back. This also aligns with pruning and quantization behavior where we leave the vision tower and projector untouched. This composes with pruning and QAD (`--student_megatron_path`) exactly as for LLMs, and the HF export reuses `--student_hf_path` (no `--student_hf_model` needed). +To recover the accuracy lost during [Post-Training Quantization](#post-training-quantization), distill the quantized model (student) from the original, unquantized model (teacher). Pass the quantized **Megatron checkpoint** produced by `quantize.py` via `--student_megatron_path` (the ModelOpt quantizers are restored automatically, so distillation trains the fake-quantized student), while `--student_hf_path` provides the student architecture and `--teacher_hf_path` points to the original unquantized model. We also use a smaller learning rate for QAD: ```bash torchrun --nproc_per_node 8 distill.py \ --tp_size 8 \ - --teacher_hf_path Qwen/Qwen3-VL-2B-Thinking \ - --student_hf_path Qwen/Qwen3-VL-2B-Thinking \ - ... + --teacher_hf_path Qwen/Qwen3-8B \ + --student_hf_path Qwen/Qwen3-8B \ + --student_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \ + --data_paths 1.0 tokenized_qwen3/data1_text_document 1.0 tokenized_qwen3/data2_text_document \ + --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \ + --seq_length 8192 \ + --gbs 768 \ + --train_iters 1000 \ + --lr 1e-5 \ + --min_lr 5e-6 \ + --output_dir /output/qwen3_8b_nvfp4_qad ``` +The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/output/qwen3_8b_nvfp4_qad/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). + +### Slurm Usage + +To run the distillation script on a Slurm cluster for multi-node training, you just need use `python` instead of `torchrun` and set the number of nodes using `#SBATCH --nodes=` clause in your Slurm script. + ### Distillation Results See [examples/pruning/](../pruning/README.md#tutorials--results) for distillation experiment results covering Minitron and Puzzletron pruning algorithms. @@ -356,6 +344,21 @@ torchrun --nproc_per_node 2 prune_minitron.py \ --output_hf_path /tmp/Qwen3.5-4B-Pruned-3B ``` +## Sanity-Check Generation + +[generate_vllm.py](generate_vllm.py) runs a quick generation check on an exported HuggingFace checkpoint using vLLM — a useful smoke test for a **quantized**, **pruned**, or **distilled** model to confirm it still produces coherent text. For quantized checkpoints, vLLM auto-detects the ModelOpt quantization from the exported `hf_quant_config.json`, so no extra flags are needed: + +```bash +# Quantized model +python generate_vllm.py --model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 --trust_remote_code + +# Pruned model +python generate_vllm.py --model /tmp/Qwen3-8B-Pruned-6B +``` + +> [!NOTE] +> `--trust_remote_code` is only needed for models that ship custom modeling code (e.g. Nemotron); Qwen models don't require it. + ## Resources - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699) From f0787db16e9d50e1e19f1ce4080895cb2c6fd28e Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 15 Jul 2026 13:15:52 -0700 Subject: [PATCH 6/8] fix(prune): restore full layer list before sampling max subnet in _temporarily_pruned _prune sets width hparams on all layers before dropping some, so dropped layers retain pruned per-layer widths. Reattach the full layer list before sample(max) so every layer -- including previously-dropped ones -- is reset to the max subnet instead of being reattached in a stale state. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/prune/plugins/mcore_minitron.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index df61713b480..f5f3fcede3c 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -443,11 +443,14 @@ def _temporarily_pruned(self, export_config: dict): self._prune(export_config) yield finally: - # reset to max subnet, revert dropped layers, and restore in-place-sliced patterns - sample(model, sample_func=max) + # Reattach the full layer list (with original numbering) BEFORE sampling the max subnet, + # so sample(max) reaches every layer -- including the dropped ones -- and resets their + # per-layer width hparams (otherwise the detached layers are reattached still holding + # their pruned widths). Then restore the in-place-sliced patterns (not hparam-controlled). for offset, layer in enumerate(all_layers): layer.layer_number = start_layer_number + offset model.decoder.layers = all_layers + sample(model, sample_func=max) for attr, pattern in saved_patterns.items(): setattr(model.config, attr, pattern) if hybrid_key: From c3f770e9c8e1bcfed3041c7e2fe48e21a76f25e5 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 16 Jul 2026 07:06:11 -0700 Subject: [PATCH 7/8] docs(megatron_bridge): clarify --student_hf_model and rename run_provider Reframe --student_hf_model as a reference HF model with a homogeneous architecture used as the export template for a heterogeneous (Puzzletron/NAS) student (arg help in distill.py / export_distilled_megatron_to_hf.py, the export_llm_to_hf docstring, and the README). Rename the distill.py local run_provider -> distill_provider for clarity. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/README.md | 2 +- examples/megatron_bridge/distill.py | 18 +++++++++--------- .../export_distilled_megatron_to_hf.py | 12 ++++++------ 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index f1a4cd3f126..e69af367a5b 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -207,7 +207,7 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \ --hf_export_path /path/to/save/distilled_hf_ckpt ``` -`--student_hf_path` builds the student and provides the exported config / tokenizer. `--student_hf_model` is only needed when the student is **heterogeneous** (from Puzzletron/NAS pruning) and so differs from `--student_hf_path`; it points to a template matching that pruned architecture. For homogeneous models and VLMs, omit it -- it defaults to `--student_hf_path`. +`--student_hf_path` builds the student and provides the exported config / tokenizer. `--student_hf_model` is a reference HF model with a **homogeneous** architecture, used as the export template only for **heterogeneous** (Puzzletron/NAS) students; for homogeneous models and VLMs, omit it -- it defaults to `--student_hf_path`. **Separate conversion** -- convert **any** saved iteration (intermediate or final) with [export_distilled_megatron_to_hf.py](export_distilled_megatron_to_hf.py): diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 947a1387e57..856aacb57e5 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -188,9 +188,9 @@ def get_args(): type=str, required=False, default=None, - help="Export template: HF model whose architecture matches the exported student. Only needed " - "for heterogeneous (Puzzletron/NAS) students whose architecture differs from --student_hf_path. " - "Defaults to --student_hf_path; unused for VLMs.", + 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.", ) args = parser.parse_args() @@ -261,7 +261,7 @@ def _build_model_provider(hf_path, load_weights=True): "VLM detected: distilling model.language_model only (vision tower / projector untouched). " "To export megatron non-quantized checkpoint, use export_distilled_megatron_to_hf.py" ) - run_provider = convert_to_distillation_provider( + distill_provider = convert_to_distillation_provider( student_provider, teacher_provider, kd_config, @@ -286,7 +286,7 @@ def _restore_student_hook(model_chunks): ) return model_chunks - run_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) + distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) # Build optimizer and scheduler optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing( @@ -318,7 +318,7 @@ def _restore_student_hook(model_chunks): # Assemble ConfigContainer and run distillation config = ConfigContainer( - model=run_provider, + model=distill_provider, train=TrainingConfig( train_iters=args.train_iters, eval_interval=args.eval_interval, @@ -351,7 +351,7 @@ def _restore_student_hook(model_chunks): wandb_exp_name=args.wandb_exp_name, ), tokenizer=TokenizerConfig( - tokenizer_type="NullTokenizer", vocab_size=run_provider.vocab_size + tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size ), checkpoint=CheckpointConfig( save_interval=args.eval_interval, @@ -377,8 +377,8 @@ def _restore_student_hook(model_chunks): # 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. - run_provider.initialize_model_parallel(seed=args.seed) - full_student = run_provider.full_model + 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( diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index ffa41825463..6daf041e575 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -71,9 +71,9 @@ def export_llm_to_hf( megatron_path: Megatron checkpoint directory (an ``iter_*`` dir or its parent). hf_export_path: Directory to write the HuggingFace checkpoint to. student_hf_path: Student HF model used for the exported config / tokenizer. - template_hf: HF model whose architecture matches the exported student. Defaults to - ``student_hf_path``; pass a distinct value only for heterogeneous (Puzzletron/NAS) - students whose architecture differs from ``student_hf_path``. + template_hf: Reference HF model with a homogeneous architecture, used as the export template + for a heterogeneous (Puzzletron/NAS) student. Defaults to ``student_hf_path`` (correct for + homogeneous students). trust_remote_code: Whether to trust remote code when loading the HF model. """ export_bridge = AutoBridge.from_hf_pretrained( @@ -140,9 +140,9 @@ def get_args() -> argparse.Namespace: "--student_hf_model", type=str, default=None, - help="Export template: HF model whose architecture matches the exported student. Only needed " - "for heterogeneous (Puzzletron/NAS) students whose architecture differs from --student_hf_path. " - "Defaults to --student_hf_path; unused for VLMs.", + 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.", ) parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code") parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size") From 5f65a35215bc180521213ecc61377929f6015d14 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 16 Jul 2026 07:19:28 -0700 Subject: [PATCH 8/8] docs(megatron_bridge): note VLM naming convention and in-memory export TODO Per review follow-ups: document the vision_config / language_model VLM-detection convention in distill.py, and add a TODO in export_llm_to_hf to unify the LLM export onto the in-memory path used for VLMs. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/distill.py | 4 ++++ examples/megatron_bridge/export_distilled_megatron_to_hf.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 856aacb57e5..dfd404fab30 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -251,6 +251,10 @@ def _build_model_provider(hf_path, load_weights=True): 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", diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index 6daf041e575..bb8726463ed 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -76,6 +76,10 @@ def export_llm_to_hf( homogeneous students). trust_remote_code: Whether to trust remote code when loading the HF model. """ + # TODO: unify with save_vlm_to_hf's in-memory export path. This LLM path re-loads the checkpoint + # from disk via export_ckpt (which reads the actual per-layer shapes, so it handles heterogeneous + # Puzzletron/NAS students); an in-memory export would need to rebuild the (possibly heterogeneous) + # student first. export_bridge = AutoBridge.from_hf_pretrained( template_hf or student_hf_path, trust_remote_code=trust_remote_code )