diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7308931d772..181828a87ca 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,11 +21,14 @@ Changelog - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. - Add a fused Triton fast path for ``local_hessian`` NVFP4 weight-scale search (the Hessian-weighted FP8-E4M3 scale sweep). For each NVFP4 block it minimizes ``dwᵀ H dw`` over the 126 candidate scales using the per-cin-block local Hessian on tensor cores, replacing the per-weight Python reference sweep — roughly **34x** faster on a single 8192x4096 weight and bit-exact with the reference for fp32/fp16 weights. Used automatically during ``local_hessian`` calibration for both dense and fused-MoE expert weights; falls back to the reference sweep on CPU, when Triton is unavailable, or via ``MODELOPT_NVFP4_TRITON_SWEEP=0``. +- Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. - Add Minitron pruning support for Megatron-Core models with the following new attention and MoE variants. For these, only ``hidden_size`` is pruned (alongside the usual ``ffn_hidden_size`` / ``num_layers`` / MoE dimensions); the variant-internal dimensions noted below are not pruned: - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned. - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned. - **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 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``. @@ -34,7 +37,6 @@ Changelog - The exported state_dict is also **reordered (decoder keys win instead of encoder)** so canonical-side keys per HF's ``_tied_weights_keys`` declaration win the data_ptr dedup; gated to the DiffusionGemma model class in ``_reorder_canonical_first``, no-op for every other model. - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes. - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. -- Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. **Bug Fixes** diff --git a/docs/source/guides/3_pruning.rst b/docs/source/guides/3_pruning.rst index 39142ad78e2..a65890625ba 100644 --- a/docs/source/guides/3_pruning.rst +++ b/docs/source/guides/3_pruning.rst @@ -16,9 +16,10 @@ These pruning methods support pruning the convolutional and linear layers, and attention heads of the model. More details on these pruning modes is as follows: #. ``mcore_minitron``: A pruning method developed by NVIDIA Research for pruning GPT, Mamba and Hybrid - Transformer Mamba models in NVIDIA Megatron-Bridge or Megatron-LM framework. It uses the activation magnitudes to prune + Transformer Mamba models (including MoE, and the language model of vision-language models) in NVIDIA + Megatron-Bridge or Megatron-LM framework. It uses the activation magnitudes to prune the embedding hidden size, mlp ffn hidden size, transformer attention heads, GQA query groups, - mamba heads and head dimension, and number of layers of the model. + mamba heads and head dimension, MoE experts, and number of layers of the model. Checkout more details of the algorithm in the `paper `_. #. ``puzzletron``: An advanced LLM/VLM pruning method by NVIDIA using Mixed Integer Programming (MIP) based NAS search algorithm. #. ``fastnas``: A pruning method recommended for Computer Vision models. Given a pretrained model, diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6b7a5d773a6..959316233fb 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -266,7 +266,6 @@ def make_calib_dataloader( batch_size=args.batch_size, num_samples=args.calib_size[0], device=device, - max_length=args.calib_seq, require_image=True, subsets=["sparsetables", "plotqa_cot", "wiki_en"], shuffle_buffer_size=10_000, diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 3f4cc17cd13..60dfed7fc7d 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -96,7 +96,15 @@ torchrun --nproc_per_node 2 export.py \ To see the full usage for advanced configurations, run `torchrun --nproc_per_node 1 quantize.py --help` (or `export.py --help`). -For VLM (vision-language model) quantization, see the Megatron-Bridge repository [here](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/quantization). +### Vision-Language Models (VLMs) + +For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automatically quantizes only the **language model** and leaves the vision tower and vision-language projector in full precision, then saves the full VLM back as a Megatron checkpoint. The calibration modality is inferred from `--calib_dataset_name`: + +- An **image-text** dataset (the default for VLMs, `nemotron_vlm_dataset_v2`) drives the full VLM forward, so the language model is calibrated on vision-conditioned activations. +- 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 @@ -307,9 +315,27 @@ torchrun --nproc_per_node 1 prune_minitron.py --help > [!NOTE] > NAS-based pruning requires ~2x the GPU memory of Manual pruning because it needs to simultaneously hold original model while evaluating each pruned candidate. +> [!NOTE] +> Multi-token-prediction (MTP) heads (e.g. Qwen3.5) are not pruned yet — they are dropped for the prune run and the saved checkpoint has no MTP. Autoregressive inference is unaffected; for speculative decoding, run a short MTP SFT on the pruned model. + > [!NOTE] > If pruning a Nemotron model and you want to save the pruned model back in HF format, please downgrade to `transformers<5` via `python -m pip install "transformers<5"` before pruning. +### Vision-Language Models (VLMs) + +For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `prune_minitron.py` automatically prunes only the **language model** and leaves the vision tower intact, then saves the full VLM back. All the pruning modes above (parameter count, active parameter count, memory footprint, and manual `export_config`) work unchanged, with two VLM-specific caveats: + +- The `--prune_target_params` / `--prune_target_active_params` / `--prune_target_memory_mb` targets (and `export_config` dimensions) apply to the **language model only** — the (unpruned) vision tower's parameters are *not* counted, so the full saved VLM will be larger than the target. +- `hidden_size` is never pruned for VLMs (it is shared with the vision projector). + +```bash +torchrun --nproc_per_node 2 prune_minitron.py \ + --pp_size 2 \ + --hf_model_name_or_path Qwen/Qwen3.5-4B \ + --prune_target_params 3e9 \ + --output_hf_path /tmp/Qwen3.5-4B-Pruned-3B +``` + ## Resources - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699) diff --git a/examples/megatron_bridge/export.py b/examples/megatron_bridge/export.py index eecf7d838f4..926b47755e5 100644 --- a/examples/megatron_bridge/export.py +++ b/examples/megatron_bridge/export.py @@ -144,6 +144,9 @@ def main(args: argparse.Namespace): print_rank_0( f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." ) + # TODO (OMNIML-5366): quantized-VLM HF export. export_mcore_gpt_to_hf's per-arch mappings don't + # cover Qwen3.5-VL / Gemma3-VL; See if Megatron-Bridge's AutoBridge.export_hf_weights_quant can be + # used instead. export_mcore_gpt_to_hf( unwrapped_model, args.hf_model_name_or_path, diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index bc3401e296f..bdd40b3ad8f 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -46,16 +46,47 @@ import torch from megatron.bridge import AutoBridge from megatron.bridge.models.mamba.mamba_provider import MambaModelProvider -from transformers import AutoConfig, AutoModelForCausalLM +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoProcessor, +) import modelopt.torch.opt as mto import modelopt.torch.prune as mtp import modelopt.torch.utils.distributed as dist from modelopt.torch.export import copy_hf_ckpt_remote_code -from modelopt.torch.utils import get_supported_datasets, print_args, print_rank_0, warn_rank_0 +from modelopt.torch.utils import ( + get_supported_datasets, + num2hrb, + print_args, + print_rank_0, + warn_rank_0, +) from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf -from modelopt.torch.utils.plugins.megatron_calibration import get_megatron_calibration_forward_loop +from modelopt.torch.utils.plugins.megatron_calibration import ( + get_megatron_calibration_forward_loop, + get_megatron_vlm_calibration_forward_loop, +) from modelopt.torch.utils.plugins.megatron_mmlu import megatron_mmlu +from modelopt.torch.utils.vlm_dataset_utils import get_supported_vlm_datasets + +# 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" + +# HF config field names that enable MTP +_MTP_HF_CONFIG_FIELDS = ("num_nextn_predict_layers", "mtp_num_hidden_layers", "mtp_num_layers") + + +def _hf_config_has_mtp(hf_cfg) -> bool: + """Whether an HF config declares MTP heads (checked top-level and under ``text_config``).""" + return any( + cfg is not None and getattr(cfg, field, 0) + for cfg in (getattr(hf_cfg, "text_config", None), hf_cfg) + for field in _MTP_HF_CONFIG_FIELDS + ) def get_args() -> argparse.Namespace: @@ -92,10 +123,13 @@ def get_args() -> argparse.Namespace: parser.add_argument( "--calib_dataset_name", type=str, - default="nemotron-post-training-dataset-v2", + default=None, help=( - f"HF Dataset name or local path for calibration (supported options: {', '.join(get_supported_datasets())}. " - "You can also pass any other dataset and see if auto-detection for your dataset works." + "Calibration dataset. If unset, it is auto-selected by model type: a text dataset " + f"({DEFAULT_TEXT_CALIB_DATASET}) for language models, and an image-text dataset " + f"({DEFAULT_VLM_CALIB_DATASET}) for VLMs. Passing a text dataset for a VLM estimates importance from text " + f"only. Text dataset options: {get_supported_datasets()}; VLM (image) dataset options: " + f"{get_supported_vlm_datasets()}." ), ) parser.add_argument( @@ -103,7 +137,12 @@ def get_args() -> argparse.Namespace: ) # TODO: Add support for pre-training dataset (pre-tokenized) parser.add_argument("--calib_batch_size", type=int, default=1, help="Calibration batch size") - parser.add_argument("--seq_length", type=int, default=4096) + parser.add_argument( + "--seq_length", + type=int, + default=4096, + help="Calibration sequence length (text only; ignored for image-text VLM calibration).", + ) # Pruning parameters parser.add_argument( "--prune_intermediate_ckpt", @@ -130,7 +169,8 @@ def get_args() -> argparse.Namespace: help=( "Target total parameter count e.g., 6e9 for 6B params. " "Uses NAS to find the best pruned model that maximizes --prune_score_func. " - "Can be combined with --prune_target_active_params and/or --prune_target_memory_mb." + "Can be combined with --prune_target_active_params and/or --prune_target_memory_mb. " + "For VLMs this targets the language-model tower only." ), ) parser.add_argument( @@ -139,7 +179,8 @@ def get_args() -> argparse.Namespace: help=( "Target active parameter count e.g., 3e9 for 3B active params (useful for MoE models). " "Uses NAS to find the best pruned model that maximizes --prune_score_func. " - "Can be combined with --prune_target_params and/or --prune_target_memory_mb." + "Can be combined with --prune_target_params and/or --prune_target_memory_mb. " + "For VLMs this targets the language-model tower only." ), ) parser.add_argument( @@ -149,7 +190,8 @@ def get_args() -> argparse.Namespace: "Target memory footprint in MB (weights + KV-cache estimated via seq_length and " "--inference_batch_size; assumes BF16). " "Uses NAS to find the best pruned model that maximizes --prune_score_func. " - "Can be combined with --prune_target_params and/or --prune_target_active_params." + "Can be combined with --prune_target_params and/or --prune_target_active_params. " + "For VLMs this targets the language-model tower only." ), ) parser.add_argument( @@ -264,6 +306,42 @@ def get_args() -> argparse.Namespace: return args +def _log_vlm_param_breakdown(unwrapped_model, language_model, stage: str) -> None: + """Log language-model / frozen-non-LM / total param counts for a VLM (rank 0).""" + + def _local(module) -> int: + # De-dup weights shared within a rank (e.g. tied embedding/output on a single stage). + seen: set[int] = set() + n = 0 + for p in module.parameters(): + if id(p) not in seen: + seen.add(id(p)) + n += p.numel() + return n + + total = dist.allreduce(_local(unwrapped_model)) # sum across pipeline ranks + lm = dist.allreduce(_local(language_model)) + # Under PP a tied embedding lives on both the first and last stage, so the sum double-counts it; + # subtract one copy (the allreduce over the first-stage-only ``word_embeddings`` gives exactly one). + if dist.size() > 1 and getattr(language_model, "share_embeddings_and_output_weights", False): + emb = dist.allreduce( + next( + ( + p.numel() + for n, p in unwrapped_model.named_parameters() + if "word_embeddings" in n + ), + 0, + ) + ) + total -= emb + lm -= emb + print_rank_0( + f"[{stage}] language_model={num2hrb(lm)} (--prune_target_* applies here) | " + f"frozen non-language-model={num2hrb(total - lm)} | full model={num2hrb(total)}" + ) + + def main(args: argparse.Namespace): assert dist.size() == args.pp_size, "Only Pipeline parallelism is supported for pruning." @@ -287,19 +365,84 @@ def main(args: argparse.Namespace): "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage, "pipeline_dtype": torch.bfloat16, "seq_length": args.seq_length, + "mtp_num_layers": 0, # MTP is not supported during calibration }, init_model_parallel=True, moe_grouped_gemm=False, ) - forward_loop = get_megatron_calibration_forward_loop( - tokenizer, - dataset_name=args.calib_dataset_name, - num_samples=args.calib_num_samples, - seq_length=args.seq_length, - batch_size=args.calib_batch_size, - # pack=True uses Megatron pretraining-style global-stream document packing - pack=True, - ) + + # TODO: Support pruning with MTP heads enabled (e.g. Qwen3.5 mtp_num_hidden_layers=1). + # Requires ModelOpt fixes for gated-attention QKV under DynamicModule during MTP calibration, + # _DynamicMCoreLanguageModel conversion/export of MTP submodules, importance hooks on MTP + # layers, mcore_param_count including MTP in --prune_target_params, and a CI test with MTP. + if _hf_config_has_mtp(bridge.hf_pretrained.config): + warn_rank_0( + "Dropping Multi-Token Prediction (MTP): calibration does not yet support MTP. Exported " + "checkpoints will not contain MTP weights. Standard autoregressive inference is unaffected. To use " + "MTP speculative decoding later, run a separate SFT phase with mtp_num_layers=1 on the pruned model." + ) + + # For VLMs (e.g. Qwen3-VL), only the language model is pruned; the vision tower is left intact. + # hidden_size is shared with the vision->LM projector, so it is skipped + language_model = getattr(unwrapped_model, "language_model", unwrapped_model) + is_vlm = language_model is not unwrapped_model + if is_vlm: + warn_rank_0( + "VLM detected: pruning model.language_model only; all non-language-model components " + "(vision/audio encoders, projectors, etc.) are frozen and excluded. --prune_target_* " + "applies to the language-model tower, not the full model (hidden_size pruning is also " + "skipped -- it is shared with the projector)." + ) + if args.prune_export_config and "hidden_size" in args.prune_export_config: + raise ValueError( + "Pruning 'hidden_size' is not supported for VLMs (shared with the vision projector)." + ) + args.hparams_to_skip = sorted({*args.hparams_to_skip, "hidden_size"}) + _log_vlm_param_breakdown(unwrapped_model, language_model, "before pruning") + + # Auto-select the calibration dataset by model type when not explicitly provided. + if args.calib_dataset_name is None: + args.calib_dataset_name = ( + DEFAULT_VLM_CALIB_DATASET if is_vlm else DEFAULT_TEXT_CALIB_DATASET + ) + + # Infer the calibration modality from the dataset: the known image-text datasets require a VLM, everything + # else is text. Passing a text dataset for a VLM estimates importance from text only (vision tower idle). + use_image_calib = args.calib_dataset_name in get_supported_vlm_datasets() + if use_image_calib and not is_vlm: + raise ValueError( + f"Calibration dataset '{args.calib_dataset_name}' is image-text and requires a VLM; " + "pass a text dataset for a language model." + ) + if is_vlm and not use_image_calib: + warn_rank_0( + f"Text-only calibration on a VLM (dataset '{args.calib_dataset_name}'): the language " + "model's pruning importance will not see vision tokens." + ) + print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}") + + # Estimate pruning importance for the language model: text-only on the LM for text datasets, or + # the full VLM forward over image-text pairs. + if use_image_calib: + processor = AutoProcessor.from_pretrained( + args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code + ) + forward_loop = get_megatron_vlm_calibration_forward_loop( + unwrapped_model, # full VLM (vision encoder + projector + language model) + processor, + dataset_name=args.calib_dataset_name, + num_samples=args.calib_num_samples, + batch_size=args.calib_batch_size, + ) + else: + forward_loop = get_megatron_calibration_forward_loop( + tokenizer, + dataset_name=args.calib_dataset_name, + num_samples=args.calib_num_samples, + seq_length=args.seq_length, + batch_size=args.calib_batch_size, + pack=True, # Megatron pretraining-style global-stream document packing + ) pruning_config = { "forward_loop": forward_loop, @@ -379,16 +522,20 @@ def score_func(m): pruning_config["seq_length"] = args.seq_length print_rank_0(f"Pruning constraints: {pruning_constraints}") - unwrapped_model, pruning_scores = mtp.prune( # in-place pruning - unwrapped_model, + # Prune the language model in place (for VLMs this mutates unwrapped_model.language_model, so the + # full wrapper is still saved below); for plain LMs language_model is unwrapped_model itself. + language_model, pruning_scores = mtp.prune( # in-place pruning + language_model, mode=[("mcore_minitron", ss_config)], # type: ignore[arg-type] constraints=pruning_constraints, dummy_input=None, config=pruning_config, ) # Remove unnecessary modelopt_state since ckpt is homogeneous - if mto.ModeloptStateManager.has_state_for_mode_type("prune", model=unwrapped_model): - mto.ModeloptStateManager.remove_state(unwrapped_model) + if mto.ModeloptStateManager.has_state_for_mode_type("prune", model=language_model): + mto.ModeloptStateManager.remove_state(language_model) + if is_vlm: + _log_vlm_param_breakdown(unwrapped_model, language_model, "after pruning") if isinstance(provider, MambaModelProvider): hybrid_key = ( "hybrid_override_pattern" @@ -425,42 +572,78 @@ def score_func(m): hf_cfg = AutoConfig.from_pretrained( args.output_hf_path, trust_remote_code=args.trust_remote_code ) - mcore_cfg = unwrapped_model.config - - hf_cfg.hidden_size = mcore_cfg.hidden_size - hf_cfg.intermediate_size = mcore_cfg.ffn_hidden_size - hf_cfg.num_attention_heads = mcore_cfg.num_attention_heads - hf_cfg.head_dim = mcore_cfg.kv_channels - hf_cfg.num_key_value_heads = mcore_cfg.num_query_groups - if hasattr(hf_cfg, "mamba_num_heads"): - hf_cfg.mamba_num_heads = mcore_cfg.mamba_num_heads - if hasattr(hf_cfg, "mamba_head_dim"): - hf_cfg.mamba_head_dim = mcore_cfg.mamba_head_dim - if hasattr(hf_cfg, "moe_intermediate_size"): - hf_cfg.moe_intermediate_size = mcore_cfg.moe_ffn_hidden_size - if hasattr(hf_cfg, "moe_shared_expert_intermediate_size"): - hf_cfg.moe_shared_expert_intermediate_size = ( - mcore_cfg.moe_shared_expert_intermediate_size - ) - if hasattr(hf_cfg, "num_experts"): - hf_cfg.num_experts = mcore_cfg.num_moe_experts - if hasattr(hf_cfg, "n_routed_experts"): - hf_cfg.n_routed_experts = mcore_cfg.num_moe_experts - if hasattr(hf_cfg, "n_shared_experts"): - hf_cfg.n_shared_experts = ( + mcore_cfg = language_model.config + # For VLMs the language-model fields live under hf_cfg.text_config; write back there. + text_cfg = getattr(hf_cfg, "text_config", hf_cfg) + + text_cfg.hidden_size = mcore_cfg.hidden_size + text_cfg.intermediate_size = mcore_cfg.ffn_hidden_size + text_cfg.num_attention_heads = mcore_cfg.num_attention_heads + text_cfg.head_dim = mcore_cfg.kv_channels + text_cfg.num_key_value_heads = mcore_cfg.num_query_groups + if hasattr(text_cfg, "mamba_num_heads"): + text_cfg.mamba_num_heads = mcore_cfg.mamba_num_heads + if hasattr(text_cfg, "mamba_head_dim"): + text_cfg.mamba_head_dim = mcore_cfg.mamba_head_dim + if hasattr(text_cfg, "moe_intermediate_size"): + text_cfg.moe_intermediate_size = mcore_cfg.moe_ffn_hidden_size + # HF names this field with or without the ``moe_`` prefix depending on the model + # (e.g. Qwen3.5-MoE uses ``shared_expert_intermediate_size``). + for shared_expert_field in ( + "moe_shared_expert_intermediate_size", + "shared_expert_intermediate_size", + ): + if hasattr(text_cfg, shared_expert_field): + setattr( + text_cfg, shared_expert_field, mcore_cfg.moe_shared_expert_intermediate_size + ) + if hasattr(text_cfg, "num_experts"): + text_cfg.num_experts = mcore_cfg.num_moe_experts + if hasattr(text_cfg, "n_routed_experts"): + text_cfg.n_routed_experts = mcore_cfg.num_moe_experts + if hasattr(text_cfg, "n_shared_experts"): + text_cfg.n_shared_experts = ( mcore_cfg.moe_shared_expert_intermediate_size // mcore_cfg.moe_ffn_hidden_size ) - if hasattr(hf_cfg, "layer_types"): - kept_layer_nums = pruning_scores["sorted_layers"][: mcore_cfg.num_layers] # 1-indexed - hf_cfg.layer_types = [ - lt for i, lt in enumerate(hf_cfg.layer_types) if i + 1 in kept_layer_nums + # Layers that survived depth pruning (1-indexed). sorted_layers is None when no layer scores + # were collected (no depth pruning) -> all layers kept. + sorted_layers = pruning_scores["sorted_layers"] + kept_layer_nums = ( + set(sorted_layers[: mcore_cfg.num_layers]) + if sorted_layers is not None + else set(range(1, mcore_cfg.num_layers + 1)) + ) + if hasattr(text_cfg, "layer_types"): + text_cfg.layer_types = [ + lt for i, lt in enumerate(text_cfg.layer_types) if i + 1 in kept_layer_nums + ] + # Qwen3-VL injects deepstack vision features at specific LM layers; remap those indices to the + # surviving layers (a dropped one snaps to the nearest survivor below; count is preserved). + vision_cfg = getattr(hf_cfg, "vision_config", None) + ds_indices = getattr(vision_cfg, "deepstack_visual_indexes", None) + if vision_cfg is not None and ds_indices: + kept_sorted = sorted(kept_layer_nums) + vision_cfg.deepstack_visual_indexes = [ + max(0, sum(k <= d + 1 for k in kept_sorted) - 1) for d in ds_indices ] + if any((d + 1) not in kept_layer_nums for d in ds_indices): + warn_rank_0( + "A deepstack vision-injection layer was dropped during depth pruning; its " + "feature was snapped to the nearest surviving layer. Text-only (LM) " + "distillation cannot recover this vision-path change -- consider full VLM " + "training/distillation instead of LM-only to recover vision quality." + ) if isinstance(provider, MambaModelProvider) and hasattr(hf_cfg, "hybrid_override_pattern"): hf_cfg.hybrid_override_pattern = getattr(unwrapped_model, hybrid_key) - hf_cfg.num_hidden_layers = mcore_cfg.num_layers + text_cfg.num_hidden_layers = mcore_cfg.num_layers + # Mark MTP as disabled on the HF text config written after pruning + for field in _MTP_HF_CONFIG_FIELDS: + if hasattr(text_cfg, field): + setattr(text_cfg, field, 0) # Save dummy pruned HF model to get the correct bridge for saving pruned weights - AutoModelForCausalLM.from_config( + dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM + dummy_model_cls.from_config( hf_cfg, trust_remote_code=args.trust_remote_code ).save_pretrained(args.output_hf_path, trust_remote_code=args.trust_remote_code) pruned_bridge = AutoBridge.from_hf_pretrained( diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index cf25f6bfc2d..d57e74277c6 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -61,6 +61,7 @@ import gc import torch +from transformers import AutoProcessor import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist @@ -69,8 +70,16 @@ from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 from modelopt.torch.utils.dataset_utils import get_supported_datasets from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf -from modelopt.torch.utils.plugins.megatron_calibration import get_megatron_calibration_forward_loop +from modelopt.torch.utils.plugins.megatron_calibration import ( + get_megatron_calibration_forward_loop, + get_megatron_vlm_calibration_forward_loop, +) from modelopt.torch.utils.plugins.megatron_generate import megatron_generate +from modelopt.torch.utils.vlm_dataset_utils import get_supported_vlm_datasets + +# Default calibration datasets when --calib_dataset_name is not set +DEFAULT_TEXT_CALIB_DATASET = "cnn_nemotron_v2_mix" # cnn_dailymail + nemotron-post-training-v2 +DEFAULT_VLM_CALIB_DATASET = "nemotron_vlm_dataset_v2" # The --quant_cfg / --kv_cache_quant CLI vocabularies are discovered from the preset # YAMLs (shared with the hf_ptq examples via modelopt.recipe.presets). --quant_cfg @@ -152,17 +161,25 @@ def get_args() -> argparse.Namespace: parser.add_argument( "--calib_dataset_name", type=str, - default="cnn_nemotron_v2_mix", # cnn_dailymail + nemotron-post-training-dataset-v2 + default=None, help=( - f"HF Dataset name or local path for calibration (supported options: {', '.join(get_supported_datasets())}. " - "You can also pass any other dataset and see if auto-detection for your dataset works." + "Calibration dataset. If unset, it is auto-selected by model type: a text dataset " + f"({DEFAULT_TEXT_CALIB_DATASET}) for language models, and an image-text dataset " + f"({DEFAULT_VLM_CALIB_DATASET}) for VLMs. Passing a text dataset for a VLM estimates importance from text " + f"only. Text dataset options: {get_supported_datasets()}; VLM (image) dataset options: " + f"{get_supported_vlm_datasets()}." ), ) parser.add_argument( "--calib_num_samples", type=int, default=1024, help="Number of samples for calibration" ) parser.add_argument("--calib_batch_size", type=int, default=1, help="Calibration batch size") - parser.add_argument("--seq_length", type=int, default=4096, help="Calibration sequence length") + parser.add_argument( + "--seq_length", + type=int, + default=4096, + help="Calibration sequence length (text only; ignored for image-text VLM calibration).", + ) # Post-quantization generation (sanity check) arguments parser.add_argument( @@ -274,8 +291,53 @@ def main(args: argparse.Namespace): init_model_parallel=True, ) + # Only the language model is quantized (vision tower + projector stay full precision) + language_model = getattr(unwrapped_model, "language_model", unwrapped_model) + is_vlm = language_model is not unwrapped_model + if is_vlm: + warn_rank_0( + "VLM detected: quantizing `model.language_model` only (vision tower left in full precision)." + ) + + # Auto-select the calibration dataset by model type when not explicitly provided. + if args.calib_dataset_name is None: + args.calib_dataset_name = ( + DEFAULT_VLM_CALIB_DATASET if is_vlm else DEFAULT_TEXT_CALIB_DATASET + ) + + # Infer the calibration modality from the dataset: the known image-text datasets require a VLM, everything + # else is text. Passing a text dataset for a VLM estimates importance from text only (vision tower idle). + use_image_calib = args.calib_dataset_name in get_supported_vlm_datasets() + if use_image_calib and not is_vlm: + raise ValueError( + f"Calibration dataset '{args.calib_dataset_name}' is image-text and requires a VLM; " + "pass a text dataset for a language model." + ) + if is_vlm and not use_image_calib: + warn_rank_0( + f"Text-only calibration on a VLM (dataset '{args.calib_dataset_name}'): the language " + "model's calibration statistics will not see vision tokens." + ) + print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}") + mtq_config = get_quant_config(args) + # Quantize only the language model: disable quantizers on every top-level submodule that is not + # the language model (vision tower + projector). Skip aliases of language-model submodules (e.g. + # Qwen's ``self.decoder = language_model.decoder``) so the LM's own layers stay enabled. + if is_vlm: + lm_module_ids = {id(m) for m in language_model.modules()} + non_lm_children = sorted( + name + for name, child in unwrapped_model.named_children() + if name != "language_model" and id(child) not in lm_module_ids + ) + for name in non_lm_children: + # Anchor to the child subtree (top-level child of the quantized root) so a short non-LM + # name cannot accidentally match a language-model quantizer path by substring. + mtq_config["quant_cfg"].append({"quantizer_name": f"{name}.*", "enable": False}) + print_rank_0(f"Disabling quantizers on non-language-model submodules: {non_lm_children}") + # KV-cache quantization is incompatible with weight compression. Validate on the *resolved* # config (KV-cache quantizers are named ``*[kv]_bmm_quantizer``) so this also covers # recipe-driven KV-cache configs, not just the --kv_cache_quant flag. @@ -288,25 +350,41 @@ def main(args: argparse.Namespace): print_rank_0(f"Quantizing the model with: {args.recipe or args.quant_cfg}") if "awq" in str(mtq_config.get("algorithm")): print_rank_0( - "AWQ calibration can take longer than other methods; " - "reduce --calib_num_samples to speed it up." + "AWQ calibration can take longer than other methods; reduce --calib_num_samples to speed it up." ) # Dynamic and weight-only configs need no activation statistics, so skip both the # (potentially expensive) calibration dataset download and the calibration forward pass. - if mtq.need_calibration(mtq_config): - forward_loop = get_megatron_calibration_forward_loop( + if not mtq.need_calibration(mtq_config): + warn_rank_0("Dynamic or weight-only quantization detected; skipping calibration.") + forward_loop = None + elif not use_image_calib: + text_forward_loop = get_megatron_calibration_forward_loop( tokenizer, dataset_name=args.calib_dataset_name, num_samples=args.calib_num_samples, seq_length=args.seq_length, batch_size=args.calib_batch_size, - # pack=True uses Megatron pretraining-style global-stream document packing - pack=True, + pack=True, # Megatron pretraining-style global-stream document packing ) + + # Run text prefill on the language model: we quantize the root (a VLM root forward expects + # vision inputs), but text calibration must drive the inner LM. For plain LMs these are the same. + def forward_loop(_model=None): + text_forward_loop(language_model) else: - warn_rank_0("Dynamic or weight-only quantization detected; skipping calibration.") - forward_loop = None + # VLMs: drive the full VLM forward on image-text pairs so the language model's quantizers + # see vision-conditioned activations (we still quantize the LM only). + processor = AutoProcessor.from_pretrained( + args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code + ) + forward_loop = get_megatron_vlm_calibration_forward_loop( + unwrapped_model, # full VLM (vision encoder + projector + language model) + processor, + dataset_name=args.calib_dataset_name, + num_samples=args.calib_num_samples, + batch_size=args.calib_batch_size, + ) if hasattr(unwrapped_model, "calibration_mode"): # Some model wrappers (e.g. distillation/speculative) gate calibration behind a flag. @@ -349,13 +427,14 @@ def main(args: argparse.Namespace): ) if not args.skip_generate and not args.compress: print_rank_0("\nTesting quantized model with custom prompts...") - unwrapped_model.eval() + # Sanity-check text generation on the quantized language model. + language_model.eval() for idx, prompt in enumerate(args.prompts.split("|")): tokens = tokenizer(prompt, return_tensors="pt") # enable_kv_cache=False avoids pre-allocating the static KV cache: this is a short sanity-check # generation and the KV-cache allocation can OOM tight quantization runs on large MoE models. generated_ids = megatron_generate( - unwrapped_model, tokens.input_ids.cuda(), osl=args.osl, enable_kv_cache=False + language_model, tokens.input_ids.cuda(), osl=args.osl, enable_kv_cache=False ) generated_texts = tokenizer.batch_decode(generated_ids) print_rank_0(f"\nPrompt {idx + 1}: {prompt}\nGenerated: {generated_texts}") diff --git a/examples/pruning/README.md b/examples/pruning/README.md index dab47e8f651..332d5cbb83e 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -174,19 +174,21 @@ If your model parameters are already sorted and you just want to prune the weigh | **Algorithm** | **Model** | **Pruning Constraints** | | :---: | :---: | :---: | -| Minitron | Megatron-core1 (M-Bridge, M-LM) based dense / MoE / hybrid Mamba-Transformer LLMs2 | **Auto:** one or more of `params`, `active_params`, `memory_mb`
**Manual:** `export_config` with width (`hidden_size`, `ffn_hidden_size`, `num_attention_heads`3, `mamba_num_heads`, `mamba_head_dim`, `num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`) and/or depth (`num_layers`) pruned values | -| Puzzletron | Hugging Face based dense / MoE / hybrid Mamba-Transformer LLMs & VLMs4 | **Target:** one or more of `target_memory`, `num_params`, `target_latency_seconds`
**Heterogeneous (per-layer) search dimensions:**5 FFN `intermediate_size` (different sizes per layer), attention `op`/`no_op` (selective attention-layer removal) and KV heads (GQA grouping), `hidden_size`, and MoE `num_experts` (expert removal) | +| Minitron | Megatron-core1 (M-Bridge, M-LM) based dense / MoE / hybrid Mamba-Transformer LLMs4 (and the language model of VLMs)2 | **Auto:** one or more of `params`, `active_params`, `memory_mb`
**Manual:** `export_config` with width (`hidden_size`, `ffn_hidden_size`, `num_attention_heads`3, `mamba_num_heads`, `mamba_head_dim`, `num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`) and/or depth (`num_layers`) pruned values | +| Puzzletron | Hugging Face based dense / MoE / hybrid Mamba-Transformer LLMs & VLMs5 | **Target:** one or more of `target_memory`, `num_params`, `target_latency_seconds`
**Heterogeneous (per-layer) search dimensions:**6 FFN `intermediate_size` (different sizes per layer), attention `op`/`no_op` (selective attention-layer removal) and KV heads (GQA grouping), `hidden_size`, and MoE `num_experts` (expert removal) | | FastNAS | Computer Vision models | `flops`, `params` | > *1.Hugging Face models can be imported into M-Bridge/M-LM format as long as they are [supported](https://docs.nvidia.com/nemo/megatron-bridge/latest/index.html#supported-models) by the framework.* -> *2.Language model part of VLMs can be pruned as well.* +> *2.The language model of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) can be pruned as well; the vision tower is left intact. `hidden_size` is not pruned for VLMs as it is shared with the vision projector. See the [Megatron-Bridge pruning example](../megatron_bridge/README.md#pruning).* > *3.`num_attention_heads` pruning is not supported for some attention types — GatedDeltaNet (linear attention), gated attention (`attention_output_gate`, e.g. Qwen3.5) and Multi-Latent Attention (MLA, e.g. DeepSeek); for these only `hidden_size` is pruned.* -> *4.Puzzletron operates on Hugging Face checkpoints via its `AnyModel` abstraction. New architectures can be added by writing a model descriptor + converter — see the [AnyModel Guide](../../modelopt/torch/puzzletron/anymodel/README.md). Available configs are in the Puzzletron [configs](../puzzletron/configs/) directory.* +> *4.Multi-token-prediction (MTP) heads (e.g. Qwen3.5) are not pruned yet — they are dropped for the prune run and the saved checkpoint has no MTP. Autoregressive inference is unaffected; for speculative decoding, run a separate MTP SFT on the pruned model.* -> *5.The MIP search produces a heterogeneous architecture (dimensions can differ per layer). Which dimensions are searched is model- and config-dependent.* +> *5.Puzzletron operates on Hugging Face checkpoints via its `AnyModel` abstraction. New architectures can be added by writing a model descriptor + converter — see the [AnyModel Guide](../../modelopt/torch/puzzletron/anymodel/README.md). Available configs are in the Puzzletron [configs](../puzzletron/configs/) directory.* + +> *6.The MIP search produces a heterogeneous architecture (dimensions can differ per layer). Which dimensions are searched is model- and config-dependent.* ## Examples diff --git a/modelopt/torch/nas/plugins/mbridge.py b/modelopt/torch/nas/plugins/mbridge.py index 051d85d9823..74feac23c9f 100644 --- a/modelopt/torch/nas/plugins/mbridge.py +++ b/modelopt/torch/nas/plugins/mbridge.py @@ -38,6 +38,7 @@ from .megatron import ( NumAttentionHeadsHp, _DynamicLanguageModelEmbedding, + _DynamicMCoreLanguageModel, _DynamicSelfAttention, _DynamicTEProjRowParallelLinear, _DynamicTERowParallelLinear, @@ -128,3 +129,28 @@ def _convert_linear_proj( num_attention_heads=num_attention_heads, hidden_size=hidden_size, ) + + +# Qwen3-VL / Qwen3.5-VL ########################################################################### +# The VLM wraps the language model as ``Qwen3VLModel.language_model`` (a ``Qwen3VLGPTModel``); prune +# the language model by passing ``model.language_model`` to ``mtp.prune``. Both the LM class and its +# full-attention class override ``forward`` (absolute mRoPE), so they need explicit registration. +# GatedDeltaNet (Qwen3.5) and gated attention are already handled in ``megatron.py``. +# Guarded import — these classes exist only in newer Megatron-Bridge builds with the Qwen3-VL bridge. +try: + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import Qwen3VLGPTModel + + _QWEN3VL_PREFIX = "megatron.bridge.models.qwen_vl.modelling_qwen3_vl" + + # Both subclasses only override forward; their dynamic setup is identical to the megatron-core + # base classes, so reuse the existing dynamic classes directly (no behavioral subclass needed). + DMRegistry.register({Qwen3VLGPTModel: f"{_QWEN3VL_PREFIX}.text_model.Qwen3VLGPTModel"})( + _DynamicMCoreLanguageModel + ) + DMRegistry.register( + {Qwen3VLSelfAttention: f"{_QWEN3VL_PREFIX}.attention.Qwen3VLSelfAttention"} + )(_DynamicSelfAttention) + +except ImportError: + pass diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 7a298cac8d8..cee43d85807 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -311,19 +311,25 @@ def before_search(self) -> None: assert isinstance(self.constraints[k], (int, float)), f"{k} must be a float!" assert self.has_score, "score_func (e.g. MMLU) is required for metric-based pruning!" export_config = None - # Sort all parameters for metric-based pruning - self.hps_to_sort = SUPPORTED_HPARAMS + # Sort only hparams that may be pruned: sorting never-pruned hparams is churn, and is + # unsafe for ``hidden_size`` on VLMs (shared with the frozen vision projector, so + # permuting it would misalign injected image features). + self.hps_to_sort = SUPPORTED_HPARAMS - set(self.config["hparams_to_skip"] or []) for n, hp in named_hparams(self.model, unique=True): hp_name = n.split(".")[-1] + hp.reset_choices() # Refresh ConcatHparam choices (recomputed after modify()) before validating if hp.is_configurable: # Make sure configurable hparams are the ones with right names else implementation needs to be fixed! assert hp_name in SUPPORTED_HPARAMS, f"[ImplError] Invalid hparam {hp_name}!" - if export_config is not None and hp_name in export_config: - assert export_config[hp_name] in hp.choices, ( - f"Invalid choice {export_config[hp_name]} for {n}! Available choices: {hp.choices}" - ) - hp.reset_choices() # Make sure ConcatHparam choices are updated after modify() + # Validate every matching hparam (even non-configurable): an out-of-range value is else + # silently ignored while model.config is overwritten -> weights mismatch the saved config. + if export_config is not None and hp_name in export_config: + assert export_config[hp_name] in hp.choices, ( + f"Invalid choice {export_config[hp_name]} for {n}! Available choices: " + f"{hp.choices}. Manual export_config values must match the search-space " + "granularity (see the *_divisor settings)." + ) assert isinstance(self.model, _DynamicMCoreLanguageModel), ( "Input should be unwrapped MCore model!" @@ -794,13 +800,42 @@ def _set_divisors(c): return config +def _inherit_base_model_rules(model: nn.Module, rules: dict) -> dict: + """Let registered model subclasses inherit their base ``SUPPORTED_MODELS`` rule. + + Model subclasses (e.g. VLM language models like ``Qwen3VLGPTModel``) are registered under their + own ``DMRegistry`` key but share the dynamic class (and thus the search-space rule schema) of + their base ``GPTModel``/``MambaModel``. ``MCoreMinitronConfig`` only defines rule fields for the + base classes, so without this a subclass module would have no matching rule and get *frozen* + during conversion (disabling width/depth pruning). Copy the base rule onto the subclass key. + """ + rules = dict(rules) + for base_cls, base_key in SUPPORTED_MODELS.items(): + base_rule = rules.get(base_key) + if base_rule is None: + continue + # GPTModel/MambaModel are siblings (not subclasses of each other), so isinstance uniquely + # assigns each module to its base; the subclass shares the base's dynamic class and hence its + # rule schema. A subclass registered under its own key but absent from rules inherits it here. + for mod in model.modules(): + if not isinstance(mod, base_cls) or type(mod) not in DMRegistry: + continue + key = DMRegistry.get_key(type(mod)) + if key not in rules: + rules[key] = base_rule + return rules + + def _convert_model_to_dynamic_space( model: nn.Module, config: ModeloptBaseConfig | None = None ) -> DynamicSpace: """Create a dynamic space for the model (in-place).""" dynamic_space = DynamicSpace(model) dynamic_space._should_be_converted = lambda mod: isinstance(mod, tuple(SUPPORTED_MODELS.keys())) - dynamic_space.convert_to_dynamic(config.model_dump() if config else None, DMRegistry) + rules = config.model_dump() if config else None + if rules is not None: + rules = _inherit_base_model_rules(model, rules) + dynamic_space.convert_to_dynamic(rules, DMRegistry) if not dynamic_space.is_configurable(): raise ApplyModeError( "The model does not contain any configurable hyperparameters! Please check the" diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 4dd52ac9422..f7c46eef29a 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -294,7 +294,7 @@ def _apply_chat_template_to_messages( "preprocess": lambda sample: sample["text"], }, "wikitext": { - "config": {"path": "wikitext", "name": "wikitext-103-v1", "split": ["train"]}, + "config": {"path": "Salesforce/wikitext", "name": "wikitext-103-v1", "split": ["train"]}, "preprocess": lambda sample: sample["text"], }, } diff --git a/modelopt/torch/utils/logging.py b/modelopt/torch/utils/logging.py index 0c560f1b2ce..d1a9fc1aef9 100644 --- a/modelopt/torch/utils/logging.py +++ b/modelopt/torch/utils/logging.py @@ -130,6 +130,9 @@ def warn_rank_0(message, *args, **kwargs): """ if dist.is_master(): kwargs["stacklevel"] = kwargs.get("stacklevel", 1) + 1 + # Yellow on a TTY only (avoid escape codes in redirected output). + if isinstance(message, str) and sys.stderr.isatty(): + message = f"\033[33m{message}\033[0m" warnings.warn(message, *args, **kwargs) diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index 02b5339431c..e53e6516992 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -17,6 +17,7 @@ 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 @@ -27,7 +28,6 @@ load_modelopt_state, ) from megatron.core.models.gpt import GPTModel -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.mamba import MambaModel from megatron.core.transformer.module import MegatronModule from megatron.core.utils import unwrap_model @@ -39,6 +39,48 @@ __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, @@ -71,6 +113,7 @@ 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, @@ -85,17 +128,14 @@ def load_mbridge_model_from_hf( assert hasattr(provider, key), f"{type(provider)} does not have attribute {key}" setattr(provider, key, value) - # Only MoE models need their layer spec overridden to disable moe_grouped_gemm (not supported - # by pruning yet). Dense models keep the bridge's native spec, which is required for models - # with custom layers (e.g. Gemma3's gemma3_layer_spec) to be built correctly. + # Pruning does not support grouped GEMM yet, so disable it for MoE models. Set the flag on the + # provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than + # replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's + # GatedDeltaNet + gated-attention or Gemma3's custom spec). if isinstance(provider, MambaModelProvider): provider.mamba_stack_spec = get_te_mamba_stack_spec(moe_grouped_gemm=moe_grouped_gemm) elif (provider.num_moe_experts or 0) > 0: - provider.transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( - num_experts=provider.num_moe_experts, - moe_grouped_gemm=moe_grouped_gemm, - qk_layernorm=provider.qk_layernorm, - ) + provider.moe_grouped_gemm = moe_grouped_gemm provider.finalize() if init_model_parallel: provider.initialize_model_parallel(seed=0) @@ -103,7 +143,13 @@ def load_mbridge_model_from_hf( model = provider.provide_distributed_model(wrap_with_ddp=False) assert len(model) == 1 unwrapped_model = unwrap_model(model[0]) - assert isinstance(unwrapped_model, (GPTModel, MambaModel)) + # VLMs (e.g. Qwen3-VL) wrap the language model as ``.language_model``; the pruning target is the + # inner GPTModel/MambaModel, but we still return the full wrapper so callers can save the VLM. + language_model = getattr(unwrapped_model, "language_model", unwrapped_model) + assert isinstance(language_model, (GPTModel, MambaModel)), ( + f"Expected a GPTModel/MambaModel (optionally wrapped as .language_model), " + f"got {type(unwrapped_model)}" + ) tokenizer = AutoTokenizer.from_pretrained( hf_model_name_or_path, trust_remote_code=trust_remote_code diff --git a/modelopt/torch/utils/plugins/megatron_calibration.py b/modelopt/torch/utils/plugins/megatron_calibration.py index 223215a2923..fd0d6ff5ce0 100644 --- a/modelopt/torch/utils/plugins/megatron_calibration.py +++ b/modelopt/torch/utils/plugins/megatron_calibration.py @@ -16,6 +16,7 @@ """Shared calibration forward-loop builder for Megatron-Core models.""" import copy +import math from collections.abc import Callable from typing import TYPE_CHECKING @@ -25,13 +26,17 @@ from modelopt.torch.utils import distributed as dist from modelopt.torch.utils.dataset_utils import get_dataset_dataloader +from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader from .megatron_generate import cp_split_sequence, megatron_prefill if TYPE_CHECKING: - from transformers import PreTrainedTokenizerBase + from transformers import PreTrainedTokenizerBase, ProcessorMixin -__all__ = ["get_megatron_calibration_forward_loop"] +__all__ = [ + "get_megatron_calibration_forward_loop", + "get_megatron_vlm_calibration_forward_loop", +] def get_megatron_calibration_forward_loop( @@ -54,8 +59,7 @@ def get_megatron_calibration_forward_loop( maps to that function's ``max_sample_length``. Returns: - A ``forward_loop(model)`` callable to pass into ``mtq.quantize``, - ``mtp.prune``, or other such APIs. + A ``forward_loop(model)`` callable to pass into ``mtq.quantize``, ``mtp.prune``, or other such APIs. """ # Deepcopy before mutating pad_token so the caller's tokenizer isn't silently changed. if getattr(tokenizer, "pad_token", None) is None: @@ -94,3 +98,76 @@ def _forward_loop(model: torch.nn.Module) -> None: megatron_prefill(model, sample["input_ids"], skip_return_logits=True) return _forward_loop + + +def get_megatron_vlm_calibration_forward_loop( + model: torch.nn.Module, + processor: "ProcessorMixin", + *, + dataset_name: str = "scienceqa", + batch_size: int = 1, + num_samples: int = 512, + device: torch.device | str | None = "cuda", + subsets: list[str] | None = None, + max_shards: int | None = None, +) -> Callable[[torch.nn.Module], None]: + """Build a Megatron-Core **multimodal** calibration ``forward_loop`` for a VLM. + + This iterates image-text pairs and drives the **full VLM** forward so the language model's activation + statistics are conditioned on the merged vision tokens. The returned loop ignores the model + argument passed to it and always runs ``model`` (the full VLM captured here) -- this lets the + caller quantize only the inner ``language_model`` while still calibrating it on real multimodal activations. + + Args: + model: The full VLM (e.g. the ``Qwen3VLModel`` wrapper) to run forward for calibration. + processor: The HF processor (e.g. ``AutoProcessor``) used to encode the image-text pairs. + dataset_name: VLM calibration dataset name (see ``vlm_dataset_utils``). + batch_size: Calibration batch size. + num_samples: Number of calibration samples. + device: Device to move the encoded tensors to. + subsets: Subsets to use (only for ``nemotron_vlm_dataset_v2``; ignored otherwise). + max_shards: Max media tar shards to download per subset (only for ``nemotron_vlm_dataset_v2``; + ignored otherwise). Caps the download for large multi-shard subsets. + + Returns: + A ``forward_loop(model)`` callable to pass into ``mtq.quantize``, ``mtp.prune``, or other such APIs. + """ + # Shard the image-text data across DP ranks + dp_size = mpu.get_data_parallel_world_size() + dataloader = get_vlm_dataset_dataloader( + dataset_name=dataset_name, + processor=processor, + subsets=subsets, + max_shards=max_shards, + batch_size=batch_size, + num_samples=num_samples, + device=device, + require_image=True, + dp_size=dp_size, + dp_rank=mpu.get_data_parallel_rank(), + ) + # tqdm total: the streaming/sharded dataloader has no __len__. + total_batches = math.ceil(math.ceil(num_samples / dp_size) / batch_size) + + def _forward_loop(_model: torch.nn.Module | None = None) -> None: + # CP would have to split each sequence across ranks, but the multimodal forward merges vision + # embeddings into the sequence (the vision tower is not CP-split), so the text-style sequence + # split would misalign them. Use DP/TP/PP, or run text-only calibration (a text + # --calib_dataset_name uses get_megatron_calibration_forward_loop, which supports CP). + cp_size = mpu.get_context_parallel_world_size() + if cp_size != 1: + raise RuntimeError( + f"get_megatron_vlm_calibration_forward_loop requires CP=1, got " + f"context_parallel_world_size={cp_size}. Run calibration without CP." + ) + for batch in tqdm(dataloader, total=total_batches, disable=not dist.is_master()): + megatron_prefill( + model, + batch["input_ids"], + pixel_values=batch.get("pixel_values"), + image_grid_thw=batch.get("image_grid_thw"), + image_sizes=batch.get("image_sizes"), + skip_return_logits=True, + ) + + return _forward_loop diff --git a/modelopt/torch/utils/plugins/megatron_generate.py b/modelopt/torch/utils/plugins/megatron_generate.py index 5bcf253c02f..d868f1ba65c 100644 --- a/modelopt/torch/utils/plugins/megatron_generate.py +++ b/modelopt/torch/utils/plugins/megatron_generate.py @@ -225,6 +225,10 @@ def megatron_prefill( else: output = model(tokens, position_ids, attention_mask, runtime_gather_output=True) + # Some VLM wrappers (e.g. Gemma3VLModel) return ``(logits, loss_mask)`` rather than a bare tensor. + if isinstance(output, (tuple, list)): + output = output[0] + # For PP non-last stages, forward activations to the next stage and return early. if is_pp and not pp_last: pp_dtype = model.config.pipeline_dtype or ( @@ -244,8 +248,13 @@ def megatron_prefill( logits_dtype = torch.float32 # All PP ranks must participate in the broadcast to stay in sync. + # VLM wrappers (e.g. Qwen3VLModel) hold the output layer on the inner language_model, so the + # vocab size lives there rather than on the wrapper itself. + vocab_size = getattr(model, "vocab_size", None) + if vocab_size is None: + vocab_size = getattr(model, "language_model", model).vocab_size result = broadcast_from_last_pipeline_stage( - [batch_size, seq_length, model.vocab_size], logits_dtype, logits + [batch_size, seq_length, vocab_size], logits_dtype, logits ) return None if skip_return_logits else result diff --git a/modelopt/torch/utils/vlm_dataset_utils.py b/modelopt/torch/utils/vlm_dataset_utils.py index 2f1ae72b91f..71240968b60 100644 --- a/modelopt/torch/utils/vlm_dataset_utils.py +++ b/modelopt/torch/utils/vlm_dataset_utils.py @@ -35,13 +35,17 @@ # Use dict to store the config for each dataset. # If we want to export more options to user like target languages, we need more standardized approach like dataclass. SUPPORTED_VLM_DATASET_CONFIG: dict[str, dict[str, Any]] = { + # Small dataset; good for tests and lightweight calibration "scienceqa": {"config": {"path": "derek-thomas/ScienceQA", "split": "train"}}, # Large multi-subset dataset (use streaming to avoid downloading the entire dataset) "nemotron_vlm_dataset_v2": { "config": {"path": "nvidia/Nemotron-VLM-Dataset-v2", "split": "train", "streaming": True}, - # Provide a sane default that (a) includes in-repo media shards and (b) is document-centric. - # Subsets like docvqa_cot/chartqa_cot are JSONL-only in the dataset repo and require --vlm_image_root. + # Document-centric subsets with in-repo media shards (charts + tables + diagrams/general), + # a reasonable default for document/figure-heavy VLM benchmarks (e.g. MMMU). Subsets like + # docvqa_cot/chartqa_cot are JSONL-only in the repo and require --vlm_image_root. "default_subsets": ["sparsetables", "plotqa_cot", "wiki_en"], + # Cap tar shards per subset so the default download stays bounded (~9.5GB vs ~49GB for all). + "max_shards": 1, }, } @@ -63,6 +67,23 @@ def __len__(self): return self._num_samples +class _ShardedIterable(torch.utils.data.IterableDataset): + """Shard an iterable dataset across ranks by strided iteration (rank ``r`` of ``world``). + + Disjoint shards require every rank to iterate the same order, so the underlying stream must be + seeded identically across ranks (as our VLM datasets are). + """ + + def __init__(self, base, rank: int, world: int): + super().__init__() + self._base = base + self._rank = rank + self._world = world + + def __iter__(self): + return itertools.islice(iter(self._base), self._rank, None, self._world) + + def _extract_text_from_messages(messages: Any) -> str | None: """Best-effort extraction of a user text prompt from a chat-style `messages` field.""" if not isinstance(messages, list): @@ -220,12 +241,14 @@ def _get_vlm_dataset( Returns: A hugging face Dataset. """ + from datasets import interleave_datasets, load_dataset + # Load the dataset if dataset_name in SUPPORTED_VLM_DATASET_CONFIG: - from datasets import load_dataset - cfg = SUPPORTED_VLM_DATASET_CONFIG[dataset_name]["config"].copy() streaming = bool(cfg.pop("streaming", False)) + if max_shards is None: + max_shards = SUPPORTED_VLM_DATASET_CONFIG[dataset_name].get("max_shards") if dataset_name == "nemotron_vlm_dataset_v2": # This dataset contains many subsets; load only the requested ones via `name=...`. @@ -273,8 +296,6 @@ def _get_vlm_dataset( for subset in subsets ] try: - from datasets import interleave_datasets - ds = interleave_datasets(streams) except Exception: # Fallback: round-robin by chaining (less balanced than interleave). @@ -289,8 +310,8 @@ def _get_vlm_dataset( f" {get_supported_vlm_datasets()}." ) - # Streaming datasets: shuffle with bounded buffer and wrap into a torch IterableDataset. - if dataset_name == "nemotron_vlm_dataset_v2": + # Streaming datasets: shuffle with a bounded buffer for sample diversity + if streaming: with contextlib.suppress(Exception): ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer_size) @@ -336,7 +357,6 @@ def get_vlm_dataset_dataloader( batch_size: int = 1, num_samples: int = 512, device: str | torch.device | None = None, - max_length: int | None = None, require_image: bool = True, subsets: list[str] | None = None, shuffle_buffer_size: int = 10_000, @@ -344,6 +364,8 @@ def get_vlm_dataset_dataloader( image_root: str | Path | None = None, use_media_shards: bool = True, max_shards: int | None = None, + dp_size: int = 1, + dp_rank: int = 0, ) -> DataLoader: """Get a dataloader with the dataset name and processor of the target model. @@ -353,8 +375,9 @@ def get_vlm_dataset_dataloader( batch_size: Batch size of the returned dataloader. num_samples: Number of samples from the dataset. device: Device to move returned tensors to. If None, keep on CPU. - max_length: Optional max length for text tokenization (if supported by the processor). require_image: If True, keep only samples that have an image field. + dp_size: Data-parallel world size; if > 1, the dataset is sharded across DP ranks. + dp_rank: This process's rank within the data-parallel group. Returns: An instance of dataloader. @@ -404,6 +427,20 @@ def _collate_fn(examples: list[dict[str, Any]]) -> dict[str, torch.Tensor] | dic # Prompt extraction prompt = None tok = getattr(processor, "tokenizer", None) + + # Datasets that ship images but no chat-style ``messages`` (e.g. ScienceQA) would + # otherwise fall back to a text-only prompt with no image placeholder, so the processor + # emits no image tokens and the VLM forward fails. Synthesize a single user turn with an + # image part + the question text so the chat template emits the image token(s). + if messages is None and img is not None: + question = ex.get("question") or "Describe this image in detail." + messages = [ + { + "role": "user", + "content": [{"type": "image"}, {"type": "text", "text": question}], + } + ] + if tok is not None and messages is not None: trimmed = _messages_up_to_last_user(messages) or [] # For some Nemotron-style templates, the image content expects an empty string. @@ -446,9 +483,7 @@ def _collate_fn(examples: list[dict[str, Any]]) -> dict[str, torch.Tensor] | dic "return_tensors": "pt", "padding": True, } - if max_length is not None and "images" not in kwargs: - kwargs.update({"truncation": True, "max_length": max_length}) - + # No truncation: it clips image tokens and breaks the VLM forward enc = processor(**kwargs) # Some processors return BatchEncoding; normalize to plain dict of tensors. @@ -463,4 +498,19 @@ def _collate_fn(examples: list[dict[str, Any]]) -> dict[str, torch.Tensor] | dic out[k] = v.to(device) return out - return DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=_collate_fn) + # Data-parallel sharding: DistributedSampler for map-style datasets, strided iteration for + # iterable/streaming ones (each rank then processes ~num_samples / dp_size samples). + sampler = None + if dp_size > 1: + # Discriminate on dataset kind, not __len__: the streaming wrapper is an IterableDataset + # that also defines __len__, and DataLoader rejects a sampler on an IterableDataset. + if isinstance(dataset, torch.utils.data.IterableDataset): + dataset = _ShardedIterable(dataset, rank=dp_rank, world=dp_size) + else: + from torch.utils.data.distributed import DistributedSampler + + sampler = DistributedSampler(dataset, num_replicas=dp_size, rank=dp_rank, shuffle=False) + + return DataLoader( + dataset, batch_size=batch_size, sampler=sampler, shuffle=False, collate_fn=_collate_fn + ) diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index 593c0f4f2e5..a946ebf2a74 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -14,6 +14,7 @@ # limitations under the License. import contextlib +from functools import partial from pathlib import Path import pytest @@ -23,11 +24,13 @@ transformers = pytest.importorskip("transformers") from transformers import ( AutoModelForCausalLM, + AutoModelForImageTextToText, AutoModelForQuestionAnswering, + AutoProcessor, AutoTokenizer, BertConfig, DeepseekV3Config, - Gemma3TextConfig, + Gemma3Config, GptOssConfig, LlamaConfig, NemotronConfig, @@ -57,87 +60,147 @@ def get_tiny_tokenizer(*, pad_side: str = "left") -> "transformers.PreTrainedTok return tokenizer -##### Qwen3 ##### -def get_tiny_qwen3(**config_kwargs) -> PreTrainedModel: - set_seed(SEED) - - kwargs = { - "dtype": torch.bfloat16, - "hidden_size": 32, - "intermediate_size": 32, - "num_hidden_layers": 2, - "num_attention_heads": 16, - "num_key_value_heads": 2, - "max_position_embeddings": 32, - "vocab_size": 32, - } - kwargs.update(**config_kwargs) - # NOTE: Use AutoModelForCausalLM.from_config() instead of Qwen3ForCausalLM() for correct dtype handling - tiny_qwen3 = AutoModelForCausalLM.from_config(Qwen3Config(**kwargs)) - - return tiny_qwen3 +def _pad_vocab_size(vocab_size: int, multiple: int = 128) -> int: + """Round a vocab size up to a multiple.""" + return ((vocab_size + multiple - 1) // multiple) * multiple -def create_tiny_qwen3_dir( - tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs +def _create_tiny_llm_dir( + dir_path: Path | str, + build_model, + *, + with_tokenizer: bool = False, + return_model: bool = False, + tokenizer_factory=get_tiny_tokenizer, + **config_kwargs, ) -> Path | tuple[Path, PreTrainedModel]: - qwen3_dir = Path(tmp_path) / "tiny_qwen3" + """Save a tiny model (and, if ``with_tokenizer``, a tokenizer sized to it) to ``dir_path``.""" + dir_path = Path(dir_path) if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(qwen3_dir) + tokenizer = tokenizer_factory() + tokenizer.save_pretrained(dir_path) config_kwargs["vocab_size"] = tokenizer.vocab_size - tiny_qwen3 = get_tiny_qwen3(**config_kwargs) - tiny_qwen3.save_pretrained(qwen3_dir) + model = build_model(**config_kwargs) + model.save_pretrained(dir_path) + return (dir_path, model) if return_model else dir_path + + +def _get_tiny_vlm_tokenizer(ref_tokenizer: "transformers.PreTrainedTokenizerBase"): + """Tiny tokenizer re-registering ``ref_tokenizer``'s named vision tokens + chat template, so the + saved processor and the model config (built from the same tokenizer) share small, matching ids.""" + # transformers' model-agnostic base special-token attributes; anything beyond these on a tokenizer + # is a model-specific named token (e.g. a VLM's image_token / vision_start_token). + _base_special_token_attrs = { + "bos_token", + "eos_token", + "unk_token", + "sep_token", + "pad_token", + "cls_token", + "mask_token", + "additional_special_tokens", + } - if return_model: - return qwen3_dir, tiny_qwen3 - else: - return qwen3_dir + named = { + attr: str(getattr(ref_tokenizer, attr)) + for attr in ref_tokenizer.SPECIAL_TOKENS_ATTRIBUTES + if attr not in _base_special_token_attrs and getattr(ref_tokenizer, attr, None) is not None + } + tokenizer = get_tiny_tokenizer() + tokenizer.add_special_tokens({"additional_special_tokens": list(named.values())}) + # Register the named tokens so ```` + ``_id`` resolve and round-trip through save. + tokenizer._set_model_specific_special_tokens(named) + tokenizer.chat_template = ref_tokenizer.chat_template + return tokenizer + + +def get_tiny_vlm_processor( + ref_model_id: str, *, trust_remote_code: bool = False +) -> "transformers.ProcessorMixin": + """Tiny-vocab VLM processor: the real ``ref_model_id`` image/video processor + chat template, + paired with a tiny tokenizer that carries the ref's vision tokens (so the vocab stays small).""" + real = AutoProcessor.from_pretrained(ref_model_id, trust_remote_code=trust_remote_code) + tokenizer = _get_tiny_vlm_tokenizer(real.tokenizer) + kwargs = {"image_processor": real.image_processor, "tokenizer": tokenizer} + if getattr(real, "video_processor", None) is not None: + kwargs["video_processor"] = real.video_processor + return type(real)(**kwargs) + + +def _create_tiny_vlm_dir( + dir_path: Path | str, + ref_model_id: str, + build_model, + *, + with_processor: bool, + return_model: bool, + **config_kwargs, +) -> Path | tuple[Path, PreTrainedModel]: + """Save a tiny VLM (and, if ``with_processor``, its small-vocab processor) to ``dir_path``.""" + dir_path = Path(dir_path) + if with_processor: + get_tiny_vlm_processor(ref_model_id).save_pretrained(dir_path) + model = build_model(**config_kwargs) + model.save_pretrained(dir_path) + return (dir_path, model) if return_model else dir_path -##### Qwen3 MoE ##### -def get_tiny_qwen3_moe(**config_kwargs) -> PreTrainedModel: +##### Qwen3 (dense or MoE) ##### +def _get_tiny_qwen3(moe: bool = False, **config_kwargs) -> PreTrainedModel: set_seed(SEED) kwargs = { "dtype": torch.bfloat16, "hidden_size": 32, "intermediate_size": 32, - "moe_intermediate_size": 32, "num_hidden_layers": 2, "num_attention_heads": 16, "num_key_value_heads": 2, "max_position_embeddings": 32, "vocab_size": 32, - "num_experts": 4, - "num_experts_per_tok": 2, - "decoder_sparse_step": 1, } - kwargs.update(**config_kwargs) - tiny_qwen3_moe = AutoModelForCausalLM.from_config(Qwen3MoeConfig(**kwargs)) - - return tiny_qwen3_moe + if moe: + kwargs.update( + { + "moe_intermediate_size": 32, + "num_experts": 4, + "num_experts_per_tok": 2, + "decoder_sparse_step": 1, + } + ) + kwargs.update(config_kwargs) + # NOTE: Use AutoModelForCausalLM.from_config() instead of Qwen3[Moe]ForCausalLM() for correct dtype handling + return AutoModelForCausalLM.from_config((Qwen3MoeConfig if moe else Qwen3Config)(**kwargs)) + + +def _create_tiny_qwen3_dir( + tmp_path: Path | str, + with_tokenizer: bool = False, + return_model: bool = False, + *, + moe: bool = False, + **config_kwargs, +) -> Path | tuple[Path, PreTrainedModel]: + return _create_tiny_llm_dir( + Path(tmp_path) / ("tiny_qwen3_moe" if moe else "tiny_qwen3"), + _get_tiny_qwen3, + with_tokenizer=with_tokenizer, + return_model=return_model, + moe=moe, + **config_kwargs, + ) -def create_tiny_qwen3_moe_dir( - tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs -) -> Path | tuple[Path, PreTrainedModel]: - qwen3_moe_dir = Path(tmp_path) / "tiny_qwen3_moe" - if with_tokenizer: - tokenizer = tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(qwen3_moe_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - get_tiny_qwen3_moe(**config_kwargs).save_pretrained(qwen3_moe_dir) - return qwen3_moe_dir +get_tiny_qwen3 = partial(_get_tiny_qwen3, moe=False) +create_tiny_qwen3_dir = partial(_create_tiny_qwen3_dir, moe=False) +get_tiny_qwen3_moe = partial(_get_tiny_qwen3, moe=True) +create_tiny_qwen3_moe_dir = partial(_create_tiny_qwen3_dir, moe=True) ##### Qwen3-VL ##### def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel: - # Lazy imports — Qwen3VL classes live under transformers.models.qwen3_vl which - # may not exist in older transformers builds, and this module is imported by - # every test that uses transformers_models.py. + # Lazy imports — Qwen3VL requires transformers>=4.57 from transformers import Qwen3VLConfig - from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration set_seed(SEED) @@ -145,6 +208,7 @@ def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel: # Pass config_kwargs to override for multi-GPU tests (e.g. num_attention_heads=num_gpus, # num_key_value_heads=num_gpus, hidden_size=num_gpus*head_dim). text_kwargs = { + "dtype": torch.bfloat16, "hidden_size": 32, "intermediate_size": 32, "num_hidden_layers": 2, @@ -167,21 +231,185 @@ def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel: "spatial_merge_size": 1, "temporal_patch_size": 1, "out_hidden_size": text_kwargs["hidden_size"], # must match text hidden_size + # Single deepstack injection (the bridge requires len <= num language-model layers, so keep + # this small for tiny models; defaults to [8, 16, 24] which needs >= 3 layers). + "deepstack_visual_indexes": [0], } - cfg = Qwen3VLConfig(text_config=text_kwargs, vision_config=vision_kwargs) - return Qwen3VLForConditionalGeneration(cfg) + cfg = Qwen3VLConfig(text_config=text_kwargs, vision_config=vision_kwargs, dtype=torch.bfloat16) + return AutoModelForImageTextToText.from_config(cfg) def create_tiny_qwen3vl_dir( - tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs -) -> Path: - qwen3vl_dir = Path(tmp_path) / "tiny_qwen3vl" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(qwen3vl_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - get_tiny_qwen3vl(**config_kwargs).save_pretrained(qwen3vl_dir) - return qwen3vl_dir + tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs +) -> Path | tuple[Path, PreTrainedModel]: + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_qwen3vl", + get_tiny_qwen3vl, + with_tokenizer=with_tokenizer, + return_model=return_model, + **config_kwargs, + ) + + +##### Gemma3-VL ##### +# Real tiny Gemma3 reused (via get_tiny_vlm_processor) for the fixture's image processor + chat +# template; the vision geometry below matches it so the processor's pixel_values fit the tiny tower. +GEMMA3_VL_REF = "hf-internal-testing/tiny-random-Gemma3ForConditionalGeneration" + + +def get_tiny_gemma3vl(**config_kwargs) -> PreTrainedModel: + set_seed(SEED) + + # Vocab + vision token ids derive from the ref tokenizer to match the saved processor. + tokenizer = _get_tiny_vlm_tokenizer(AutoTokenizer.from_pretrained(GEMMA3_VL_REF)) + + # layer_types auto-generates (sliding/full attention) so the pruning code path is exercised. + text_kwargs = { + "dtype": torch.bfloat16, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "max_position_embeddings": 1024, # >= 256 image tokens + text for image calibration + "sliding_window": 16, + "vocab_size": _pad_vocab_size(len(tokenizer)), + } + text_kwargs.update(config_kwargs) + text_kwargs.setdefault("query_pre_attn_scalar", text_kwargs["head_dim"]) + # Tiny SigLIP vision tower; the multimodal projector maps it to the text hidden size. + # Match Megatron-Bridge Gemma3VL (and real Gemma3 checkpoints): no SigLIP pooling head. + vision_kwargs = { + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + # Real patch geometry so the processor's pixel_values feed the tiny vision tower. + "image_size": 896, + "patch_size": 14, + "num_channels": 3, + "vision_use_head": False, + } + cfg = Gemma3Config( + text_config=text_kwargs, + vision_config=vision_kwargs, + mm_tokens_per_image=256, # 896 / 14 -> 4096 patches pooled to 256 image tokens + # Token ids (from the tiny tokenizer) must match the processor for vision-token merging. + image_token_index=tokenizer.image_token_id, + boi_token_index=tokenizer.boi_token_id, + eoi_token_index=tokenizer.eoi_token_id, + dtype=torch.bfloat16, + ) + return AutoModelForImageTextToText.from_config(cfg) + + +def create_tiny_gemma3vl_dir( + tmp_path: Path | str, with_processor: bool = False, return_model: bool = False, **config_kwargs +) -> Path | tuple[Path, PreTrainedModel]: + return _create_tiny_vlm_dir( + Path(tmp_path) / "tiny_gemma3vl", + GEMMA3_VL_REF, + get_tiny_gemma3vl, + with_processor=with_processor, + return_model=return_model, + **config_kwargs, + ) + + +##### Qwen3.5-VL (dense or MoE, hybrid GatedDeltaNet + gated attention) ##### +QWEN3_5_VL_REF = "Qwen/Qwen3.5-0.8B" + + +def _get_tiny_qwen3_5_vl(moe: bool = False, **config_kwargs) -> PreTrainedModel: + # Lazy imports — Qwen3.5-VL requires a recent transformers version. + from transformers import Qwen3_5Config, Qwen3_5MoeConfig + + set_seed(SEED) + + # Vocab + vision token ids derive from the ref tokenizer to match the saved processor. + tokenizer = _get_tiny_vlm_tokenizer(AutoTokenizer.from_pretrained(QWEN3_5_VL_REF)) + + # Hybrid GatedDeltaNet (linear attention) + gated full-attention (layer_types auto-generated). + text_kwargs = { + "dtype": torch.bfloat16, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "max_position_embeddings": 32, + # Pad to a multiple of 128 (Megatron pads the vocab to make_vocab_size_divisible_by=128; + # matching here avoids an embedding-size mismatch on import). Token ids stay < len(tokenizer). + "vocab_size": _pad_vocab_size(len(tokenizer)), + # GatedDeltaNet linear-attention dims (kept small for tiny models). + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 16, + "linear_value_head_dim": 16, + "linear_conv_kernel_dim": 4, + } + if moe: + # Replace the dense MLP with a tiny MoE (Qwen3.5-MoE: hybrid GatedDeltaNet + gated attention + MoE). + text_kwargs.update( + { + "num_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 64, + "shared_expert_intermediate_size": 64, + } + ) + text_kwargs.update(config_kwargs) + vision_kwargs = { + "hidden_size": 16, + "intermediate_size": 16, + "depth": 2, + "num_heads": 2, + # Real patch params so the processor's pixel_values fit the tiny tower; only hidden dims shrink. + "patch_size": 16, + "temporal_patch_size": 2, + "in_channels": 3, + "spatial_merge_size": 2, + "out_hidden_size": text_kwargs["hidden_size"], # must match text hidden_size + "deepstack_visual_indexes": [], # no deepstack injection (unlike Qwen3-VL) + } + cfg = (Qwen3_5MoeConfig if moe else Qwen3_5Config)( + text_config=text_kwargs, + vision_config=vision_kwargs, + dtype=torch.bfloat16, + # Token ids (from the tiny tokenizer) must match the processor for vision-token merging. + image_token_id=tokenizer.image_token_id, + video_token_id=tokenizer.video_token_id, + vision_start_token_id=tokenizer.vision_bos_token_id, + vision_end_token_id=tokenizer.vision_eos_token_id, + ) + return AutoModelForImageTextToText.from_config(cfg) + + +def _create_tiny_qwen3_5_vl_dir( + tmp_path: Path | str, + with_processor: bool = False, + return_model: bool = False, + *, + moe: bool = False, + **config_kwargs, +) -> Path | tuple[Path, PreTrainedModel]: + return _create_tiny_vlm_dir( + Path(tmp_path) / ("tiny_qwen3_5_moe_vl" if moe else "tiny_qwen3_5_vl"), + QWEN3_5_VL_REF, + _get_tiny_qwen3_5_vl, + with_processor=with_processor, + return_model=return_model, + moe=moe, + **config_kwargs, + ) + + +get_tiny_qwen3_5_vl = partial(_get_tiny_qwen3_5_vl, moe=False) +create_tiny_qwen3_5_vl_dir = partial(_create_tiny_qwen3_5_vl_dir, moe=False) +get_tiny_qwen3_5_moe_vl = partial(_get_tiny_qwen3_5_vl, moe=True) +create_tiny_qwen3_5_moe_vl_dir = partial(_create_tiny_qwen3_5_vl_dir, moe=True) ##### NEMOTRON ##### @@ -200,20 +428,19 @@ def get_tiny_nemotron(**config_kwargs) -> PreTrainedModel: "max_position_embeddings": 32, "vocab_size": 32, } - kwargs.update(**config_kwargs) + kwargs.update(config_kwargs) return AutoModelForCausalLM.from_config(NemotronConfig(**kwargs)) def create_tiny_nemotron_dir( tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs ) -> Path: - nemotron_dir = Path(tmp_path) / "tiny_nemotron" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(nemotron_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - get_tiny_nemotron(**config_kwargs).save_pretrained(nemotron_dir) - return nemotron_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_nemotron", + get_tiny_nemotron, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### NEMOTRON-H (Mamba + Attention + MoE/MLP hybrid) ##### @@ -249,20 +476,19 @@ def get_tiny_nemotron_h(**config_kwargs) -> PreTrainedModel: "vocab_size": 32, "max_position_embeddings": 32, } - kwargs.update(**config_kwargs) + kwargs.update(config_kwargs) return AutoModelForCausalLM.from_config(NemotronHConfig(**kwargs)) def create_tiny_nemotron_h_dir( tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs ) -> Path: - nemotron_h_dir = Path(tmp_path) / "tiny_nemotron_h" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(nemotron_h_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - get_tiny_nemotron_h(**config_kwargs).save_pretrained(nemotron_h_dir) - return nemotron_h_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_nemotron_h", + get_tiny_nemotron_h, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### DeepSeek V3 ##### @@ -290,7 +516,7 @@ def get_tiny_deepseek_v3(**config_kwargs) -> PreTrainedModel: # Required so vLLM allocates ``gate.e_score_correction_bias`` (HF saves it unconditionally). "topk_method": "noaux_tc", } - kwargs.update(**config_kwargs) + kwargs.update(config_kwargs) cfg = DeepseekV3Config(**kwargs) # Survive transformers versions that drop unknown kwargs from the dataclass. cfg.topk_method = kwargs["topk_method"] @@ -300,13 +526,12 @@ def get_tiny_deepseek_v3(**config_kwargs) -> PreTrainedModel: def create_tiny_deepseek_v3_dir( tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs ) -> Path: - deepseek_dir = Path(tmp_path) / "tiny_deepseek_v3" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(deepseek_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - get_tiny_deepseek_v3(**config_kwargs).save_pretrained(deepseek_dir) - return deepseek_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_deepseek_v3", + get_tiny_deepseek_v3, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### GPT-OSS ##### @@ -324,66 +549,19 @@ def get_tiny_gpt_oss(**config_kwargs) -> PreTrainedModel: "num_attention_heads": 2, "num_key_value_heads": 1, } - kwargs.update(**config_kwargs) - tiny_gpt_oss = AutoModelForCausalLM.from_config(GptOssConfig(**kwargs)) - - return tiny_gpt_oss + kwargs.update(config_kwargs) + return AutoModelForCausalLM.from_config(GptOssConfig(**kwargs)) def create_tiny_gpt_oss_dir( tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs ) -> Path: - gpt_oss_dir = Path(tmp_path) / "tiny_gpt_oss" - if with_tokenizer: - tokenizer = tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(gpt_oss_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - get_tiny_gpt_oss(**config_kwargs).save_pretrained(gpt_oss_dir) - return gpt_oss_dir - - -##### Gemma3 ##### -def get_tiny_gemma3(**config_kwargs) -> PreTrainedModel: - set_seed(SEED) - - # head_dim is independent of hidden_size / num_attention_heads in Gemma3. - # layer_types is left to auto-generate (all `sliding_attention` for tiny layer - # counts) so the sliding/full attention layer_types code path is still exercised. - kwargs = { - "dtype": torch.bfloat16, - "hidden_size": 32, - "intermediate_size": 32, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 8, - "max_position_embeddings": 32, - "sliding_window": 16, - "vocab_size": 32, - } - kwargs.update(**config_kwargs) - # query_pre_attn_scalar sets the attention scale (1/sqrt(query_pre_attn_scalar)); default it - # to head_dim (Gemma3's convention for all sizes except 27B) unless the caller overrides it, - # so overriding head_dim alone stays consistent. - kwargs.setdefault("query_pre_attn_scalar", kwargs["head_dim"]) - return AutoModelForCausalLM.from_config(Gemma3TextConfig(**kwargs)) - - -def create_tiny_gemma3_dir( - tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs -) -> Path | tuple[Path, PreTrainedModel]: - gemma3_dir = Path(tmp_path) / "tiny_gemma3" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(gemma3_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - tiny_gemma3 = get_tiny_gemma3(**config_kwargs) - tiny_gemma3.save_pretrained(gemma3_dir) - - if return_model: - return gemma3_dir, tiny_gemma3 - else: - return gemma3_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_gpt_oss", + get_tiny_gpt_oss, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### LLAMA ##### @@ -399,23 +577,19 @@ def get_tiny_llama(**config_kwargs) -> PreTrainedModel: "max_position_embeddings": 32, "vocab_size": 32, } - kwargs.update(**config_kwargs) - tiny_llama = AutoModelForCausalLM.from_config(LlamaConfig(**kwargs)) - - return tiny_llama + kwargs.update(config_kwargs) + return AutoModelForCausalLM.from_config(LlamaConfig(**kwargs)) def create_tiny_llama_dir( tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs ) -> Path: - llama_dir = Path(tmp_path) / "tiny_llama" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(llama_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - - get_tiny_llama(**config_kwargs).save_pretrained(llama_dir) - return llama_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_llama", + get_tiny_llama, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### T5 ##### @@ -433,22 +607,20 @@ def get_tiny_t5(**config_kwargs) -> PreTrainedModel: "relative_attention_max_distance": 32, "decoder_start_token_id": 0, } - kwargs.update(**config_kwargs) - t5_model = T5ForConditionalGeneration(T5Config(**kwargs)).to(torch.bfloat16) - - return t5_model + kwargs.update(config_kwargs) + return T5ForConditionalGeneration(T5Config(**kwargs)).to(torch.bfloat16) def create_tiny_t5_dir(tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs) -> Path: - set_seed(SEED) - t5_dir = Path(tmp_path) / "tiny_t5" - if with_tokenizer: - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-T5Model") - tokenizer.save_pretrained(t5_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - - get_tiny_t5(**config_kwargs).save_pretrained(t5_dir) - return t5_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_t5", + get_tiny_t5, + with_tokenizer=with_tokenizer, + tokenizer_factory=lambda: AutoTokenizer.from_pretrained( + "hf-internal-testing/tiny-random-T5Model" + ), + **config_kwargs, + ) ##### BERT ##### @@ -464,17 +636,12 @@ def get_tiny_bert(**config_kwargs) -> PreTrainedModel: "max_position_embeddings": 32, "vocab_size": 32, } - kwargs.update(**config_kwargs) - tiny_bert = AutoModelForQuestionAnswering.from_config(BertConfig(**kwargs)) - - return tiny_bert + kwargs.update(config_kwargs) + return AutoModelForQuestionAnswering.from_config(BertConfig(**kwargs)) def create_tiny_bert_dir(tmp_path: Path | str, **config_kwargs) -> Path: - set_seed(SEED) - bert_dir = Path(tmp_path) / "tiny_bert" - get_tiny_bert(**config_kwargs).save_pretrained(bert_dir) - return bert_dir + return _create_tiny_llm_dir(Path(tmp_path) / "tiny_bert", get_tiny_bert, **config_kwargs) ##### TESTERS ##### diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 33588234dec..0a0503a5a53 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -14,34 +14,54 @@ # limitations under the License. """Tests for prune_minitron.py and distill.py scripts.""" -from pathlib import Path - import pytest +import torch from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.transformers_models import create_tiny_gemma3_dir, create_tiny_qwen3_dir -from transformers import AutoModelForCausalLM +from _test_utils.torch.transformers_models import ( + create_tiny_gemma3vl_dir, + create_tiny_nemotron_h_dir, + create_tiny_qwen3_5_moe_vl_dir, + create_tiny_qwen3_dir, +) +from transformers import AutoModelForCausalLM, AutoModelForImageTextToText @pytest.mark.parametrize( - "create_tiny_model_dir", + ("create_teacher", "megatron_format"), [ - create_tiny_qwen3_dir, - # Gemma3 exercises the sliding/full attention ``layer_types`` pruning path. - create_tiny_gemma3_dir, + # Dense Qwen3 LM, exported back to HF (reloadable to verify the pruned param count). + pytest.param( + lambda tmp_path, num_gpus: create_tiny_qwen3_dir( + tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=num_gpus + ), + False, + id="qwen3", + ), + # NemotronH (nemotron-3-nano): Mamba + attention + MoE hybrid. Saved in Megatron checkpoint + # format because HF export of a pruned NemotronH requires transformers<5. + pytest.param( + lambda tmp_path, num_gpus: create_tiny_nemotron_h_dir( + tmp_path, with_tokenizer=True, return_model=True + ), + True, + id="nemotron_h", + ), ], ) -def test_prune_minitron(tmp_path: Path, num_gpus, create_tiny_model_dir): - teacher_hf_path, teacher_model = create_tiny_model_dir( - tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=num_gpus - ) +def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): + teacher_hf_path, teacher_model = create_teacher(tmp_path, num_gpus) teacher_params = sum(p.numel() for p in teacher_model.parameters()) prune_target_params = int(teacher_params * 0.8) - pruned_model_path = tmp_path / "pruned" + pruned_path = tmp_path / "pruned" + output_kwarg = ( + {"output_megatron_path": pruned_path} + if megatron_format + else {"output_hf_path": pruned_path} + ) prune_command_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"], hf_model_name_or_path=teacher_hf_path, - output_hf_path=pruned_model_path, pp_size=num_gpus, calib_dataset_name="cnn_dailymail", calib_num_samples=8, @@ -51,10 +71,87 @@ def test_prune_minitron(tmp_path: Path, num_gpus, create_tiny_model_dir): ss_channel_divisor=4, hparams_to_skip="num_attention_heads", top_k=1, + **output_kwarg, + ) + run_example_command(prune_command_parts, example_path="megatron_bridge") + + if megatron_format: + # HF reload of a pruned NemotronH needs transformers<5; just verify the Megatron checkpoint. + assert (pruned_path / "latest_checkpointed_iteration.txt").exists() + else: + assert (pruned_path / "config.json").exists() + pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path) + assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params + + +@pytest.mark.parametrize( + "create_teacher", + [ + # gemma3vl: sliding/full attention dense LM with a multimodal projector. + pytest.param( + lambda tmp_path, num_gpus: create_tiny_gemma3vl_dir( + tmp_path, + with_processor=True, + return_model=True, + num_hidden_layers=4 * num_gpus, + intermediate_size=128, + max_position_embeddings=1024, + ), + id="gemma3vl", + ), + # qwen3.5-VL MoE: hybrid GatedDeltaNet + gated-attention MoE LM (prunes depth + MoE dims). + pytest.param( + lambda tmp_path, num_gpus: create_tiny_qwen3_5_moe_vl_dir( + tmp_path, + with_processor=True, + return_model=True, + num_hidden_layers=4 * num_gpus, + max_position_embeddings=1024, + ), + id="qwen3_5_moe_vl", + ), + ], +) +def test_prune_minitron_vlm(tmp_path, num_gpus, create_teacher): + # >= 4 layers per PP stage so depth is prunable and stays balanced; max_position_embeddings + # >= image tokens + text so the image-text calibration sequence is not truncated. + teacher_hf_path, teacher_model = create_teacher(tmp_path, num_gpus) + language_model_params = sum(p.numel() for p in teacher_model.model.language_model.parameters()) + prune_target_params = int(language_model_params * 0.7) + + pruned_model_path = tmp_path / "pruned" + prune_command_parts = extend_cmd_parts( + ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"], + hf_model_name_or_path=teacher_hf_path, + output_hf_path=pruned_model_path, + pp_size=num_gpus, + calib_dataset_name="scienceqa", # image-text calibration runs the full VLM forward + calib_num_samples=8, + seq_length=1024, + prune_target_params=prune_target_params, + prune_score_func="mmlu_1pct_bs32", + ss_channel_divisor=4, + # Allow depth pruning (the primary param lever once hidden_size is fixed for VLMs). + max_depth_pruning=0.6, + hparams_to_skip="num_attention_heads", + top_k=1, ) run_example_command(prune_command_parts, example_path="megatron_bridge") assert (pruned_model_path / "config.json").exists() - pruned_model = AutoModelForCausalLM.from_pretrained(pruned_model_path) - pruned_params = sum(p.numel() for p in pruned_model.parameters()) - assert pruned_params <= prune_target_params + # from_pretrained (default strict load) verifies the saved weights match the pruned config. + pruned_model = AutoModelForImageTextToText.from_pretrained(pruned_model_path) + pruned_lm_params = sum(p.numel() for p in pruned_model.model.language_model.parameters()) + # Language model is pruned to the param target. + assert pruned_lm_params <= prune_target_params + # Everything outside the language model (vision tower, projector, lm_head) is byte-identical + assert hasattr(pruned_model.config, "vision_config") + teacher_non_lm = { + n: p for n, p in teacher_model.named_parameters() if "language_model." not in n + } + pruned_non_lm = {n: p for n, p in pruned_model.named_parameters() if "language_model." not in n} + assert pruned_non_lm.keys() == teacher_non_lm.keys() + for name, expected in teacher_non_lm.items(): + torch.testing.assert_close( + pruned_non_lm[name].detach().cpu(), expected.detach().cpu(), rtol=0, atol=0 + ) diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index a1ddaed9593..8e9402337f7 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -17,7 +17,7 @@ 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_dir +from _test_utils.torch.transformers_models import create_tiny_qwen3_5_vl_dir, create_tiny_qwen3_dir def test_quantize_and_export(tmp_path: Path, num_gpus): @@ -80,3 +80,41 @@ 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" + )