diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..7dd96e6ac69 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 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``. @@ -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..e69af367a5b 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)\] |
@@ -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 ${MEGATRON_BRIDGE_SRC_DIR}:/opt/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 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. -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,15 +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. - -## 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: - -```bash -python generate_vllm.py --model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 --trust_remote_code -``` +> 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. ## Distillation @@ -184,6 +183,41 @@ torchrun --nproc_per_node 8 distill.py \ --output_dir /tmp/test_distill ``` +### Vision-Language Models (VLMs) + +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-VL-2B-Thinking \ + --student_hf_path Qwen/Qwen3-VL-2B-Thinking \ + ... +``` + +### Converting to Hugging Face format (optional) + +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: + +```bash +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 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): + +```bash +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 +``` + ### Quantization Aware Distillation (QAD) 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: @@ -204,38 +238,12 @@ 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 `/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. -### 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: - -**Inline** -- add `--hf_export_path` and `--student_hf_model` 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 -``` - -`--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`. - -**Separate conversion** -- convert any saved iteration using the Megatron-Bridge conversion script: - -```bash -uv run python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ - --hf-model \ - --megatron-path /checkpoints/iter_ \ - --hf-path -``` - -For more details, see the [Megatron-Bridge conversion README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion). - ### Distillation Results See [examples/pruning/](../pruning/README.md#tutorials--results) for distillation experiment results covering Minitron and Puzzletron pruning algorithms. @@ -336,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) diff --git a/examples/megatron_bridge/_distillation_provider.py b/examples/megatron_bridge/_distillation_provider.py new file mode 100644 index 00000000000..937d3254ae0 --- /dev/null +++ b/examples/megatron_bridge/_distillation_provider.py @@ -0,0 +1,108 @@ +# 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`` 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 +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). +""" + +import inspect + +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 _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 + 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 + + +# 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 283504f1ec2..dfd404fab30 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="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() @@ -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) @@ -347,23 +247,50 @@ 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 ) + + # VLM detection convention: HF VLM configs expose a ``vision_config``, and Megatron-Bridge nests + # the text model under the ``language_model`` submodule (used as ``distill_submodule`` below). If a + # future model breaks either convention, the ``getattr(model, "language_model")`` in the provider + # will error loudly rather than silently distilling the wrong module. + is_vlm = hasattr( + AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), + "vision_config", + ) + + if is_vlm: + warn_rank_0( + "VLM detected: distilling model.language_model only (vision tower / projector untouched). " + "To export megatron non-quantized checkpoint, use export_distilled_megatron_to_hf.py" + ) distill_provider = convert_to_distillation_provider( - student_provider, teacher_provider, kd_config + 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 + + 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( @@ -450,7 +377,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. + distill_provider.initialize_model_parallel(seed=args.seed) + full_student = distill_provider.full_model + # Strip the distillation wrapper -> plain trained language model (in place; reassign to be safe). + full_student.language_model = mtd.export(full_student.language_model) + save_vlm_to_hf( + full_student, + args.hf_export_path, + args.student_hf_path, + trust_remote_code=args.trust_remote_code, + ) + print_rank_0(f"Saved distilled VLM to {args.hf_export_path} in HF format") + elif args.hf_export_path: 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 +402,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..bb8726463ed --- /dev/null +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -0,0 +1,226 @@ +# 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: 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. + """ + # 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 + ) + 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="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") + 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/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index bdd40b3ad8f..5bead2491ba 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" @@ -544,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" @@ -567,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/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index d57e74277c6..3454da00441 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 @@ -413,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.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/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/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 * 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..f5f3fcede3c 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,55 @@ 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: + # 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: + 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 +470,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 +612,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/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/_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_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..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, @@ -109,6 +110,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..6ccf97a52a2 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -17,28 +17,71 @@ 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_qwen3_dir +from _test_utils.torch.transformers_models import ( + create_tiny_gemma3vl_dir, + create_tiny_qwen3_5_moe_vl_dir, + create_tiny_qwen3_dir, +) -@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" +@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 +90,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 +116,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" - ) 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):