From b7f984c799509905203c194b8fbceb865b6072ff Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 22 Jun 2026 13:50:40 -0700 Subject: [PATCH 01/12] Add Minitron pruning support for VLM language models Prune the language model of Megatron-Bridge VLMs (Qwen3-VL, Qwen3.5-VL, Gemma3-VL) while leaving the vision tower intact. The example extracts model.language_model, prunes it in place (ffn/MoE dims + depth; hidden_size is skipped since it's shared with the vision projector), and saves the full VLM back. - nas/plugins/mbridge.py: register Qwen3VLGPTModel -> _DynamicMCoreLanguageModel and Qwen3VLSelfAttention -> _DynamicSelfAttention. - utils/plugins/mbridge.py: loader accepts VLM wrappers (validates the inner .language_model, returns the full wrapper for saving). - prune/plugins/mcore_minitron.py: protected_layers option to keep VLM deepstack-injected layers during depth pruning; inherit base-model search-space rules for registered subclasses; validate export_config values. - megatron_generate.py: read vocab_size from the inner language_model for VLMs. - examples/megatron_bridge/prune_minitron.py: VLM detection, hidden_size skip, deepstack index protection + remap, and VLM-aware HF save. - tests: tiny Qwen3-VL / Gemma3-VL / Qwen3.5-VL (dense) fixtures + parametrized test_prune_minitron_vlm. - Docs: CHANGELOG (0.46), pruning + megatron_bridge READMEs, pruning guide. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + docs/source/guides/3_pruning.rst | 5 +- examples/megatron_bridge/README.md | 12 ++ examples/megatron_bridge/prune_minitron.py | 120 ++++++++--- examples/pruning/README.md | 4 +- modelopt/torch/nas/plugins/mbridge.py | 26 +++ .../torch/prune/plugins/mcore_minitron.py | 70 ++++++- modelopt/torch/utils/plugins/mbridge.py | 8 +- .../torch/utils/plugins/megatron_generate.py | 7 +- .../_test_utils/torch/transformers_models.py | 190 +++++++++++++----- .../megatron_bridge/test_prune_minitron.py | 77 +++++-- 11 files changed, 410 insertions(+), 110 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7308931d772..93b2f66d531 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -26,6 +26,7 @@ Changelog - **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-VL, 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. - Add NVFP4 Four-Over-Six (4/6) weight quantization (``mtq.NVFP4_FOUR_OVER_SIX_CFG``): MSE weight calibration picks, per block, between an M=6 and an M=4 dynamic range (the choice is folded into the FP8 per-block scales), with the ``four_over_six: true`` flag normalizing those scales by 256 (vs 448) for M=4 headroom. Supported via ``mtq.quantize`` and HF / Megatron export only -- **not** ``mtq.compress``, which does not preserve the per-block M=4/M=6 choice - Add dLLM (tied-weight PTQ and HF-checkpoint export) support for diffusion-based encoder-decoder LLMs (e.g. DiffusionGemma) whose encoder/decoder stacks share parameters via HF ``_tied_weights_keys``. diff --git a/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/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 3f4cc17cd13..e1b7367731e 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -310,6 +310,18 @@ torchrun --nproc_per_node 1 prune_minitron.py --help > [!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-VL, 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 — the only difference is that `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-VL-8B-Instruct \ + --prune_target_params 6e9 \ + --output_hf_path /tmp/Qwen3-VL-8B-Pruned +``` + ## Resources - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699) diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index bc3401e296f..e6ecedfd0ed 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -46,7 +46,7 @@ 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 import modelopt.torch.opt as mto import modelopt.torch.prune as mtp @@ -291,6 +291,34 @@ def main(args: argparse.Namespace): init_model_parallel=True, moe_grouped_gemm=False, ) + + # 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 never pruned here (follow-up). + language_model = getattr(unwrapped_model, "language_model", unwrapped_model) + is_vlm = language_model is not unwrapped_model + # 1-indexed LM layer numbers that depth pruning must keep (see protected_layers handling below). + protected_layers: set[int] = set() + if is_vlm: + print_rank_0( + "VLM detected: pruning model.language_model only (hidden_size pruning skipped)." + ) + 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"}) + + # Deepstack-injected LM layers (Qwen3-VL): the vision tower feeds features into specific LM + # layers (``vision_config.deepstack_visual_indexes``, 0-indexed). Those layers must survive + # depth pruning so the (unpruned) vision tower's deepstack mergers still line up; protect + # them (as 1-indexed layer numbers) and remap the indices after pruning. + vision_cfg = getattr(bridge.hf_pretrained.config, "vision_config", None) + deepstack_indexes = list(getattr(vision_cfg, "deepstack_visual_indexes", None) or []) + protected_layers = {idx + 1 for idx in deepstack_indexes} + + # Calibration runs the pruning target (the language model) directly. For VLMs this is the + # extracted ``language_model``; text-only calibration uses its own embedding and works under + # pipeline parallelism via ``megatron_prefill`` (which handles inter-stage send/recv). forward_loop = get_megatron_calibration_forward_loop( tokenizer, dataset_name=args.calib_dataset_name, @@ -304,6 +332,9 @@ def main(args: argparse.Namespace): pruning_config = { "forward_loop": forward_loop, "checkpoint": args.prune_intermediate_ckpt, + # Layers that depth pruning must keep (empty set => no constraint). For VLMs this protects + # the deepstack-injected LM layers so the vision tower stays consistent. + "protected_layers": protected_layers or None, } if args.prune_export_config is not None: # Less restrictive search space for manual pruning @@ -379,16 +410,18 @@ 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 isinstance(provider, MambaModelProvider): hybrid_key = ( "hybrid_override_pattern" @@ -425,42 +458,65 @@ 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 = 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 + if hasattr(text_cfg, "moe_shared_expert_intermediate_size"): + text_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 = ( + 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 + # 1-indexed original layer numbers kept after depth pruning, in ascending (original) order. + # sorted_layers is None when no layer scores were collected (depth pruning impossible), in + # which case all layers are kept and the mapping is the identity. + sorted_layers = pruning_scores["sorted_layers"] + kept_layer_nums = ( + sorted(sorted_layers[: mcore_cfg.num_layers]) + if sorted_layers is not None + else list(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 + ] + # Remap VLM deepstack indices (0-indexed LM layers) to the pruned/renumbered layers. The + # protected_layers constraint above guarantees every deepstack layer survived, so each index + # maps to its new compacted position; the vision tower's mergers stay aligned. + vision_cfg = getattr(hf_cfg, "vision_config", None) + if ( + is_vlm + and vision_cfg is not None + and getattr(vision_cfg, "deepstack_visual_indexes", None) + ): + vision_cfg.deepstack_visual_indexes = [ + kept_layer_nums.index(idx + 1) for idx in vision_cfg.deepstack_visual_indexes ] 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 # 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/pruning/README.md b/examples/pruning/README.md index dab47e8f651..c43fc9a423c 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -174,13 +174,13 @@ 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 | +| Minitron | Megatron-core1 (M-Bridge, M-LM) based dense / MoE / hybrid Mamba-Transformer LLMs (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 & 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) | | 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-VL, 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.* 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..3ccc27df48d 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -225,6 +225,10 @@ class MCoreMinitronSearcher(BaseSearcher): - `max_depth_pruning`: Maximum fraction per depth hyperparameter to prune (default: 0.20). Only top (1 - max_depth_pruning) choices will be considered. - `hparams_to_skip`: List of hparams to skip during the search (default: None). + - `protected_layers`: Set of 1-indexed layer numbers that must never be dropped during depth + pruning (default: None). Used e.g. for VLM deepstack-injection layers, which the vision + tower feeds into and so must be kept. Works for both ``export_config`` and metric-based + pruning. Requires the target ``num_layers`` >= ``len(protected_layers)``. - `top_k`: Number of candidates to consider for score_func validation (default: 10). - `seq_length`: Sequence length for KV-cache memory estimate (default: 4096). Only used with the ``memory_mb`` constraint. @@ -250,6 +254,7 @@ def default_search_config(self) -> SearchConfig: "max_width_pruning": 0.40, "max_depth_pruning": 0.20, "hparams_to_skip": None, + "protected_layers": None, "top_k": 10, # Memory footprint config (only used with memory_mb constraint) "seq_length": 4096, @@ -319,10 +324,16 @@ def before_search(self) -> None: 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}" - ) + # Validate requested export_config values are achievable for *every* matching hparam, + # including non-configurable ones. A value outside `choices` (e.g. not aligned to the + # search-space divisor) would otherwise be silently ignored while model.config is still + # overwritten, producing a checkpoint whose weights don't match its 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)." + ) hp.reset_choices() # Make sure ConcatHparam choices are updated after modify() assert isinstance(self.model, _DynamicMCoreLanguageModel), ( @@ -366,6 +377,20 @@ def run_search(self) -> None: for layer, _ in sorted(self.layer_scores.items(), key=lambda x: x[1], reverse=True) ] assert sorted(self.sorted_layers) == list(range(1, self.model.config.num_layers + 1)) + + # Force protected layers (e.g. VLM deepstack-injection layers) to the front of the + # importance ranking so depth pruning never drops them (drop set is sorted_layers[N:]). + # Relative order within each group is preserved (stable partition). + protected = self.config["protected_layers"] + if protected: + protected = set(protected) + assert protected <= set(self.sorted_layers), ( + f"protected_layers {sorted(protected)} must be 1-indexed layer numbers in " + f"1..{self.model.config.num_layers}" + ) + self.sorted_layers = [ln for ln in self.sorted_layers if ln in protected] + [ + ln for ln in self.sorted_layers if ln not in protected + ] else: assert ( self.constraints.keys() == {"export_config"} @@ -422,6 +447,12 @@ def _prune(self, export_config: dict, prune_depth: bool = True) -> None: if num_layers_hp.active != num_layers_hp.max: assert self.sorted_layers is not None layers_to_drop = self.sorted_layers[num_layers_hp.active :] + protected = self.config["protected_layers"] + if protected: + assert not (set(layers_to_drop) & set(protected)), ( + f"Cannot depth-prune to {num_layers_hp.active} layers while protecting " + f"{sorted(protected)}: target num_layers must be >= {len(set(protected))}." + ) drop_mcore_language_model_layers(self.model, layers_to_drop=layers_to_drop) # Update model config with pruned architecture @@ -794,13 +825,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/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index 02b5339431c..c671bc338b0 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -103,7 +103,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_generate.py b/modelopt/torch/utils/plugins/megatron_generate.py index 5bcf253c02f..8eec92bafdf 100644 --- a/modelopt/torch/utils/plugins/megatron_generate.py +++ b/modelopt/torch/utils/plugins/megatron_generate.py @@ -244,8 +244,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/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index 593c0f4f2e5..8000bb2ac99 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -23,11 +23,12 @@ transformers = pytest.importorskip("transformers") from transformers import ( AutoModelForCausalLM, + AutoModelForImageTextToText, AutoModelForQuestionAnswering, AutoTokenizer, BertConfig, DeepseekV3Config, - Gemma3TextConfig, + Gemma3Config, GptOssConfig, LlamaConfig, NemotronConfig, @@ -133,11 +134,8 @@ def create_tiny_qwen3_moe_dir( ##### 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 +143,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 +166,146 @@ 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: + tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs +) -> Path | tuple[Path, PreTrainedModel]: 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 + tiny_qwen3vl = get_tiny_qwen3vl(**config_kwargs) + tiny_qwen3vl.save_pretrained(qwen3vl_dir) + + if return_model: + return qwen3vl_dir, tiny_qwen3vl + else: + return qwen3vl_dir + + +##### Gemma3-VL ##### +def get_tiny_gemma3vl(**config_kwargs) -> PreTrainedModel: + set_seed(SEED) + + # Gemma3 language model config (config_kwargs override the text side, e.g. num_hidden_layers). + # layer_types auto-generates (sliding/full attention) so that 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": 32, + "sliding_window": 16, + "vocab_size": 32, + } + 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. + vision_kwargs = { + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "image_size": 16, + "patch_size": 8, + "num_channels": 3, + } + cfg = Gemma3Config( + text_config=text_kwargs, + vision_config=vision_kwargs, + mm_tokens_per_image=4, # (image_size / patch_size) ** 2 = (16 / 8) ** 2 + image_token_index=1, # must be < vocab_size; not exercised by text-only calibration + boi_token_index=2, + eoi_token_index=3, + dtype=torch.bfloat16, + ) + return AutoModelForImageTextToText.from_config(cfg) + + +def create_tiny_gemma3vl_dir( + tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs +) -> Path | tuple[Path, PreTrainedModel]: + gemma3vl_dir = Path(tmp_path) / "tiny_gemma3vl" + if with_tokenizer: + tokenizer = get_tiny_tokenizer() + tokenizer.save_pretrained(gemma3vl_dir) + config_kwargs["vocab_size"] = tokenizer.vocab_size + tiny_gemma3vl = get_tiny_gemma3vl(**config_kwargs) + tiny_gemma3vl.save_pretrained(gemma3vl_dir) + + if return_model: + return gemma3vl_dir, tiny_gemma3vl + else: + return gemma3vl_dir + + +##### Qwen3.5-VL (dense, hybrid GatedDeltaNet + gated attention) ##### +def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: + from transformers import Qwen3_5Config + + set_seed(SEED) + + # Dense Qwen3.5-VL language model: hybrid GatedDeltaNet (linear attention) + gated full-attention + # layers (layer_types auto-generates every-4th full attention). config_kwargs override the text side. + 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, + "vocab_size": 32, + # 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, + } + text_kwargs.update(config_kwargs) + vision_kwargs = { + "hidden_size": 16, + "intermediate_size": 16, + "depth": 2, + "num_heads": 2, + "patch_size": 16, + "out_hidden_size": text_kwargs["hidden_size"], # must match text hidden_size + # Qwen3.5-VL does not use deepstack injection (unlike Qwen3-VL); empty list => no deepstack + # mergers (the attribute must still be present for the bridge's vision config). + "deepstack_visual_indexes": [], + } + cfg = Qwen3_5Config(text_config=text_kwargs, vision_config=vision_kwargs, dtype=torch.bfloat16) + return AutoModelForImageTextToText.from_config(cfg) + + +def create_tiny_qwen3_5_vl_dir( + tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs +) -> Path | tuple[Path, PreTrainedModel]: + qwen3_5_vl_dir = Path(tmp_path) / "tiny_qwen3_5_vl" + if with_tokenizer: + tokenizer = get_tiny_tokenizer() + tokenizer.save_pretrained(qwen3_5_vl_dir) + config_kwargs["vocab_size"] = tokenizer.vocab_size + tiny_qwen3_5_vl = get_tiny_qwen3_5_vl(**config_kwargs) + tiny_qwen3_5_vl.save_pretrained(qwen3_5_vl_dir) + + if return_model: + return qwen3_5_vl_dir, tiny_qwen3_5_vl + else: + return qwen3_5_vl_dir ##### NEMOTRON ##### @@ -342,50 +466,6 @@ def create_tiny_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 - - ##### LLAMA ##### def get_tiny_llama(**config_kwargs) -> PreTrainedModel: set_seed(SEED) diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 33588234dec..9a0638d7a5f 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -18,20 +18,17 @@ import pytest 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_qwen3_5_vl_dir, + create_tiny_qwen3_dir, + create_tiny_qwen3vl_dir, +) +from transformers import AutoModelForCausalLM, AutoModelForImageTextToText -@pytest.mark.parametrize( - "create_tiny_model_dir", - [ - create_tiny_qwen3_dir, - # Gemma3 exercises the sliding/full attention ``layer_types`` pruning path. - create_tiny_gemma3_dir, - ], -) -def test_prune_minitron(tmp_path: Path, num_gpus, create_tiny_model_dir): - teacher_hf_path, teacher_model = create_tiny_model_dir( +def test_prune_minitron(tmp_path: Path, num_gpus): + teacher_hf_path, teacher_model = create_tiny_qwen3_dir( tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=num_gpus ) teacher_params = sum(p.numel() for p in teacher_model.parameters()) @@ -58,3 +55,59 @@ def test_prune_minitron(tmp_path: Path, num_gpus, create_tiny_model_dir): 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 + + +@pytest.mark.parametrize( + "create_tiny_vlm_dir", + [ + # Qwen3-VL: standard attention LM + deepstack vision injection (indices must be preserved + # during depth pruning and remapped to the renumbered layers). + create_tiny_qwen3vl_dir, + # Gemma3-VL: sliding/full attention LM (exercises the ``layer_types`` pruning path) with a + # multimodal projector (no deepstack). + create_tiny_gemma3vl_dir, + # Qwen3.5-VL (dense): hybrid GatedDeltaNet (linear attention) + gated full-attention LM. + create_tiny_qwen3_5_vl_dir, + ], +) +def test_prune_minitron_vlm(tmp_path: Path, num_gpus, create_tiny_vlm_dir): + teacher_hf_path, teacher_model = create_tiny_vlm_dir( + tmp_path, + with_tokenizer=True, + return_model=True, + num_hidden_layers=4 * num_gpus, # >= 4 per PP stage so depth is prunable and stays balanced + intermediate_size=128, + ) + 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="cnn_dailymail", + calib_num_samples=8, + seq_length=16, + 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() + + # from_pretrained (default strict load) verifies the saved weights match the pruned config. + pruned_model = AutoModelForImageTextToText.from_pretrained(pruned_model_path) + # Language model pruned to the param target; vision tower preserved. + pruned_lm_params = sum(p.numel() for p in pruned_model.model.language_model.parameters()) + assert pruned_lm_params <= prune_target_params + assert hasattr(pruned_model.config, "vision_config") + # deepstack indices (if any, e.g. Qwen3-VL) must remain valid for the (depth-pruned) LM. + deepstack = getattr(pruned_model.config.vision_config, "deepstack_visual_indexes", None) + if deepstack: + assert all(0 <= i < pruned_model.config.text_config.num_hidden_layers for i in deepstack) From 6863c224fe29f875ee1e5c28e1b471cab6397f6b Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 22 Jun 2026 13:51:15 -0700 Subject: [PATCH 02/12] Drop protected_layers and Qwen3-VL from VLM pruning support Qwen3.5-VL supersedes Qwen3-VL and removed deepstack vision injection, which was the only consumer of protected_layers (no other supported VLM injects vision features at specific LM layers). Remove the protected_layers search option and the deepstack index protection/remap for simplicity. - prune/plugins/mcore_minitron.py: remove the protected_layers config option, the sorted-layers reorder, and the depth-prune assert. - examples/megatron_bridge/prune_minitron.py: remove the deepstack protected_layers computation and the deepstack-index remap on save. - tests: drop the Qwen3-VL parametrization from test_prune_minitron_vlm (Gemma3-VL + Qwen3.5-VL remain; helper kept for the export test). - Docs: drop Qwen3-VL from the supported VLM lists; the example prunes a smaller Qwen3.5-4B for faster runs. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/megatron_bridge/README.md | 8 +++--- examples/megatron_bridge/prune_minitron.py | 25 ------------------- examples/pruning/README.md | 2 +- .../torch/prune/plugins/mcore_minitron.py | 25 ------------------- .../megatron_bridge/test_prune_minitron.py | 4 --- 6 files changed, 6 insertions(+), 60 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 93b2f66d531..df89320000f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -26,8 +26,8 @@ Changelog - **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-VL, 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. - 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 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. - 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``. - **Deduplicate the modules shared at source** in the quantized export step: ``_export_quantized_weight`` and ``_export_fused_experts`` now alias bit-identical packed ``weight`` / ``weight_scale`` / ``weight_scale_2`` buffers across modules sharing a source weight ``data_ptr()`` so the downstream ``postprocess_state_dict`` dedup catches them (~42% storage reduction on ``nvfp4_experts_only`` for tied 26B MoE checkpoints). diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index e1b7367731e..d13d40011e4 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -312,14 +312,14 @@ torchrun --nproc_per_node 1 prune_minitron.py --help ### Vision-Language Models (VLMs) -For a vision-language model (e.g. Qwen3-VL, 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 — the only difference is that `hidden_size` is never pruned for VLMs (it is shared with the vision projector). +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 — the only difference is that `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-VL-8B-Instruct \ - --prune_target_params 6e9 \ - --output_hf_path /tmp/Qwen3-VL-8B-Pruned + --hf_model_name_or_path Qwen/Qwen3.5-4B \ + --prune_target_params 3e9 \ + --output_hf_path /tmp/Qwen3.5-4B-Pruned-3B ``` ## Resources diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index e6ecedfd0ed..87dcd21a0fb 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -296,8 +296,6 @@ def main(args: argparse.Namespace): # hidden_size is shared with the vision->LM projector, so it is never pruned here (follow-up). language_model = getattr(unwrapped_model, "language_model", unwrapped_model) is_vlm = language_model is not unwrapped_model - # 1-indexed LM layer numbers that depth pruning must keep (see protected_layers handling below). - protected_layers: set[int] = set() if is_vlm: print_rank_0( "VLM detected: pruning model.language_model only (hidden_size pruning skipped)." @@ -308,14 +306,6 @@ def main(args: argparse.Namespace): ) args.hparams_to_skip = sorted({*args.hparams_to_skip, "hidden_size"}) - # Deepstack-injected LM layers (Qwen3-VL): the vision tower feeds features into specific LM - # layers (``vision_config.deepstack_visual_indexes``, 0-indexed). Those layers must survive - # depth pruning so the (unpruned) vision tower's deepstack mergers still line up; protect - # them (as 1-indexed layer numbers) and remap the indices after pruning. - vision_cfg = getattr(bridge.hf_pretrained.config, "vision_config", None) - deepstack_indexes = list(getattr(vision_cfg, "deepstack_visual_indexes", None) or []) - protected_layers = {idx + 1 for idx in deepstack_indexes} - # Calibration runs the pruning target (the language model) directly. For VLMs this is the # extracted ``language_model``; text-only calibration uses its own embedding and works under # pipeline parallelism via ``megatron_prefill`` (which handles inter-stage send/recv). @@ -332,9 +322,6 @@ def main(args: argparse.Namespace): pruning_config = { "forward_loop": forward_loop, "checkpoint": args.prune_intermediate_ckpt, - # Layers that depth pruning must keep (empty set => no constraint). For VLMs this protects - # the deepstack-injected LM layers so the vision tower stays consistent. - "protected_layers": protected_layers or None, } if args.prune_export_config is not None: # Less restrictive search space for manual pruning @@ -498,18 +485,6 @@ def score_func(m): text_cfg.layer_types = [ lt for i, lt in enumerate(text_cfg.layer_types) if i + 1 in kept_layer_nums ] - # Remap VLM deepstack indices (0-indexed LM layers) to the pruned/renumbered layers. The - # protected_layers constraint above guarantees every deepstack layer survived, so each index - # maps to its new compacted position; the vision tower's mergers stay aligned. - vision_cfg = getattr(hf_cfg, "vision_config", None) - if ( - is_vlm - and vision_cfg is not None - and getattr(vision_cfg, "deepstack_visual_indexes", None) - ): - vision_cfg.deepstack_visual_indexes = [ - kept_layer_nums.index(idx + 1) for idx in vision_cfg.deepstack_visual_indexes - ] if isinstance(provider, MambaModelProvider) and hasattr(hf_cfg, "hybrid_override_pattern"): hf_cfg.hybrid_override_pattern = getattr(unwrapped_model, hybrid_key) text_cfg.num_hidden_layers = mcore_cfg.num_layers diff --git a/examples/pruning/README.md b/examples/pruning/README.md index c43fc9a423c..91ba5ac866d 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -180,7 +180,7 @@ If your model parameters are already sorted and you just want to prune the weigh > *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.The language model of vision-language models (e.g. Qwen3-VL, 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).* +> *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.* diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 3ccc27df48d..64b8eb49f87 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -225,10 +225,6 @@ class MCoreMinitronSearcher(BaseSearcher): - `max_depth_pruning`: Maximum fraction per depth hyperparameter to prune (default: 0.20). Only top (1 - max_depth_pruning) choices will be considered. - `hparams_to_skip`: List of hparams to skip during the search (default: None). - - `protected_layers`: Set of 1-indexed layer numbers that must never be dropped during depth - pruning (default: None). Used e.g. for VLM deepstack-injection layers, which the vision - tower feeds into and so must be kept. Works for both ``export_config`` and metric-based - pruning. Requires the target ``num_layers`` >= ``len(protected_layers)``. - `top_k`: Number of candidates to consider for score_func validation (default: 10). - `seq_length`: Sequence length for KV-cache memory estimate (default: 4096). Only used with the ``memory_mb`` constraint. @@ -254,7 +250,6 @@ def default_search_config(self) -> SearchConfig: "max_width_pruning": 0.40, "max_depth_pruning": 0.20, "hparams_to_skip": None, - "protected_layers": None, "top_k": 10, # Memory footprint config (only used with memory_mb constraint) "seq_length": 4096, @@ -377,20 +372,6 @@ def run_search(self) -> None: for layer, _ in sorted(self.layer_scores.items(), key=lambda x: x[1], reverse=True) ] assert sorted(self.sorted_layers) == list(range(1, self.model.config.num_layers + 1)) - - # Force protected layers (e.g. VLM deepstack-injection layers) to the front of the - # importance ranking so depth pruning never drops them (drop set is sorted_layers[N:]). - # Relative order within each group is preserved (stable partition). - protected = self.config["protected_layers"] - if protected: - protected = set(protected) - assert protected <= set(self.sorted_layers), ( - f"protected_layers {sorted(protected)} must be 1-indexed layer numbers in " - f"1..{self.model.config.num_layers}" - ) - self.sorted_layers = [ln for ln in self.sorted_layers if ln in protected] + [ - ln for ln in self.sorted_layers if ln not in protected - ] else: assert ( self.constraints.keys() == {"export_config"} @@ -447,12 +428,6 @@ def _prune(self, export_config: dict, prune_depth: bool = True) -> None: if num_layers_hp.active != num_layers_hp.max: assert self.sorted_layers is not None layers_to_drop = self.sorted_layers[num_layers_hp.active :] - protected = self.config["protected_layers"] - if protected: - assert not (set(layers_to_drop) & set(protected)), ( - f"Cannot depth-prune to {num_layers_hp.active} layers while protecting " - f"{sorted(protected)}: target num_layers must be >= {len(set(protected))}." - ) drop_mcore_language_model_layers(self.model, layers_to_drop=layers_to_drop) # Update model config with pruned architecture diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 9a0638d7a5f..7c95e4f8826 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -22,7 +22,6 @@ create_tiny_gemma3vl_dir, create_tiny_qwen3_5_vl_dir, create_tiny_qwen3_dir, - create_tiny_qwen3vl_dir, ) from transformers import AutoModelForCausalLM, AutoModelForImageTextToText @@ -60,9 +59,6 @@ def test_prune_minitron(tmp_path: Path, num_gpus): @pytest.mark.parametrize( "create_tiny_vlm_dir", [ - # Qwen3-VL: standard attention LM + deepstack vision injection (indices must be preserved - # during depth pruning and remapped to the renumbered layers). - create_tiny_qwen3vl_dir, # Gemma3-VL: sliding/full attention LM (exercises the ``layer_types`` pruning path) with a # multimodal projector (no deepstack). create_tiny_gemma3vl_dir, From b3ee71a0e293c83978ee90a3d49b9782977b4067 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 22 Jun 2026 14:14:53 -0700 Subject: [PATCH 03/12] VLM pruning cleanups - Clarify in the README that VLM pruning targets (params / memory / export_config) apply to the language model only; the vision tower is not counted. - Drop the now-dead deepstack assertion from the VLM test and assert the non-language-model params (vision tower, projector, lm_head) are unchanged. - Simplify the layer_types kept-layer computation in prune_minitron.py. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/README.md | 5 ++++- examples/megatron_bridge/prune_minitron.py | 17 ++++++++--------- .../megatron_bridge/test_prune_minitron.py | 17 +++++++---------- 3 files changed, 19 insertions(+), 20 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index d13d40011e4..ff936b4eb24 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -312,7 +312,10 @@ torchrun --nproc_per_node 1 prune_minitron.py --help ### 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 — the only difference is that `hidden_size` is never pruned for VLMs (it is shared with the vision projector). +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 \ diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 87dcd21a0fb..89acbc55048 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -472,16 +472,15 @@ def score_func(m): text_cfg.n_shared_experts = ( mcore_cfg.moe_shared_expert_intermediate_size // mcore_cfg.moe_ffn_hidden_size ) - # 1-indexed original layer numbers kept after depth pruning, in ascending (original) order. - # sorted_layers is None when no layer scores were collected (depth pruning impossible), in - # which case all layers are kept and the mapping is the identity. - sorted_layers = pruning_scores["sorted_layers"] - kept_layer_nums = ( - sorted(sorted_layers[: mcore_cfg.num_layers]) - if sorted_layers is not None - else list(range(1, mcore_cfg.num_layers + 1)) - ) if hasattr(text_cfg, "layer_types"): + # Keep layer_types for the 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)) + ) text_cfg.layer_types = [ lt for i, lt in enumerate(text_cfg.layer_types) if i + 1 in kept_layer_nums ] diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 7c95e4f8826..31c1c526d74 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -59,11 +59,8 @@ def test_prune_minitron(tmp_path: Path, num_gpus): @pytest.mark.parametrize( "create_tiny_vlm_dir", [ - # Gemma3-VL: sliding/full attention LM (exercises the ``layer_types`` pruning path) with a - # multimodal projector (no deepstack). - create_tiny_gemma3vl_dir, - # Qwen3.5-VL (dense): hybrid GatedDeltaNet (linear attention) + gated full-attention LM. - create_tiny_qwen3_5_vl_dir, + create_tiny_gemma3vl_dir, # sliding/full attention LM with a multimodal projector + create_tiny_qwen3_5_vl_dir, # dense - hybrid GatedDeltaNet (linear attention) + gated full-attention LM ], ) def test_prune_minitron_vlm(tmp_path: Path, num_gpus, create_tiny_vlm_dir): @@ -99,11 +96,11 @@ def test_prune_minitron_vlm(tmp_path: Path, num_gpus, create_tiny_vlm_dir): # from_pretrained (default strict load) verifies the saved weights match the pruned config. pruned_model = AutoModelForImageTextToText.from_pretrained(pruned_model_path) - # Language model pruned to the param target; vision tower preserved. 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 left untouched. assert hasattr(pruned_model.config, "vision_config") - # deepstack indices (if any, e.g. Qwen3-VL) must remain valid for the (depth-pruned) LM. - deepstack = getattr(pruned_model.config.vision_config, "deepstack_visual_indexes", None) - if deepstack: - assert all(0 <= i < pruned_model.config.text_config.num_hidden_layers for i in deepstack) + teacher_non_lm = sum(p.numel() for p in teacher_model.parameters()) - language_model_params + pruned_non_lm = sum(p.numel() for p in pruned_model.parameters()) - pruned_lm_params + assert pruned_non_lm == teacher_non_lm From 4c485e15aed2fa2178ff5038461e639530e84e95 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 24 Jun 2026 09:54:14 -0700 Subject: [PATCH 04/12] Add VLM PTQ and image-text calibration for Megatron-Bridge examples - quantize.py: quantize only a VLM's language model (vision tower + projector stay full precision; ModelOpt state on the root so the Megatron save works). Calibration modality is inferred from --calib_dataset_name: an image-text dataset drives the full VLM forward (vision-conditioned activations); a text dataset runs text-only LM calibration (ablations). - prune_minitron.py: estimate pruning importance from image-text calibration (full VLM forward) by default, text-only otherwise; same dataset inference. - Shared get_megatron_vlm_calibration_forward_loop (megatron_prefill-based, unwraps tuple outputs) + vlm_dataset_utils (scienceqa, nemotron_vlm_dataset_v2 with config-driven subsets/shard cap). - Tiny VLM test fixtures (qwen3.5-vl, gemma3-vl) with vision tokens derived dynamically from the ref processor; VLM prune + quantize example tests. - README + CHANGELOG; TODO in export.py noting VLM HF export is unsupported. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 5 +- examples/megatron_bridge/README.md | 10 +- examples/megatron_bridge/export.py | 8 + examples/megatron_bridge/prune_minitron.py | 82 +++- examples/megatron_bridge/quantize.py | 94 ++++- .../utils/plugins/megatron_calibration.py | 85 +++- .../torch/utils/plugins/megatron_generate.py | 4 + modelopt/torch/utils/vlm_dataset_utils.py | 70 +++- .../_test_utils/torch/transformers_models.py | 377 +++++++++++------- .../megatron_bridge/test_prune_minitron.py | 27 +- .../megatron_bridge/test_quantize_export.py | 40 +- 11 files changed, 590 insertions(+), 212 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df89320000f..181828a87ca 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,13 +21,15 @@ 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 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. - 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``. - **Deduplicate the modules shared at source** in the quantized export step: ``_export_quantized_weight`` and ``_export_fused_experts`` now alias bit-identical packed ``weight`` / ``weight_scale`` / ``weight_scale_2`` buffers across modules sharing a source weight ``data_ptr()`` so the downstream ``postprocess_state_dict`` dedup catches them (~42% storage reduction on ``nvfp4_experts_only`` for tied 26B MoE checkpoints). @@ -35,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/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index ff936b4eb24..236c45f744e 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 diff --git a/examples/megatron_bridge/export.py b/examples/megatron_bridge/export.py index eecf7d838f4..5634a6e1959 100644 --- a/examples/megatron_bridge/export.py +++ b/examples/megatron_bridge/export.py @@ -144,6 +144,14 @@ def main(args: argparse.Namespace): print_rank_0( f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." ) + # TODO: Support exporting quantized VLMs. export_mcore_gpt_to_hf uses modelopt's own per-arch + # mcore<->HF mappings (mcore_qwen.py, ...), which don't cover Qwen3.5-VL / Gemma3-VL. Rather than + # authoring new per-model mappings, route the megatron->HF quant export through Megatron-Bridge's + # AutoBridge.export_hf_weights_quant(quantization_checker, quant_fn, quant_block_size): it reuses + # the bridge's per-model mapping (covering these VLMs plus the vision tower/projector, which the + # checker leaves in full precision), so modelopt only supplies the checker + pack/scale fn and the + # hf_quant_config (KV-cache scales need a separate path). quantize.py already saves the + # quantized-LM VLM Megatron checkpoint. 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 89acbc55048..ffb00db1c53 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -46,7 +46,12 @@ import torch from megatron.bridge import AutoBridge from megatron.bridge.models.mamba.mamba_provider import MambaModelProvider -from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoProcessor, +) import modelopt.torch.opt as mto import modelopt.torch.prune as mtp @@ -54,8 +59,16 @@ 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.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" def get_args() -> argparse.Namespace: @@ -92,10 +105,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( @@ -306,18 +322,50 @@ def main(args: argparse.Namespace): ) args.hparams_to_skip = sorted({*args.hparams_to_skip, "hidden_size"}) - # Calibration runs the pruning target (the language model) directly. For VLMs this is the - # extracted ``language_model``; text-only calibration uses its own embedding and works under - # pipeline parallelism via ``megatron_prefill`` (which handles inter-stage send/recv). - 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, - ) + # 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, + seq_length=args.seq_length, + 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, diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index cf25f6bfc2d..a247d2582dc 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,10 +161,13 @@ 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( @@ -274,8 +286,51 @@ 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: + print_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 pruning importance 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: + 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. @@ -294,16 +349,34 @@ def main(args: argparse.Namespace): # 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 mtq.need_calibration(mtq_config) and use_image_calib: + # 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, + seq_length=args.seq_length, + batch_size=args.calib_batch_size, + ) + elif mtq.need_calibration(mtq_config): + 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 @@ -349,13 +422,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/modelopt/torch/utils/plugins/megatron_calibration.py b/modelopt/torch/utils/plugins/megatron_calibration.py index 223215a2923..29f2a6dfc8a 100644 --- a/modelopt/torch/utils/plugins/megatron_calibration.py +++ b/modelopt/torch/utils/plugins/megatron_calibration.py @@ -25,13 +25,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 +58,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 +97,77 @@ 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, + seq_length: 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. + seq_length: Max text length for the processor (maps to ``max_length``). + 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, + max_length=seq_length, + require_image=True, + dp_size=dp_size, + dp_rank=mpu.get_data_parallel_rank(), + ) + + 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, 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 8eec92bafdf..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 ( diff --git a/modelopt/torch/utils/vlm_dataset_utils.py b/modelopt/torch/utils/vlm_dataset_utils.py index 2f1ae72b91f..b32a0992407 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) @@ -344,6 +365,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. @@ -355,6 +378,8 @@ def get_vlm_dataset_dataloader( 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 +429,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. @@ -463,4 +502,17 @@ 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: + if hasattr(dataset, "__len__"): + from torch.utils.data.distributed import DistributedSampler + + sampler = DistributedSampler(dataset, num_replicas=dp_size, rank=dp_rank, shuffle=False) + else: + dataset = _ShardedIterable(dataset, rank=dp_rank, world=dp_size) + + 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 8000bb2ac99..5241b0b45e2 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -25,6 +25,7 @@ AutoModelForCausalLM, AutoModelForImageTextToText, AutoModelForQuestionAnswering, + AutoProcessor, AutoTokenizer, BertConfig, DeepseekV3Config, @@ -58,6 +59,91 @@ def get_tiny_tokenizer(*, pad_side: str = "left") -> "transformers.PreTrainedTok return tokenizer +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_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]: + """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 = tokenizer_factory() + tokenizer.save_pretrained(dir_path) + config_kwargs["vocab_size"] = tokenizer.vocab_size + 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", + } + + 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 ##### def get_tiny_qwen3(**config_kwargs) -> PreTrainedModel: set_seed(SEED) @@ -72,28 +158,21 @@ def get_tiny_qwen3(**config_kwargs) -> PreTrainedModel: "max_position_embeddings": 32, "vocab_size": 32, } - kwargs.update(**config_kwargs) + 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 + return AutoModelForCausalLM.from_config(Qwen3Config(**kwargs)) def create_tiny_qwen3_dir( tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs ) -> Path | tuple[Path, PreTrainedModel]: - qwen3_dir = Path(tmp_path) / "tiny_qwen3" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(qwen3_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - tiny_qwen3 = get_tiny_qwen3(**config_kwargs) - tiny_qwen3.save_pretrained(qwen3_dir) - - if return_model: - return qwen3_dir, tiny_qwen3 - else: - return qwen3_dir + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_qwen3", + get_tiny_qwen3, + with_tokenizer=with_tokenizer, + return_model=return_model, + **config_kwargs, + ) ##### Qwen3 MoE ##### @@ -114,22 +193,19 @@ def get_tiny_qwen3_moe(**config_kwargs) -> PreTrainedModel: "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 + kwargs.update(config_kwargs) + return AutoModelForCausalLM.from_config(Qwen3MoeConfig(**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 + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_qwen3_moe", + get_tiny_qwen3_moe, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### Qwen3-VL ##### @@ -177,26 +253,28 @@ def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel: def create_tiny_qwen3vl_dir( tmp_path: Path | str, with_tokenizer: bool = False, return_model: bool = False, **config_kwargs ) -> Path | tuple[Path, PreTrainedModel]: - 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 - tiny_qwen3vl = get_tiny_qwen3vl(**config_kwargs) - tiny_qwen3vl.save_pretrained(qwen3vl_dir) - - if return_model: - return qwen3vl_dir, tiny_qwen3vl - else: - return qwen3vl_dir + 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) - # Gemma3 language model config (config_kwargs override the text side, e.g. num_hidden_layers). - # layer_types auto-generates (sliding/full attention) so that pruning code path is exercised. + # 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, @@ -205,59 +283,66 @@ def get_tiny_gemma3vl(**config_kwargs) -> PreTrainedModel: "num_attention_heads": 4, "num_key_value_heads": 2, "head_dim": 8, - "max_position_embeddings": 32, + "max_position_embeddings": 1024, # >= 256 image tokens + text for image calibration "sliding_window": 16, - "vocab_size": 32, + "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, - "image_size": 16, - "patch_size": 8, + # 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=4, # (image_size / patch_size) ** 2 = (16 / 8) ** 2 - image_token_index=1, # must be < vocab_size; not exercised by text-only calibration - boi_token_index=2, - eoi_token_index=3, + 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_tokenizer: bool = False, return_model: bool = False, **config_kwargs + tmp_path: Path | str, with_processor: bool = False, return_model: bool = False, **config_kwargs ) -> Path | tuple[Path, PreTrainedModel]: - gemma3vl_dir = Path(tmp_path) / "tiny_gemma3vl" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(gemma3vl_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - tiny_gemma3vl = get_tiny_gemma3vl(**config_kwargs) - tiny_gemma3vl.save_pretrained(gemma3vl_dir) - - if return_model: - return gemma3vl_dir, tiny_gemma3vl - else: - return gemma3vl_dir + 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, hybrid GatedDeltaNet + gated attention) ##### +# Real Qwen3.5-VL reused (via get_tiny_vlm_processor) for the fixture's image processor + chat +# template; the vision patch params below match it so the processor's pixel_values fit the tiny tower. +QWEN3_5_VL_REF = "Qwen/Qwen3.5-0.8B" + + def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: from transformers import Qwen3_5Config set_seed(SEED) - # Dense Qwen3.5-VL language model: hybrid GatedDeltaNet (linear attention) + gated full-attention - # layers (layer_types auto-generates every-4th full attention). config_kwargs override the text side. + # 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, @@ -267,7 +352,9 @@ def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: "num_key_value_heads": 2, "head_dim": 16, "max_position_embeddings": 32, - "vocab_size": 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, @@ -281,31 +368,38 @@ def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: "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 - # Qwen3.5-VL does not use deepstack injection (unlike Qwen3-VL); empty list => no deepstack - # mergers (the attribute must still be present for the bridge's vision config). - "deepstack_visual_indexes": [], + "deepstack_visual_indexes": [], # no deepstack injection (unlike Qwen3-VL) } - cfg = Qwen3_5Config(text_config=text_kwargs, vision_config=vision_kwargs, dtype=torch.bfloat16) + cfg = 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_tokenizer: bool = False, return_model: bool = False, **config_kwargs + tmp_path: Path | str, with_processor: bool = False, return_model: bool = False, **config_kwargs ) -> Path | tuple[Path, PreTrainedModel]: - qwen3_5_vl_dir = Path(tmp_path) / "tiny_qwen3_5_vl" - if with_tokenizer: - tokenizer = get_tiny_tokenizer() - tokenizer.save_pretrained(qwen3_5_vl_dir) - config_kwargs["vocab_size"] = tokenizer.vocab_size - tiny_qwen3_5_vl = get_tiny_qwen3_5_vl(**config_kwargs) - tiny_qwen3_5_vl.save_pretrained(qwen3_5_vl_dir) - - if return_model: - return qwen3_5_vl_dir, tiny_qwen3_5_vl - else: - return qwen3_5_vl_dir + return _create_tiny_vlm_dir( + Path(tmp_path) / "tiny_qwen3_5_vl", + QWEN3_5_VL_REF, + get_tiny_qwen3_5_vl, + with_processor=with_processor, + return_model=return_model, + **config_kwargs, + ) ##### NEMOTRON ##### @@ -324,20 +418,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) ##### @@ -373,20 +466,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 ##### @@ -414,7 +506,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"] @@ -424,13 +516,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 ##### @@ -448,22 +539,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 + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_gpt_oss", + get_tiny_gpt_oss, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) ##### LLAMA ##### @@ -479,23 +567,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 ##### @@ -513,22 +597,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 ##### @@ -544,17 +626,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 31c1c526d74..88c733aad5f 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -16,13 +16,8 @@ from pathlib import Path -import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.transformers_models import ( - create_tiny_gemma3vl_dir, - create_tiny_qwen3_5_vl_dir, - create_tiny_qwen3_dir, -) +from _test_utils.torch.transformers_models import create_tiny_gemma3vl_dir, create_tiny_qwen3_dir from transformers import AutoModelForCausalLM, AutoModelForImageTextToText @@ -56,20 +51,16 @@ def test_prune_minitron(tmp_path: Path, num_gpus): assert pruned_params <= prune_target_params -@pytest.mark.parametrize( - "create_tiny_vlm_dir", - [ - create_tiny_gemma3vl_dir, # sliding/full attention LM with a multimodal projector - create_tiny_qwen3_5_vl_dir, # dense - hybrid GatedDeltaNet (linear attention) + gated full-attention LM - ], -) -def test_prune_minitron_vlm(tmp_path: Path, num_gpus, create_tiny_vlm_dir): - teacher_hf_path, teacher_model = create_tiny_vlm_dir( +def test_prune_minitron_vlm(tmp_path: Path, num_gpus): + # gemma3vl: sliding/full attention LM with a multimodal projector (qwen3_5_vl is covered by the + # VLM quantize test instead, to keep each VLM exercised once across the two examples). + teacher_hf_path, teacher_model = create_tiny_gemma3vl_dir( tmp_path, - with_tokenizer=True, + with_processor=True, return_model=True, num_hidden_layers=4 * num_gpus, # >= 4 per PP stage so depth is prunable and stays balanced intermediate_size=128, + max_position_embeddings=1024, # >= 256 image tokens + text, no truncation ) language_model_params = sum(p.numel() for p in teacher_model.model.language_model.parameters()) prune_target_params = int(language_model_params * 0.7) @@ -80,9 +71,9 @@ def test_prune_minitron_vlm(tmp_path: Path, num_gpus, create_tiny_vlm_dir): hf_model_name_or_path=teacher_hf_path, output_hf_path=pruned_model_path, pp_size=num_gpus, - calib_dataset_name="cnn_dailymail", + calib_dataset_name="scienceqa", # image-text calibration runs the full VLM forward calib_num_samples=8, - seq_length=16, + seq_length=1024, prune_target_params=prune_target_params, prune_score_func="mmlu_1pct_bs32", ss_channel_divisor=4, 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" + ) From 89b3c49fac771e164e429dfa9fe16604ff7b18de Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 25 Jun 2026 02:37:48 -0700 Subject: [PATCH 05/12] Address PR review comments and fix VLM pruning correctness bugs Review-comment fixes: - vlm_dataset_utils: discriminate DP sharding on IterableDataset type rather than __len__ (the streaming wrapper defines __len__), so DP>1 calibration works for the default streaming VLM dataset instead of crashing in DataLoader. - mcore_minitron: reset hparam choices before validating export_config so ConcatHparam choices are refreshed after modify(). - quantize: anchor the non-LM quantizer-disable pattern to the child subtree (f"{name}.*") so a short non-LM name cannot match an LM quantizer path; fix copy-pasted "pruning importance" wording (calibration, not pruning). - transformers_models: document the Qwen3.5-VL lazy import. - test_prune_minitron_vlm: compare the vision tower/projector/lm_head by value (byte-identical), not just counts. Bug fixes: - mcore_minitron: do not sort (importance-permute) hparams listed in hparams_to_skip. For VLMs hidden_size is skipped because the language model shares its residual dimension with the un-pruned vision projector; permuting it (even when hidden_size is not reduced) misaligned the injected image features and produced a permuted lm_head/embedding. Now the residual order is preserved and the pruned VLM stays vision-aligned. - prune_minitron: write the pruned shared-expert size back under both HF field names (Qwen3.5-MoE uses shared_expert_intermediate_size), so the exported config matches the pruned weights and reloads via from_pretrained. Test-util refactor: - Merge the dense/MoE tiny-model builders behind a moe flag exposed via partial aliases (get_tiny_qwen3[/_moe], create_tiny_qwen3_dir[/_moe], and the Qwen3.5-VL equivalents), with no caller churn. Verified test_prune_minitron_vlm[gemma3vl] and [qwen3_5_moe_vl] pass on 2 GPUs (nemo:26.06) with byte-identical vision tower, projector, and lm_head. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/prune_minitron.py | 16 ++- examples/megatron_bridge/quantize.py | 6 +- .../torch/prune/plugins/mcore_minitron.py | 10 +- modelopt/torch/utils/plugins/mbridge.py | 58 +++++++-- modelopt/torch/utils/vlm_dataset_utils.py | 8 +- .../_test_utils/torch/transformers_models.py | 108 ++++++++-------- .../megatron_bridge/test_prune_minitron.py | 116 +++++++++++++----- 7 files changed, 223 insertions(+), 99 deletions(-) diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index ffb00db1c53..09c10e0f2bf 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -309,7 +309,7 @@ def main(args: argparse.Namespace): ) # 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 never pruned here (follow-up). + # 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: @@ -508,10 +508,16 @@ def score_func(m): 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 - if hasattr(text_cfg, "moe_shared_expert_intermediate_size"): - text_cfg.moe_shared_expert_intermediate_size = ( - mcore_cfg.moe_shared_expert_intermediate_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"): diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index a247d2582dc..5c9051177c0 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -311,7 +311,7 @@ def main(args: argparse.Namespace): 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." + "model's calibration statistics will not see vision tokens." ) print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}") @@ -328,7 +328,9 @@ def main(args: argparse.Namespace): if name != "language_model" and id(child) not in lm_module_ids ) for name in non_lm_children: - mtq_config["quant_cfg"].append({"quantizer_name": f"*{name}*", "enable": False}) + # 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* diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 64b8eb49f87..95a9cfc79fc 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -311,11 +311,16 @@ 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 all parameters that may be pruned. Skip the ones explicitly excluded from the + # search: sorting permutes a dimension by importance, and for hparams that are never + # pruned this permutation is pure churn. It is also unsafe for ``hidden_size`` on VLMs -- + # the language model shares its residual dimension with the (un-pruned) vision projector, + # so permuting it would misalign the 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}!" @@ -329,7 +334,6 @@ def before_search(self) -> None: f"{hp.choices}. Manual export_config values must match the search-space " "granularity (see the *_divisor settings)." ) - hp.reset_choices() # Make sure ConcatHparam choices are updated after modify() assert isinstance(self.model, _DynamicMCoreLanguageModel), ( "Input should be unwrapped MCore model!" diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index c671bc338b0..c7fea098dca 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -27,7 +27,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 +38,49 @@ __all__ = ["load_mbridge_model_from_hf", "load_modelopt_megatron_checkpoint"] +def _patch_qwen35_moe_sequential_expert_mappings() -> None: + """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 (26.06.01 onwards). + """ + try: + from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping + 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) diff --git a/modelopt/torch/utils/vlm_dataset_utils.py b/modelopt/torch/utils/vlm_dataset_utils.py index b32a0992407..6b945b7b55e 100644 --- a/modelopt/torch/utils/vlm_dataset_utils.py +++ b/modelopt/torch/utils/vlm_dataset_utils.py @@ -506,12 +506,14 @@ def _collate_fn(examples: list[dict[str, Any]]) -> dict[str, torch.Tensor] | dic # iterable/streaming ones (each rank then processes ~num_samples / dp_size samples). sampler = None if dp_size > 1: - if hasattr(dataset, "__len__"): + # 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) - else: - dataset = _ShardedIterable(dataset, rank=dp_rank, world=dp_size) 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 5241b0b45e2..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 @@ -144,8 +145,8 @@ def _create_tiny_vlm_dir( return (dir_path, model) if return_model else dir_path -##### Qwen3 ##### -def get_tiny_qwen3(**config_kwargs) -> PreTrainedModel: +##### Qwen3 (dense or MoE) ##### +def _get_tiny_qwen3(moe: bool = False, **config_kwargs) -> PreTrainedModel: set_seed(SEED) kwargs = { @@ -158,54 +159,42 @@ def get_tiny_qwen3(**config_kwargs) -> PreTrainedModel: "max_position_embeddings": 32, "vocab_size": 32, } + 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 Qwen3ForCausalLM() for correct dtype handling - return AutoModelForCausalLM.from_config(Qwen3Config(**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, **config_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", - get_tiny_qwen3, + 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, ) -##### Qwen3 MoE ##### -def get_tiny_qwen3_moe(**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) - return AutoModelForCausalLM.from_config(Qwen3MoeConfig(**kwargs)) - - -def create_tiny_qwen3_moe_dir( - tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs -) -> Path | tuple[Path, PreTrainedModel]: - return _create_tiny_llm_dir( - Path(tmp_path) / "tiny_qwen3_moe", - get_tiny_qwen3_moe, - with_tokenizer=with_tokenizer, - **config_kwargs, - ) +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 ##### @@ -328,14 +317,13 @@ def create_tiny_gemma3vl_dir( ) -##### Qwen3.5-VL (dense, hybrid GatedDeltaNet + gated attention) ##### -# Real Qwen3.5-VL reused (via get_tiny_vlm_processor) for the fixture's image processor + chat -# template; the vision patch params below match it so the processor's pixel_values fit the tiny tower. +##### 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(**config_kwargs) -> PreTrainedModel: - from transformers import Qwen3_5Config +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) @@ -362,6 +350,16 @@ def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: "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, @@ -376,7 +374,7 @@ def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: "out_hidden_size": text_kwargs["hidden_size"], # must match text hidden_size "deepstack_visual_indexes": [], # no deepstack injection (unlike Qwen3-VL) } - cfg = Qwen3_5Config( + cfg = (Qwen3_5MoeConfig if moe else Qwen3_5Config)( text_config=text_kwargs, vision_config=vision_kwargs, dtype=torch.bfloat16, @@ -389,19 +387,31 @@ def get_tiny_qwen3_5_vl(**config_kwargs) -> PreTrainedModel: return AutoModelForImageTextToText.from_config(cfg) -def create_tiny_qwen3_5_vl_dir( - tmp_path: Path | str, with_processor: bool = False, return_model: bool = False, **config_kwargs +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_vl", + Path(tmp_path) / ("tiny_qwen3_5_moe_vl" if moe else "tiny_qwen3_5_vl"), QWEN3_5_VL_REF, - get_tiny_qwen3_5_vl, + _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 ##### def get_tiny_nemotron(**config_kwargs) -> PreTrainedModel: set_seed(SEED) diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 88c733aad5f..0a0503a5a53 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -14,25 +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_gemma3vl_dir, create_tiny_qwen3_dir +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 -def test_prune_minitron(tmp_path: Path, num_gpus): - teacher_hf_path, teacher_model = create_tiny_qwen3_dir( - tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=num_gpus - ) +@pytest.mark.parametrize( + ("create_teacher", "megatron_format"), + [ + # 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, 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, @@ -42,26 +71,51 @@ def test_prune_minitron(tmp_path: Path, num_gpus): 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") - 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 + 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 -def test_prune_minitron_vlm(tmp_path: Path, num_gpus): - # gemma3vl: sliding/full attention LM with a multimodal projector (qwen3_5_vl is covered by the - # VLM quantize test instead, to keep each VLM exercised once across the two examples). - teacher_hf_path, teacher_model = create_tiny_gemma3vl_dir( - tmp_path, - with_processor=True, - return_model=True, - num_hidden_layers=4 * num_gpus, # >= 4 per PP stage so depth is prunable and stays balanced - intermediate_size=128, - max_position_embeddings=1024, # >= 256 image tokens + text, no truncation - ) +@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) @@ -90,8 +144,14 @@ def test_prune_minitron_vlm(tmp_path: Path, num_gpus): 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 left untouched. + # Everything outside the language model (vision tower, projector, lm_head) is byte-identical assert hasattr(pruned_model.config, "vision_config") - teacher_non_lm = sum(p.numel() for p in teacher_model.parameters()) - language_model_params - pruned_non_lm = sum(p.numel() for p in pruned_model.parameters()) - pruned_lm_params - assert pruned_non_lm == teacher_non_lm + 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 + ) From db4d052d45ee210197189e1e9f9a393c8d33c6f2 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 25 Jun 2026 05:07:35 -0700 Subject: [PATCH 06/12] Clarify VLM prune target scope and remap deepstack indices on export - Log a language-model / frozen-non-LM / full-model param breakdown (via num2hrb) before and after pruning, and clarify in the --prune_target_* help text and the runtime message that the target applies to the language-model tower only for VLMs (vision/audio encoders, projectors, etc. are frozen and excluded). Param counts use the dist utility (allreduce) and correct for tied embeddings under PP. - Remap vision_config.deepstack_visual_indexes (Qwen3-VL) to the surviving layers after depth pruning so the exported config stays valid: renumber for dropped lower layers and snap a dropped injection layer to the nearest survivor (count preserved so the frozen vision projector still lines up). Warn -- recommending full VLM training over LM-only -- when a deepstack layer is dropped, since text-only distillation cannot recover the vision path. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/prune_minitron.py | 96 +++++++++++++++++++--- 1 file changed, 83 insertions(+), 13 deletions(-) diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 09c10e0f2bf..7156e575d00 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -57,7 +57,13 @@ 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, @@ -146,7 +152,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( @@ -155,7 +162,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( @@ -165,7 +173,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( @@ -280,6 +289,43 @@ 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 pipeline parallelism a tied embedding is materialized on both the first (word_embeddings) + # and last (output_layer) stage, so the sum double-counts it; subtract one copy. Only the first + # stage owns ``word_embeddings``, so the allreduce-sum below yields exactly one copy. + 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." @@ -314,13 +360,17 @@ def main(args: argparse.Namespace): is_vlm = language_model is not unwrapped_model if is_vlm: print_rank_0( - "VLM detected: pruning model.language_model only (hidden_size pruning skipped)." + "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: @@ -457,6 +507,8 @@ def score_func(m): # Remove unnecessary modelopt_state since ckpt is homogeneous 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" @@ -526,18 +578,36 @@ def score_func(m): text_cfg.n_shared_experts = ( mcore_cfg.moe_shared_expert_intermediate_size // mcore_cfg.moe_ffn_hidden_size ) + # 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"): - # Keep layer_types for the 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)) - ) 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 (0-indexed) LM layers. Remap those + # indices to the surviving layers so the exported config stays valid (no out-of-range index); + # a dropped injection layer snaps to the nearest survivor below it (count is preserved so the + # frozen vision projector still lines up). + 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) text_cfg.num_hidden_layers = mcore_cfg.num_layers From b47f8833a6e0f969c9c9a32f7703cba8db41963a Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 25 Jun 2026 13:35:06 -0700 Subject: [PATCH 07/12] disable qwen35_moe test Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/utils/dataset_utils.py | 2 +- .../megatron_bridge/test_prune_minitron.py | 22 +++++++++---------- 2 files changed, 12 insertions(+), 12 deletions(-) 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/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 0a0503a5a53..77bc3b852f6 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -20,7 +20,6 @@ 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 @@ -99,17 +98,18 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): ), id="gemma3vl", ), + # TODO: Enable once CI uses the nemo:26.06 container # 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", - ), + # 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): From 009e448312b5ead04ee58b81ab346231564e1184 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 26 Jun 2026 12:33:59 -0700 Subject: [PATCH 08/12] Re-enable Qwen3.5MoE pruning test Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/utils/plugins/mbridge.py | 6 ++--- .../megatron_bridge/test_prune_minitron.py | 22 +++++++++---------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index c7fea098dca..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 @@ -39,16 +40,15 @@ def _patch_qwen35_moe_sequential_expert_mappings() -> None: - """Add sequential (non-grouped) expert mappings to Megatron-Bridge's Qwen3.5 MoE bridge. + """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 (26.06.01 onwards). + TODO: Remove once Megatron-Bridge maps sequential Qwen3.5 MoE experts natively (patched in 26.06.01). """ try: - from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping from megatron.bridge.models.qwen.qwen35_bridge import Qwen35MoEBridge except ImportError: return diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 77bc3b852f6..0a0503a5a53 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -20,6 +20,7 @@ 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 @@ -98,18 +99,17 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): ), id="gemma3vl", ), - # TODO: Enable once CI uses the nemo:26.06 container # 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", - # ), + 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): From cc5dbc801b236f336e29817f0a745053f1b9d19e Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 29 Jun 2026 03:55:18 -0700 Subject: [PATCH 09/12] Address PR review: trim verbose comments, reorder quantize calib branch - Trim verbose comments in prune_minitron.py, mcore_minitron.py, export.py (PP tied-embedding, export-config validation, hidden_size-sort, deepstack) - quantize.py: put the text/LLM calibration branch first, VLM image-text last - export.py: shorten quantized-VLM-export TODO and link OMNIML-5366 Co-Authored-By: Claude Opus 4.8 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/export.py | 11 ++---- examples/megatron_bridge/prune_minitron.py | 11 +++--- examples/megatron_bridge/quantize.py | 34 +++++++++---------- .../torch/prune/plugins/mcore_minitron.py | 14 +++----- 4 files changed, 29 insertions(+), 41 deletions(-) diff --git a/examples/megatron_bridge/export.py b/examples/megatron_bridge/export.py index 5634a6e1959..926b47755e5 100644 --- a/examples/megatron_bridge/export.py +++ b/examples/megatron_bridge/export.py @@ -144,14 +144,9 @@ def main(args: argparse.Namespace): print_rank_0( f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." ) - # TODO: Support exporting quantized VLMs. export_mcore_gpt_to_hf uses modelopt's own per-arch - # mcore<->HF mappings (mcore_qwen.py, ...), which don't cover Qwen3.5-VL / Gemma3-VL. Rather than - # authoring new per-model mappings, route the megatron->HF quant export through Megatron-Bridge's - # AutoBridge.export_hf_weights_quant(quantization_checker, quant_fn, quant_block_size): it reuses - # the bridge's per-model mapping (covering these VLMs plus the vision tower/projector, which the - # checker leaves in full precision), so modelopt only supplies the checker + pack/scale fn and the - # hf_quant_config (KV-cache scales need a separate path). quantize.py already saves the - # quantized-LM VLM Megatron checkpoint. + # 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 7156e575d00..3255b8893e3 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -304,9 +304,8 @@ def _local(module) -> int: total = dist.allreduce(_local(unwrapped_model)) # sum across pipeline ranks lm = dist.allreduce(_local(language_model)) - # Under pipeline parallelism a tied embedding is materialized on both the first (word_embeddings) - # and last (output_layer) stage, so the sum double-counts it; subtract one copy. Only the first - # stage owns ``word_embeddings``, so the allreduce-sum below yields exactly one copy. + # 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( @@ -590,10 +589,8 @@ def score_func(m): 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 (0-indexed) LM layers. Remap those - # indices to the surviving layers so the exported config stays valid (no out-of-range index); - # a dropped injection layer snaps to the nearest survivor below it (count is preserved so the - # frozen vision projector still lines up). + # 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: diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 5c9051177c0..3a976d78734 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -351,21 +351,10 @@ def main(args: argparse.Namespace): # 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) and use_image_calib: - # 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, - seq_length=args.seq_length, - batch_size=args.calib_batch_size, - ) - elif mtq.need_calibration(mtq_config): + 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, @@ -380,8 +369,19 @@ def main(args: argparse.Namespace): 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, + seq_length=args.seq_length, + batch_size=args.calib_batch_size, + ) if hasattr(unwrapped_model, "calibration_mode"): # Some model wrappers (e.g. distillation/speculative) gate calibration behind a flag. diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 95a9cfc79fc..cee43d85807 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -311,11 +311,9 @@ 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 that may be pruned. Skip the ones explicitly excluded from the - # search: sorting permutes a dimension by importance, and for hparams that are never - # pruned this permutation is pure churn. It is also unsafe for ``hidden_size`` on VLMs -- - # the language model shares its residual dimension with the (un-pruned) vision projector, - # so permuting it would misalign the injected image features. + # 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): @@ -324,10 +322,8 @@ def before_search(self) -> None: 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}!" - # Validate requested export_config values are achievable for *every* matching hparam, - # including non-configurable ones. A value outside `choices` (e.g. not aligned to the - # search-space divisor) would otherwise be silently ignored while model.config is still - # overwritten, producing a checkpoint whose weights don't match its config. + # 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: " From 878a74eff32dc220c2a16561ce0a6a780dd501c0 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 29 Jun 2026 12:52:34 -0700 Subject: [PATCH 10/12] Drop MTP layers before pruning Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/README.md | 3 +++ examples/megatron_bridge/prune_minitron.py | 28 ++++++++++++++++++++++ examples/pruning/README.md | 10 ++++---- 3 files changed, 37 insertions(+), 4 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 236c45f744e..60dfed7fc7d 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -315,6 +315,9 @@ 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. diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 3255b8893e3..8fb444006bf 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -76,6 +76,18 @@ 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: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) @@ -348,11 +360,23 @@ 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, ) + # 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) @@ -608,6 +632,10 @@ def score_func(m): if isinstance(provider, MambaModelProvider) and hasattr(hf_cfg, "hybrid_override_pattern"): hf_cfg.hybrid_override_pattern = getattr(unwrapped_model, hybrid_key) 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 dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM diff --git a/examples/pruning/README.md b/examples/pruning/README.md index 91ba5ac866d..332d5cbb83e 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -174,8 +174,8 @@ 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 LLMs (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 & 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.* @@ -184,9 +184,11 @@ If your model parameters are already sorted and you just want to prune the weigh > *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 From 21cf2f9db448bbbf321d21bc44a18a2e6d95b306 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 29 Jun 2026 12:58:52 -0700 Subject: [PATCH 11/12] Show warnings in yellow color Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/prune_minitron.py | 2 +- examples/megatron_bridge/quantize.py | 5 ++--- modelopt/torch/utils/logging.py | 3 +++ 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 8fb444006bf..532fc4b3568 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -382,7 +382,7 @@ def main(args: argparse.Namespace): language_model = getattr(unwrapped_model, "language_model", unwrapped_model) is_vlm = language_model is not unwrapped_model if is_vlm: - print_rank_0( + 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 " diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 3a976d78734..01075c7b414 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -290,7 +290,7 @@ def main(args: argparse.Namespace): language_model = getattr(unwrapped_model, "language_model", unwrapped_model) is_vlm = language_model is not unwrapped_model if is_vlm: - print_rank_0( + warn_rank_0( "VLM detected: quantizing `model.language_model` only (vision tower left in full precision)." ) @@ -345,8 +345,7 @@ 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 diff --git a/modelopt/torch/utils/logging.py b/modelopt/torch/utils/logging.py index 0c560f1b2ce..8bca33d5620 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 + # Render the message in yellow on a terminal (skip escape codes for redirected output / logs). + if isinstance(message, str) and sys.stderr.isatty(): + message = f"\033[33m{message}\033[0m" warnings.warn(message, *args, **kwargs) From 335ac967a81854faa46cc338d119222b74249bb6 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 29 Jun 2026 14:49:44 -0700 Subject: [PATCH 12/12] Cleanup VLM dataset utils Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/hf_ptq/hf_ptq.py | 1 - examples/megatron_bridge/prune_minitron.py | 8 ++++++-- examples/megatron_bridge/quantize.py | 8 ++++++-- modelopt/torch/utils/logging.py | 2 +- modelopt/torch/utils/plugins/megatron_calibration.py | 8 ++++---- modelopt/torch/utils/vlm_dataset_utils.py | 6 +----- 6 files changed, 18 insertions(+), 15 deletions(-) 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/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 532fc4b3568..bdd40b3ad8f 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -137,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", @@ -427,7 +432,6 @@ def main(args: argparse.Namespace): processor, dataset_name=args.calib_dataset_name, num_samples=args.calib_num_samples, - seq_length=args.seq_length, batch_size=args.calib_batch_size, ) else: diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 01075c7b414..d57e74277c6 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -174,7 +174,12 @@ def get_args() -> argparse.Namespace: "--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( @@ -378,7 +383,6 @@ def forward_loop(_model=None): processor, dataset_name=args.calib_dataset_name, num_samples=args.calib_num_samples, - seq_length=args.seq_length, batch_size=args.calib_batch_size, ) diff --git a/modelopt/torch/utils/logging.py b/modelopt/torch/utils/logging.py index 8bca33d5620..d1a9fc1aef9 100644 --- a/modelopt/torch/utils/logging.py +++ b/modelopt/torch/utils/logging.py @@ -130,7 +130,7 @@ def warn_rank_0(message, *args, **kwargs): """ if dist.is_master(): kwargs["stacklevel"] = kwargs.get("stacklevel", 1) + 1 - # Render the message in yellow on a terminal (skip escape codes for redirected output / logs). + # 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/megatron_calibration.py b/modelopt/torch/utils/plugins/megatron_calibration.py index 29f2a6dfc8a..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 @@ -106,7 +107,6 @@ def get_megatron_vlm_calibration_forward_loop( dataset_name: str = "scienceqa", batch_size: int = 1, num_samples: int = 512, - seq_length: int = 512, device: torch.device | str | None = "cuda", subsets: list[str] | None = None, max_shards: int | None = None, @@ -124,7 +124,6 @@ def get_megatron_vlm_calibration_forward_loop( dataset_name: VLM calibration dataset name (see ``vlm_dataset_utils``). batch_size: Calibration batch size. num_samples: Number of calibration samples. - seq_length: Max text length for the processor (maps to ``max_length``). 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``; @@ -143,11 +142,12 @@ def get_megatron_vlm_calibration_forward_loop( batch_size=batch_size, num_samples=num_samples, device=device, - max_length=seq_length, 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 @@ -160,7 +160,7 @@ def _forward_loop(_model: torch.nn.Module | None = None) -> None: 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, disable=not dist.is_master()): + for batch in tqdm(dataloader, total=total_batches, disable=not dist.is_master()): megatron_prefill( model, batch["input_ids"], diff --git a/modelopt/torch/utils/vlm_dataset_utils.py b/modelopt/torch/utils/vlm_dataset_utils.py index 6b945b7b55e..71240968b60 100644 --- a/modelopt/torch/utils/vlm_dataset_utils.py +++ b/modelopt/torch/utils/vlm_dataset_utils.py @@ -357,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, @@ -376,7 +375,6 @@ 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. @@ -485,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.