From 9aa05fe877b5c8a11b78dff3323d6d87d18067cc Mon Sep 17 00:00:00 2001 From: skierat Date: Thu, 30 Jul 2026 03:15:28 +0200 Subject: [PATCH 01/11] add Qwen3-VL support for DFlash training (#1975) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: new feature Adds online DFlash training support for Qwen3-VL–style vision-language models. Changes include: - Load VLMs through the Transformers 5 `AutoModelForImageTextToText` API, while retaining compatibility with the legacy VLM auto-model API. - Run the base model through its top-level multimodal forward when image/video inputs are present, ensuring vision embeddings are injected before collecting DFlash target hidden states. - Extend `VisionLanguageDataCollator` to: - propagate `answer_only_loss`, chat-template, and DFlash label-alignment settings; - apply `VLM_MIN_PIXELS` / `VLM_MAX_PIXELS` processor limits; - derive assistant-only masks from ChatML/Llama chat boundaries when processor generation masks are unavailable; - enforce the fixed `training_seq_len` required by DFlash block training. - Preserve the existing text-only DFlash path. ### Usage ```bash python -m torch.distributed.run \ --nproc_per_node 4 \ examples/speculative_decoding/main.py \ --config modelopt_recipes/general/speculative_decoding/dflash.yaml \ model.model_name_or_path=/path/to/qwen3-vl-model \ model.trust_remote_code=true \ data.data_path=/path/to/train.jsonl \ data.vlm_processor=/path/to/qwen3-vl-model \ data.vlm_img_dir=/path/to/image/root \ training.training_seq_len=4096 \ training.answer_only_loss=true \ dflash.dflash_block_size=8 \ dflash.dflash_mask_token_id=151669 ### Testing - git diff --check - Parsed all modified Python modules successfully. - Ran iterative multi-node Slurm smoke tests with a Qwen3-VL-family model and mixed multimodal data: - validated VLM model loading with Transformers 5; - validated distributed initialization, DFlash conversion, and VLM collation paths; - identified and addressed processor padding/truncation behavior required by fixed-size DFlash blocks. This PR remains draft pending a completed end-to-end training smoke test and automated regression coverage. ### Before your PR is "Ready for review" Make sure you read and follow Contributor guidelines (https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (git commit -s -S). Make sure you read and follow the Security Best Practices (https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A - Did you write any new necessary tests?: ❌ — automated Qwen3-VL/DFlash regression coverage still needs to be added before review. - Did you update Changelog (https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — evaluate and add an entry before marking ready for review if this is considered user-facing speculative-decoding support. - Did you get Claude approval on this PR?: N/A ### Additional Information The PR intentionally excludes local Slurm launch scripts, logs, model paths, datasets, and environment-specific configuration. ## Summary by CodeRabbit * **New Features** * Expanded VLM data-collation controls, including `shift_labels` and more robust `answer_only_loss` masking. * Improved Qwen3-VL speculative decoding for Transformers 5.3+ with correct video frame grouping and safer position-id handling. * Improved DFlash RoPE export to reliably read `rope_theta` from newer config formats. * **Bug Fixes** * Hardened multimodal preprocessing and training loss masking to keep label/attention alignment consistent. * Improved behavior when anchor sampling yields no valid blocks. * More resilient VLM model loading when certain Transformers auto classes are unavailable. * **Tests** * Added coverage for RoPE export, Qwen3-VL position-id logic across Transformers versions, and VLM label-mode/collator options. --------- Signed-off-by: Slawomir Kierat Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-authored-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/speculative_decoding/eagle_utils.py | 3 + .../torch/export/plugins/hf_spec_export.py | 29 ++- .../torch/speculative/plugins/hf_dflash.py | 243 ++++++++++++++++-- modelopt/torch/speculative/utils.py | 8 +- .../utils/plugins/transformers_dataset.py | 132 +++++++++- .../torch/export/test_hf_spec_rope_export.py | 13 + .../speculative/plugins/test_fakebase.py | 28 ++ .../speculative/plugins/test_hf_dflash.py | 217 ++++++++++++++++ .../plugins/test_hf_speculative_offline.py | 55 +++- 9 files changed, 701 insertions(+), 27 deletions(-) diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index b12b9da1a52..68c6db45235 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -141,6 +141,9 @@ def make_speculative_data_module( train_len=train_len, local_image_path=data_args.vlm_img_dir, return_labels=True, + answer_only_loss=answer_only_loss, + shift_labels=shift_labels, + chat_template=chat_template, ) else: diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index caa93db3634..255b1d9ab04 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -29,6 +29,23 @@ ALL_SPEC_MODES = ["eagle", "dflash"] + +def _get_rope_theta(config, default=None): + """Get RoPE theta from either legacy or Transformers 5 config fields.""" + rope_theta = getattr(config, "rope_theta", None) + if rope_theta is not None: + return rope_theta + + # Transformers 5 stores this under rope_parameters (and exposes the same + # data through rope_scaling for backwards compatibility). + for attr in ("rope_parameters", "rope_scaling"): + rope_config = getattr(config, attr, None) + if isinstance(rope_config, dict) and rope_config.get("rope_theta") is not None: + return rope_config["rope_theta"] + + return default + + LLAMA_EAGLE_SINGLE_LAYER = { "required": { "layers.0.self_attn.q_proj", @@ -376,14 +393,10 @@ def _export_config(self): "initializer_range": getattr(base_config, "initializer_range", 0.02), "attention_bias": getattr(draft_config, "attention_bias", False), "attention_dropout": getattr(draft_config, "attention_dropout", 0.0), - # Inherit the target's rope_theta: DFlash injects the target's KV into every - # draft layer, so the draft's RoPE base must match the target's. (The draft - # arch config carries no rope_theta of its own.) - "rope_theta": ( - getattr(base_config, "rope_theta", None) - if getattr(base_config, "rope_theta", None) is not None - else getattr(draft_config, "rope_theta", 1000000.0) - ), + # Inherit the target's RoPE base: DFlash injects target KV into every draft + # layer, so their RoPE bases must match. Transformers 5 stores rope_theta + # in rope_parameters rather than a top-level config attribute. + "rope_theta": _get_rope_theta(base_config, _get_rope_theta(draft_config, 1000000.0)), # YaRN long-context scaling is injected below (see the rope_scaling block). "rope_scaling": None, "tie_word_embeddings": False, diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 7679ff0020b..e0d63bde136 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -72,9 +72,11 @@ """ import logging +from typing import Any import torch import torch.nn.functional as F +import transformers from transformers import PreTrainedModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config from transformers.trainer_pt_utils import LabelSmoother @@ -100,6 +102,54 @@ __all__ = ["HFDFlashModel"] +_QWEN3_VL_MROPE_WORKAROUND_VERSION = "5.3.0" +_MULTIMODAL_FORWARD_KWARGS = frozenset( + { + "pixel_values", + "pixel_values_videos", + "image_grid_thw", + "video_grid_thw", + "mm_token_type_ids", + "image_sizes", + "images", + "videos", + } +) + + +def _multimodal_forward_kwargs(model_kwargs: dict) -> dict: + """Return collator fields accepted by Hugging Face multimodal forwards.""" + return { + name: value + for name, value in model_kwargs.items() + if name in _MULTIMODAL_FORWARD_KWARGS and value is not None + } + + +def _expand_qwen3_video_grid_thw(video_grid_thw: torch.Tensor) -> torch.Tensor: + """Return the per-frame video grid representation used by Qwen3-VL RoPE. + + Qwen3-VL's video processor emits one ``[T, H, W]`` row per source video, but + its rendered prompt contains a separate visual-token group for every temporal + frame. Transformers 5.3's ``get_rope_index`` consumes one grid row per + rendered group, while the vision encoder still requires the original one-row- + per-video representation. This helper is therefore used *only* for mRoPE + position construction; callers must keep the original tensor for the model + forward. + """ + if video_grid_thw.ndim != 2 or video_grid_thw.shape[-1] != 3: + raise ValueError( + "Qwen3-VL video_grid_thw must have shape [num_videos, 3], got " + f"{tuple(video_grid_thw.shape)}." + ) + if torch.any(video_grid_thw[:, 0] <= 0): + raise ValueError("Qwen3-VL video_grid_thw temporal lengths must be positive.") + + expanded_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + expanded_grid_thw[:, 0] = 1 + return expanded_grid_thw + + def _dpace_position_weights( confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None ) -> torch.Tensor: @@ -183,6 +233,125 @@ def _base_llm_config(self): or self.config ) + def _qwen3_vl_position_ids( + self, + input_ids, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + model_kwargs, + ): + """Precompute Qwen3-VL mRoPE positions for Transformers 5.3.0 batches. + + The video encoder consumes one grid row per source video, whereas mRoPE + consumes one row per rendered temporal-frame group. Calling the + top-level model with the original video grid makes the two contracts + conflict. Construct the mRoPE positions with a frame-expanded copy, + then pass the original grid to the vision encoder in ``forward``. + + Transformers 5.4.0 performs this frame expansion in ``get_rope_index`` + itself; only 5.3.0 needs the external workaround. See + https://github.com/huggingface/transformers/blob/v5.4.0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py + + Prefer ``get_rope_index`` over ``compute_3d_position_ids``. The latter + writes ``rope_deltas`` into the base model even though DFlash training + never supplies a cache. Keeping this calculation side-effect free is + important when the frozen target is reused for consecutive training + batches or validation. + """ + model_type = str(getattr(self.config, "model_type", "")) + if ( + position_ids is not None + or not model_type.startswith("qwen3_vl") + # Cached decoding uses the base model's rope_deltas path. DFlash + # training has no cache and is the only path that needs the + # frame-expanded construction below. + or past_key_values is not None + ): + return position_ids + + image_grid_thw = model_kwargs.get("image_grid_thw") + video_grid_thw = model_kwargs.get("video_grid_thw") + if not isinstance(image_grid_thw, torch.Tensor) and not isinstance( + video_grid_thw, torch.Tensor + ): + return position_ids + + if transformers.__version__ != _QWEN3_VL_MROPE_WORKAROUND_VERSION: + if transformers.__version__.startswith("5.3."): + raise RuntimeError( + "Qwen3-VL DFlash mRoPE supports Transformers 5.3.0 or >=5.4.0; " + f"got {transformers.__version__}. A 5.3.x patch release may already " + "expand video_grid_thw internally." + ) + return position_ids + + mm_token_type_ids = model_kwargs.get("mm_token_type_ids") + backbone = getattr(self, "model", None) + # Probed dynamically: which one exists depends on the Transformers version. + get_rope_index: Any = getattr(backbone, "get_rope_index", None) + compute_position_ids: Any = getattr(backbone, "compute_3d_position_ids", None) + if ( + not isinstance(mm_token_type_ids, torch.Tensor) + or input_ids is None + or (not callable(get_rope_index) and not callable(compute_position_ids)) + ): + raise ValueError( + "Qwen3-VL DFlash training requires input_ids, mm_token_type_ids, and " + "a Qwen3-VL model with get_rope_index or compute_3d_position_ids. " + "Use the Qwen3-VL AutoProcessor without dropping mm_token_type_ids." + ) + + if mm_token_type_ids.shape != input_ids.shape: + raise ValueError( + "Qwen3-VL mm_token_type_ids must have the same shape as input_ids, got " + f"{tuple(mm_token_type_ids.shape)} and {tuple(input_ids.shape)}." + ) + + rope_video_grid_thw = video_grid_thw + if isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.numel() > 0: + video_token_mask = mm_token_type_ids == 2 + if isinstance(attention_mask, torch.Tensor): + video_token_mask = video_token_mask & attention_mask.bool() + video_group_starts = video_token_mask.clone() + video_group_starts[:, 1:] &= ~video_token_mask[:, :-1] + expected_video_groups = int(video_grid_thw[:, 0].sum()) + actual_video_groups = int(video_group_starts.sum()) + if actual_video_groups != expected_video_groups: + raise ValueError( + "Qwen3-VL video frame groups do not match video_grid_thw: " + f"expected {expected_video_groups}, found {actual_video_groups}." + ) + rope_video_grid_thw = _expand_qwen3_video_grid_thw(video_grid_thw) + + rope_kwargs = { + "input_ids": input_ids, + "image_grid_thw": image_grid_thw, + "video_grid_thw": rope_video_grid_thw, + "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, + } + if callable(get_rope_index): + position_ids, _ = get_rope_index(**rope_kwargs) + else: + position_ids = compute_position_ids( + **rope_kwargs, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + ) + + expected_shape = (3, *input_ids.shape) + valid_position_ids = isinstance(position_ids, torch.Tensor) and ( + tuple(position_ids.shape) == expected_shape + ) + if not valid_position_ids: + raise RuntimeError( + "Qwen3-VL produced invalid mRoPE position ids: expected shape " + f"{expected_shape}, got {getattr(position_ids, 'shape', None)}." + ) + return position_ids + def _find_base_model_parts(self): """Locate base model submodules (backbone, embeddings, lm_head) by probing known paths. @@ -592,6 +761,16 @@ def forward( - Label alignment: position k predicts token at anchor+k - Optional loss decay weighting """ + if self.training: + position_ids = self._qwen3_vl_position_ids( + input_ids, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + kwargs, + ) + if not self.training: if self.dflash_offline: raise RuntimeError( @@ -638,12 +817,37 @@ def forward( ) target_hidden = base_outputs.target_hidden else: - # TODO: For co-training the base model, remove no_grad and eval() switch. + # Multimodal models need the top-level conditional-generation forward so their + # image/video features are inserted before the language model runs. Keep the + # long-standing narrow call for text-only models. + base_forward_kwargs = _multimodal_forward_kwargs(kwargs) + use_top_level_forward = bool(base_forward_kwargs) with torch.no_grad(): - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, + if use_top_level_forward: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=False, + output_attentions=output_attentions, + output_hidden_states=True, + cache_position=cache_position, + return_dict=True, + **base_forward_kwargs, + ) + else: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + + if not getattr(raw_outputs, "hidden_states", None): + raise RuntimeError( + "The base model did not return hidden states required for DFlash training. " + "Ensure its top-level multimodal forward supports output_hidden_states=True." ) offset = 1 selected = [raw_outputs.hidden_states[lid + offset] for lid in self.target_layer_ids] @@ -652,16 +856,15 @@ def forward( target_hidden=target_hidden, logits=raw_outputs.logits ) - # 2. Build loss mask. - # When labels are provided (answer_only_loss), they already encode both - # assistant masking and padding (-100 for both). When labels are not - # provided, fall back to attention_mask for padding only. + # 2. Build loss mask. Labels carry optional answer-only masking, but do + # not in general mark padded tokens with -100 (the VLM collator creates + # them from padded input_ids). Always intersect with attention_mask so + # anchor sampling and loss never include the padded tail. + loss_mask = torch.ones(bsz, seq_len, device=device) if labels is not None: - loss_mask = (labels != LabelSmoother.ignore_index).float() - elif attention_mask is not None: - loss_mask = attention_mask.float() - else: - loss_mask = torch.ones(bsz, seq_len, device=device) + loss_mask = loss_mask * (labels != LabelSmoother.ignore_index).float() + if attention_mask is not None: + loss_mask = loss_mask * attention_mask.float() # In offline training, assistant mask is dumped and passed as kwarg. if kwargs.get("loss_mask") is not None: @@ -674,8 +877,16 @@ def forward( n_blocks = anchor_positions.shape[1] if n_blocks == 0 or not block_keep_mask.any(): - # Zero loss that still flows through dflash_module for DDP gradient sync - dummy = self.dflash_module.fc.weight.sum() * 0.0 + # Keep all trainable draft parameters in the graph so DDP can reduce a rank + # that receives an all-masked answer-only batch. + dummy = sum( + ( + parameter.reshape(-1)[0] * 0.0 + for parameter in self.dflash_module.parameters() + if parameter.requires_grad + ), + torch.zeros((), device=device), + ) return ModelOutput(loss=dummy, logits=base_outputs.logits, train_acc=[[0.0]]) # 4. Build draft inputs diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 6a3c19b993d..8c5418bb6b8 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -610,7 +610,13 @@ def load_vlm_or_llm( return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code) if _is_vlm: - model_cls = transformers.AutoModelForVision2Seq + # Transformers 5 renamed AutoModelForVision2Seq to + # AutoModelForImageTextToText. Prefer the pre-5 name so this loader + # continues to support the Transformers 4 environments used by older + # speculative-decoding jobs. + model_cls = getattr(transformers, "AutoModelForVision2Seq", None) + if model_cls is None: + model_cls = transformers.AutoModelForImageTextToText else: model_cls = transformers.AutoModelForCausalLM diff --git a/modelopt/torch/utils/plugins/transformers_dataset.py b/modelopt/torch/utils/plugins/transformers_dataset.py index c27a3d09aea..97ae4ea2d14 100644 --- a/modelopt/torch/utils/plugins/transformers_dataset.py +++ b/modelopt/torch/utils/plugins/transformers_dataset.py @@ -325,6 +325,7 @@ def __init__( chat_template: str | None = None, add_generation_prompt: bool = False, answer_only_loss: bool = False, + shift_labels: bool = True, local_image_path: str = "", return_labels: bool = False, ): @@ -340,10 +341,97 @@ def __init__( chat_template=chat_template, add_generation_prompt=add_generation_prompt, answer_only_loss=answer_only_loss, + shift_labels=shift_labels, return_labels=return_labels, ) + def _verify_generation_tags(self): + """Accept VLM templates whose assistant spans have stable chat markers. + + Cosmos/Qwen ChatML templates do not necessarily use Hugging Face's + ``{% generation %}`` tags. For those templates we derive the same + assistant-only loss mask from the tokenized assistant boundaries. + """ + if self._assistant_marker_specs(): + return + super()._verify_generation_tags() + + def _assistant_marker_specs(self): + """Return tokenized assistant start/end boundaries for supported templates.""" + if hasattr(self, "_cached_assistant_marker_specs"): + return self._cached_assistant_marker_specs + + template = self.tokenizer.chat_template or "" + specs = [] + if "<|im_start|>" in template and "<|im_end|>" in template: + specs.append( + ( + self.tokenizer("<|im_start|>assistant\n", add_special_tokens=False)[ + "input_ids" + ], + [ + self.tokenizer("<|im_end|>\n", add_special_tokens=False)["input_ids"], + self.tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"], + ], + ) + ) + self._cached_assistant_marker_specs = [ + (start, [end for end in ends if end]) for start, ends in specs if start and any(ends) + ] + return self._cached_assistant_marker_specs + + @staticmethod + def _find_subsequence(values, pattern, start=0, stop=None): + stop = len(values) if stop is None else stop + if not pattern or start >= stop: + return -1 + for index in range(start, stop - len(pattern) + 1): + if values[index : index + len(pattern)] == pattern: + return index + return -1 + + def _build_assistant_masks(self, tokenized_messages): + """Build assistant-content masks from ChatML boundaries.""" + input_ids = tokenized_messages["input_ids"] + attention_mask = tokenized_messages.get("attention_mask") + assistant_masks = torch.zeros_like(input_ids) + + for row_index, row in enumerate(input_ids): + tokens = row.tolist() + if isinstance(attention_mask, torch.Tensor): + active = attention_mask[row_index].nonzero(as_tuple=False).flatten() + if active.numel() == 0: + continue + sequence_start, sequence_end = int(active[0]), int(active[-1]) + 1 + else: + sequence_start, sequence_end = 0, len(tokens) + + for start_marker, end_markers in self._assistant_marker_specs(): + search_from = sequence_start + while search_from < sequence_end: + start = self._find_subsequence(tokens, start_marker, search_from, sequence_end) + if start == -1: + break + content_start = start + len(start_marker) + end_positions = [ + position + for marker in end_markers + if ( + position := self._find_subsequence( + tokens, marker, content_start, sequence_end + ) + ) + != -1 + ] + content_end = min(end_positions) if end_positions else sequence_end + if content_start < content_end: + assistant_masks[row_index, content_start:content_end] = 1 + search_from = max(content_start + 1, content_end + 1) + + return assistant_masks + def _process_multimodal_sample(self, examples): + derive_masks_from_markers = self.answer_only_loss and bool(self._assistant_marker_specs()) tokenized_messages = self.processor.apply_chat_template( examples, tokenize=True, @@ -353,9 +441,36 @@ def _process_multimodal_sample(self, examples): truncation=True, max_length=self.train_len, add_generation_prompt=self.add_generation_prompt, - return_assistant_tokens_mask=self.answer_only_loss, + return_assistant_tokens_mask=self.answer_only_loss and not derive_masks_from_markers, ) + if derive_masks_from_markers: + tokenized_messages["assistant_masks"] = self._build_assistant_masks(tokenized_messages) + + if self.return_labels: + input_ids = tokenized_messages["input_ids"] + labels = input_ids.new_full(input_ids.shape, IGNORE_TOKEN_ID) + if self.shift_labels: + labels[..., :-1] = input_ids[..., 1:] + else: + # DFlash predicts the token at the current position rather + # than the next autoregressive token. + labels[:] = input_ids + + if self.answer_only_loss: + if "assistant_masks" not in tokenized_messages: + raise ValueError( + "answer_only_loss requires assistant_masks from the VLM chat template." + ) + assistant_mask = tokenized_messages["assistant_masks"] + if not isinstance(assistant_mask, torch.Tensor) or not assistant_mask.any(): + labels[:] = IGNORE_TOKEN_ID + elif self.shift_labels: + labels[..., :-1][assistant_mask[..., 1:] == 0] = IGNORE_TOKEN_ID + else: + labels[assistant_mask == 0] = IGNORE_TOKEN_ID + tokenized_messages["labels"] = labels + return tokenized_messages def __call__(self, examples): @@ -385,6 +500,21 @@ def __call__(self, examples): msg["content"] = [{"type": "text", "text": msg["content"]}] for ctn in msg["content"]: + # Some JSONL producers use a fixed multimodal-part schema + # (text/image/video/fps on every part) so Arrow can load + # heterogeneous image and video datasets together. Drop + # the inactive placeholders before handing a part to the + # processor, which expects only fields relevant to its type. + content_type = ctn.get("type") + if content_type != "text" and ctn.get("text") == "": + del ctn["text"] + if content_type != "image" and ctn.get("image") == "": + del ctn["image"] + if content_type != "video": + if ctn.get("video") == "": + del ctn["video"] + if ctn.get("fps") == 0: + del ctn["fps"] if ctn["type"] == "image" and "image" in ctn: ctn["image"] = os.path.abspath( os.path.join(self.local_image_path, ctn["image"]) diff --git a/tests/unit/torch/export/test_hf_spec_rope_export.py b/tests/unit/torch/export/test_hf_spec_rope_export.py index 720082bc617..fbeb218793e 100644 --- a/tests/unit/torch/export/test_hf_spec_rope_export.py +++ b/tests/unit/torch/export/test_hf_spec_rope_export.py @@ -139,3 +139,16 @@ def test_dflash_rope_theta_inherits_base(): """rope_theta is inherited from the target/base config (draft drafts for the base).""" config = _make_dflash_exporter(base_rope_theta=5000000.0)._export_config() assert config["rope_theta"] == 5000000.0 + + +def test_dflash_rope_theta_inherits_base_rope_parameters(): + """Transformers 5 stores the target RoPE base in rope_parameters.""" + exporter = _make_dflash_exporter(base_rope_theta=None) + exporter.model.config.rope_parameters = { + "rope_type": "default", + "rope_theta": 5000000.0, + } + + config = exporter._export_config() + + assert config["rope_theta"] == 5000000.0 diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index 2880bf4ef1c..cf6dfe1a6bc 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -134,3 +134,31 @@ def _fake_from_pretrained(*args, **kwargs): model = load_vlm_or_llm("fake-model", use_offline_training=True, use_fake_base=False) assert captured_kwargs.get("num_hidden_layers") == 0 assert model.config.num_orig_hidden_layers == 4 + + +def test_load_vlm_or_llm_uses_transformers5_vlm_auto_class(monkeypatch): + """Transformers 5 loads VLMs through AutoModelForImageTextToText.""" + cfg = transformers.PretrainedConfig() + cfg.model_type = "qwen3_vl" + cfg.text_config = object() + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **kw: cfg) + + captured = {} + + class _FakeVLM: + @staticmethod + def from_pretrained(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return object() + + # ``transformers`` exposes auto classes lazily, so deleting this attribute + # lets its module-level ``__getattr__`` recreate the legacy class. An + # explicit ``None`` models its absence and reliably exercises the v5 + # fallback. + monkeypatch.setattr(transformers, "AutoModelForVision2Seq", None, raising=False) + monkeypatch.setattr(transformers, "AutoModelForImageTextToText", _FakeVLM) + + assert load_vlm_or_llm("qwen3-vl", dtype="auto") is not None + assert captured["args"] == ("qwen3-vl",) + assert captured["kwargs"]["torch_dtype"] == "auto" diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ef2eec2ad07..bd243421d2c 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -35,6 +35,7 @@ import modelopt.torch.opt as mto import modelopt.torch.speculative as mtsp +import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import ( DFlashAttention, @@ -119,6 +120,222 @@ def test_convert_sets_mask_token_id(self): assert model.mask_token_id == 0 +def test_qwen3_vl_transformers_530_position_ids_expand_video_grid(monkeypatch): + """Only mRoPE receives a per-frame video grid on Transformers 5.3.0.""" + original_grid = torch.tensor([[3, 4, 5], [2, 6, 7]]) + expected_position_ids = torch.ones(3, 1, 12, dtype=torch.long) + get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) + compute_position_ids = MagicMock() + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace( + get_rope_index=get_rope_index, + compute_3d_position_ids=compute_position_ids, + ), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": original_grid, + "mm_token_type_ids": torch.tensor([[2, 0, 2, 0, 2, 0, 2, 0, 2, 0, 0, 0]]), + }, + ) + + assert position_ids is expected_position_ids + assert not compute_position_ids.called + assert torch.equal(original_grid, torch.tensor([[3, 4, 5], [2, 6, 7]])) + assert torch.equal( + get_rope_index.call_args.kwargs["video_grid_thw"], + torch.tensor([[1, 4, 5], [1, 4, 5], [1, 4, 5], [1, 6, 7], [1, 6, 7]]), + ) + + +def test_qwen3_vl_moe_transformers_530_position_ids_expand_video_grid(monkeypatch): + """Qwen3-VL family variants use the same 5.3.0 mRoPE workaround.""" + expected_position_ids = torch.ones(3, 1, 4, dtype=torch.long) + get_rope_index = MagicMock(return_value=(expected_position_ids, torch.zeros(1, 1))) + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl_moe"), + model=SimpleNamespace(get_rope_index=get_rope_index), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[2, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 0, 2, 0]]), + }, + ) + + assert position_ids is expected_position_ids + assert torch.equal( + get_rope_index.call_args.kwargs["video_grid_thw"], + torch.tensor([[1, 4, 4], [1, 4, 4]]), + ) + + +def test_qwen3_vl_transformers_53_patch_release_raises(monkeypatch): + """Avoid double expansion when a 5.3 patch backports the upstream fix.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.1") + + with pytest.raises(RuntimeError, match=r"5\.3\.0 or >=5\.4\.0"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]), + }, + ) + + +def test_qwen3_vl_transformers_54_uses_native_position_ids(monkeypatch): + """Transformers 5.4+ performs the grid expansion inside get_rope_index.""" + get_rope_index = MagicMock() + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=get_rope_index), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.4.0") + + position_ids = HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 0, 0, 0]]), + }, + ) + + assert position_ids is None + assert not get_rope_index.called + + +def test_qwen3_vl_transformers_530_rejects_bad_video_frame_groups(monkeypatch): + """Fail before mRoPE construction when processor and video-grid contracts differ.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="video frame groups"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 4, dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "video_grid_thw": torch.tensor([[2, 4, 4]]), + "mm_token_type_ids": torch.tensor([[2, 2, 0, 0]]), + }, + ) + + +def test_multimodal_forward_kwargs_exclude_non_model_inputs(): + """Do not forward Trainer or collator-only fields to Hugging Face models.""" + pixel_values = torch.ones(1) + mm_token_type_ids = torch.zeros(1, 4, dtype=torch.long) + + forwarded = hf_dflash._multimodal_forward_kwargs( + { + "pixel_values": pixel_values, + "mm_token_type_ids": mm_token_type_ids, + "assistant_masks": torch.ones(1, 4), + "loss_mask": torch.ones(1, 4), + "num_items_in_batch": 4, + "unexpected_dataset_column": "drop me", + } + ) + + assert set(forwarded) == {"pixel_values", "mm_token_type_ids"} + assert forwarded["pixel_values"] is pixel_values + assert forwarded["mm_token_type_ids"] is mm_token_type_ids + + +def test_eval_does_not_precompute_qwen3_vl_position_ids(monkeypatch): + """Evaluation delegates mRoPE construction to the base model and its cache.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_dflash_config())]) + precompute_position_ids = MagicMock() + monkeypatch.setattr(model, "_qwen3_vl_position_ids", precompute_position_ids) + + model.eval() + model(input_ids=torch.tensor([[1, 2, 3, 4]])) + + precompute_position_ids.assert_not_called() + + +def test_qwen3_vl_transformers_53_position_ids_require_mm_token_types(monkeypatch): + """Never silently fall back to one-dimensional positions for a visual batch.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="mm_token_type_ids"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={"image_grid_thw": torch.tensor([[1, 4, 4]])}, + ) + + +def test_qwen3_vl_transformers_53_position_ids_reject_bad_mm_token_shape(monkeypatch): + """Keep processor-produced modality ids aligned with the padded text sequence.""" + fake_model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_vl"), + model=SimpleNamespace(get_rope_index=MagicMock()), + ) + monkeypatch.setattr(hf_dflash.transformers, "__version__", "5.3.0") + + with pytest.raises(ValueError, match="same shape as input_ids"): + HFDFlashModel._qwen3_vl_position_ids( + fake_model, + input_ids=torch.ones(1, 12, dtype=torch.long), + attention_mask=torch.ones(1, 12, dtype=torch.long), + position_ids=None, + past_key_values=None, + inputs_embeds=None, + model_kwargs={ + "image_grid_thw": torch.tensor([[1, 4, 4]]), + "mm_token_type_ids": torch.zeros(1, 11, dtype=torch.long), + }, + ) + + class TestDPaceWeights: """Test the D-PACE position-weighting objective (arXiv:2605.18810).""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py index 2deefadd9e3..5abaa124d68 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py +++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py @@ -30,7 +30,7 @@ import pytest import torch -from _test_utils.torch.transformers_models import get_tiny_llama +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.eagle.default_config import default_eagle_config @@ -38,6 +38,7 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) +from modelopt.torch.utils.plugins import transformers_dataset _mock_scripts = types.ModuleType("scripts") _mock_ar = types.ModuleType("scripts.ar_validate") @@ -57,6 +58,58 @@ make_speculative_data_module = _eagle_utils.make_speculative_data_module +# --------------------------------------------------------------------------- +# online VLM data-module wiring +# --------------------------------------------------------------------------- + + +def test_vlm_data_module_passes_dflash_label_mode(monkeypatch): + """VLM batches must use unshifted labels for DFlash and preserve OSL settings.""" + data_args = argparse.Namespace( + mode="online", + data_path="unused.jsonl", + vlm_processor="dummy-vlm-processor", + vlm_img_dir="/images", + chat_template=None, + ) + collator = MagicMock() + monkeypatch.setattr(_eagle_utils, "ShardedDataset", MagicMock()) + monkeypatch.setattr(_eagle_utils, "VisionLanguageDataCollator", collator) + + module = make_speculative_data_module( + MagicMock(), data_args, train_len=16, answer_only_loss=True, shift_labels=False + ) + + collator.assert_called_once_with( + processor="dummy-vlm-processor", + train_len=16, + local_image_path="/images", + return_labels=True, + answer_only_loss=True, + shift_labels=False, + chat_template=None, + ) + assert module["data_collator"] is collator.return_value + + +def test_vlm_data_collator_accepts_unshifted_labels(monkeypatch): + """The real VLM collator must support DFlash's unshifted labels.""" + processor = types.SimpleNamespace(tokenizer=get_tiny_tokenizer()) + monkeypatch.setattr( + transformers_dataset.transformers.AutoProcessor, + "from_pretrained", + lambda *_args, **_kwargs: processor, + ) + + collator = transformers_dataset.VisionLanguageDataCollator( + processor="dummy-vlm-processor", + chat_template="{{ messages }}", + shift_labels=False, + ) + + assert collator.shift_labels is False + + # --------------------------------------------------------------------------- # sample_size truncation tests # --------------------------------------------------------------------------- From ae88d67d0bed57859575ee6c7a0df4a467827bca Mon Sep 17 00:00:00 2001 From: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Date: Thu, 6 Aug 2026 11:26:06 +0800 Subject: [PATCH 02/11] fix(export): honor sub-model scope_prefix in quant-aware reverse rename (NVBug 6525511) (#2076) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What does this PR do? **Type of change:** Bug fix Fixes [NVBug 6525511](https://nvbugspro.nvidia.com/bug/6525511) / [OMNIML-5599](https://jirasw.nvidia.com/browse/OMNIML-5599): FP8 PTQ of `llava-1.5-13b` on `transformers>=5.12` produces a checkpoint vLLM refuses to load: ``` ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration ``` A cross-architecture sweep run for this PR shows **`google/gemma-3-4b-it` is broken the same way** (884/884 keys mangled) and is fixed by the same change, though no bug was filed for it. ### Root cause transformers collects conversion mappings **recursively** and tags each sub-model's transforms with the sub-module path they belong to, then matches only keys under that prefix (`WeightTransform._scoped_match`: strip `scope_prefix` → match → re-attach): ```python transform.scope_prefix = scope_prefix transform.base_model_prefix = model.base_model_prefix ``` `LlavaForConditionalGeneration` therefore carries the **vision tower's own** `PrefixChange` — "add a `vision_model.` prefix" — scoped to `model.vision_tower`: ``` [5] rev=PrefixChange scope_prefix = 'model.vision_tower' source_patterns = ['^(?:(?!vision_model\.))(.+)$'] target_patterns = ['vision_model.\1'] ``` `_build_reverse_rules` read `rev.source_patterns` / `rev.target_patterns` raw and **discarded `scope_prefix`**. Read as an unscoped regex, that pattern means "any key not already starting with `vision_model.`" — i.e. everything. All 758 llava-1.5-13b tensors were moved under a bogus top-level `vision_model.`: ``` model.language_model.layers.0.self_attn.q_proj.weight -> vision_model.model.language_model.layers.0.self_attn.q_proj.weight lm_head.weight -> vision_model.lm_head.weight ``` This is the same class of defect as #2032 (NVBug 6525534), but that fix's shadowing heuristic cannot catch it: it drops a rule whose target is an **existing** namespace, whereas this rule **invents** one (`vision_model` exists nowhere in the module tree). ### The fix `RenameRule` carries `scope_prefixes`; `_sub_scoped` applies a scoped rule only to keys under one of them — stripping the prefix before the match and re-attaching after, mirroring transformers' own semantics (trying `base_model_prefix + scope_prefix` before `scope_prefix`). The same scoping flows through `build_reverse_name_mapper`, so `exclude_modules` (which lists the BF16 vision tower) stays aligned with the weights. `_drop_shadowed_prefix_renames` skips scoped rules, which are already confined to their subtree. Converter-derived rules (`_expert_leaf_renames`, `_dense_split_rule`) match by **module suffix** rather than an anchored pattern, so they cannot be confined to a subtree the same way. A survey of 9 architectures (LLaVA, LLaVA-Next, Gemma-3, Gemma-4, Qwen2-VL, Qwen3-VL-MoE, Llama-4-Scout, Mixtral, DeepSeek-V2-Lite) found **every `WeightConverter` has `scope_prefix=None`** — transformers only scopes `WeightRenaming`/`PrefixChange` — so the case is unreachable today. Rather than emit rules that could silently reach a sibling namespace if that ever changes, a scoped converter now raises `QuantConversionUnsupportedError` and the caller falls back to in-memory names with a warning. ## Testing **Name-level oracle against the real hub checkpoint.** Exported names must equal the original checkpoint's keys. All 758 `llava-hf/llava-1.5-13b-hf` state-dict keys round-trip exactly: ``` MISSING (hub key not produced): 0 SPURIOUS (name not in hub) : 0 exported: {language_model: 363, vision_tower: 391, multi_modal_projector: 4} hub : {language_model: 363, vision_tower: 391, multi_modal_projector: 4} ``` Pre-fix, 758/758 keys were mangled. This check involves no vLLM. **Cross-architecture regression sweep** — old (unscoped) vs new (scoped) mapping over each model's real state dict: | Result | Models | |---|---| | Identical (no behavior change) | Llama-3.2, Qwen2.5, Qwen3, Mistral, Phi-3, SmolLM2, gpt-oss-20b, DeepSeek-V2-Lite, Mixtral-8x7B, gemma-2-2b, gemma-4-31B, Qwen2-VL-2B | | Differs (fixed) | llava-1.5-7b (686/686), gemma-3-4b-it (884/884) | Note Qwen2-VL carries a scoped rule yet is unchanged — the fix only bites where a scoped rule would have wrongly matched. **End-to-end**, on the exact image from the bug report (`vllm/vllm-openai:v0.26.0`, verified `vllm.__version__ == 0.26.0`, transformers 5.14.1): real FP8 PTQ (`general/ptq/fp8_default-kv_fp8_cast`, `--calib_size 512`) on llava-1.5-13b, then served with the bug's exact `api_server` command. ``` PTQ_EXIT=0 top-level namespaces: {'language_model': 923, 'multi_modal_projector': 4, 'vision_tower': 391} KEYS UNDER BOGUS vision_model.* : 0 occurrences of "no module or parameter named 'vision_model'": 0 Loading weights took 5.58 seconds Model loading took 13.2 GiB and 7.45 seconds ``` The reported failure is gone and weight loading completes. - `pytest tests/unit/torch/export` → **116 passed, 1 skipped** (3 new regression tests) - `pre-commit run --files ...` → all hooks pass ## Additional Information **Out of scope, for whoever picks up the QA ticket.** After this fix the bug's exact repro hits a *different* error: llava-1.5-13b ships `"dtype": "float16"` and vLLM's FP8 kernel requires BF16 output (`RuntimeError: For FP8 input, output must have dtype BF16`). This is not an export defect — it also occurs with `kv_cache_dtype=auto`, i.e. with the FP8 KV cache entirely out of the picture, and the same checkpoint loads and generates correctly under `--dtype bfloat16` (`" Paris. with a population of about 2,249,03"`). QA will need `--dtype bfloat16`. **[NVBug 6525597](https://nvbugspro.nvidia.com/bug/6525597)** (`gemma-4-31B-it`, `assert layer.k_scale > 0.0`) is unrelated to this PR: Gemma4 exposes zero transformers conversions, so this change provably does not touch it (confirmed identical in the sweep above). Re-tested separately on vLLM 0.26.0 with that bug's environment (transformers 5.5.0, TP=1, batch_size 8) it did not reproduce — the server reached `Application startup complete` with zero asserts — but that is tracked outside this PR. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit - **Bug Fixes** - Improved Hugging Face exports for multimodal models with nested prefix conversions. - Ensured vision-tower prefixes apply only within the correct model scope. - Prevented unrelated language-model and head parameters from being incorrectly renamed. - Preserved scoped behavior for module mappings and wildcard exclusions. - Added safer handling for unsupported scoped weight conversions. - **Tests** - Added regression coverage for scoped prefix handling across weights, module names, exclusions, and unsupported conversion scenarios. --------- Signed-off-by: James Shen Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../torch/export/quant_aware_conversion.py | 106 +++++++++++-- .../export/test_quant_aware_conversion.py | 145 ++++++++++++++++++ 3 files changed, 242 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b51831f402f..1bc8ebad9fe 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -83,6 +83,7 @@ Changelog - Fix ``examples/vllm_serve`` serving shared experts uncalibrated: their ``gate_proj``/``up_proj`` quantizer keys were not merged into ``gate_up_proj`` on reload, so they matched no module and were dropped. - Fix Qwen3-VL MoE PTQ failing on ``transformers>=5.12`` with ``AttributeError: 'QuantQwen3VLMoeTextExperts' object has no attribute 'hidden_size'`` (NVBug 6518551). transformers 5.12 moved ``Qwen3VLMoeTextExperts`` onto the standard ``@use_experts_implementation`` fused layout (``hidden_size``/``expert_dim`` renamed to ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` transposed to ``(num_experts, 2*intermediate_dim, hidden_dim)``, two ``F.linear`` calls per expert), but the legacy ``_QuantQwen3VLMoeTextExperts`` wrapper stayed statically registered and shadowed on-the-fly detection. The new layout is now left to ``register_fused_experts_on_the_fly``, which claims it with the generic ``_QuantFusedExperts``; the legacy wrapper is still registered on ``transformers<5.12``, whose ``torch.bmm``-based forward the generic wrapper cannot intercept. - Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors. +- Fix unified HF export of multimodal models whose vision tower carries its own ``PrefixChange`` conversion (``LlavaForConditionalGeneration`` on ``transformers>=5.12`` — NVBug 6525511). transformers collects conversion mappings recursively and scopes each sub-model's transforms to that sub-module via ``scope_prefix``, matching only keys under that prefix. ModelOpt's quant-aware reverse conversion read the raw patterns and ignored ``scope_prefix``, so the vision tower's "add a ``vision_model.`` prefix" rule was applied to *every* key in the state dict: an exported llava-1.5-13b checkpoint had all 758 tensors moved under a bogus top-level ``vision_model.`` namespace (``vision_model.language_model.*``, ``vision_model.lm_head.*``), and vLLM rejected it with ``ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration``. Reverse rename rules now carry their scope and are applied only to keys under it, matching transformers' own ``WeightTransform._scoped_match`` semantics. ``Gemma3ForConditionalGeneration`` was affected identically and is fixed by the same change. 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index ece6e335ca8..2e32f123869 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -94,10 +94,18 @@ class QuantConversionUnsupportedError(Exception): @dataclass(frozen=True) class RenameRule: - """Reverse of a ``WeightRenaming``: ``re.sub(pattern, repl, key)`` on every key.""" + """Reverse of a ``WeightRenaming``: ``re.sub(pattern, repl, key)`` on every key. + + ``scope_prefixes`` mirrors transformers' ``WeightTransform._scoped_match``: a rule + collected from a *sub-model* carries the sub-module path it was scoped to, and its + patterns are written relative to that sub-model's own root. Such a rule must only be + applied to keys under one of these prefixes, with the prefix stripped before the + match and re-attached after. Empty tuple means the rule is unscoped (whole-model). + """ pattern: str repl: str + scope_prefixes: tuple[str, ...] = () @dataclass(frozen=True) @@ -157,6 +165,35 @@ def _apply_split_rule(state_dict: dict[str, torch.Tensor], rule: SplitRule) -> N state_dict[target_key] = _split_leaf_tensor(leaf, tensor, n, idx, rule.dim) +def _compile_rename_rules(rename_rules: list[RenameRule]): + """Pre-compile rename rules into ``(compiled_pattern, repl, scope_prefixes)`` triples.""" + return [(re.compile(r.pattern), r.repl, r.scope_prefixes) for r in rename_rules] + + +def _sub_scoped(pattern: re.Pattern, repl: str, key: str, scope_prefixes: tuple[str, ...]) -> str: + """Apply one rename rule to ``key``, honoring the rule's sub-model scope. + + Mirrors transformers' ``WeightTransform._scoped_match``: for a scoped rule, the first + matching prefix is stripped, the pattern is applied to the remaining suffix, and the + prefix is re-attached. A scoped rule that matches no prefix never applies -- this is + what keeps a sub-model's rule (e.g. the vision tower's "add ``vision_model.``" prefix + change) from rewriting sibling namespaces of the parent multimodal model. + """ + if not scope_prefixes: + return pattern.sub(repl, key) + for prefix in scope_prefixes: + if key.startswith(prefix): + return prefix + pattern.sub(repl, key[len(prefix) :]) + return key + + +def _apply_rename_rules(key: str, compiled) -> str: + """Apply all compiled rename rules to ``key``, in order.""" + for pattern, repl, scope_prefixes in compiled: + key = _sub_scoped(pattern, repl, key, scope_prefixes) + return key + + def apply_reverse_rules( state_dict: dict[str, torch.Tensor], split_rules: list[SplitRule], @@ -171,12 +208,10 @@ def apply_reverse_rules( for rule in split_rules: _apply_split_rule(out, rule) - compiled = [(re.compile(r.pattern), r.repl) for r in rename_rules] + compiled = _compile_rename_rules(rename_rules) renamed: dict[str, torch.Tensor] = {} for key, value in out.items(): - new_key = key - for pattern, repl in compiled: - new_key = pattern.sub(repl, new_key) + new_key = _apply_rename_rules(key, compiled) if new_key in renamed: raise QuantConversionUnsupportedError(f"rename collision on '{new_key}'") renamed[new_key] = value @@ -218,7 +253,7 @@ def build_reverse_name_mapper(model): _, rename_rules, _ = _build_reverse_rules(model) if not rename_rules: return None - compiled = [(re.compile(r.pattern), r.repl) for r in rename_rules] + compiled = _compile_rename_rules(rename_rules) # The rename patterns are anchored on full weight keys and use ``.`` (any char) as a # path separator, so a trailing glob wildcard in an exclude pattern would be consumed # (e.g. ``...mlp.shared_experts.`` -> ``...`` would eat the ``*``). Append a sentinel @@ -227,9 +262,7 @@ def build_reverse_name_mapper(model): _sentinel = ".\x00modelopt_name_sentinel" def _apply(text: str) -> str: - for pattern, repl in compiled: - text = pattern.sub(repl, text) - return text + return _apply_rename_rules(text, compiled) def _map(name: str) -> str: base, suffix = name, "" @@ -287,6 +320,32 @@ def _assert_experts_pre_expanded( ) +def _scope_prefixes(rev) -> tuple[str, ...]: + """Candidate key prefixes a scoped sub-model transform may apply under. + + transformers tags a conversion collected from a sub-model with ``scope_prefix`` (the + sub-module path) and ``base_model_prefix``, then matches keys against + ``base_model_prefix.scope_prefix.`` first and ``scope_prefix.`` second (see + ``WeightTransform._scoped_match``). Returned in that same priority order, each with a + trailing dot. Empty tuple when the transform is unscoped (owned by the root model), + in which case its patterns already address the full key space. + """ + scope = getattr(rev, "scope_prefix", None) + if scope is None: + return () + scope_dot = f"{scope}." if scope != "" else "" + base = getattr(rev, "base_model_prefix", None) or "" + base_dot = f"{base}." if base != "" else "" + # Deduplicate while preserving order. An empty candidate is kept: it only arises for + # ``scope_prefix == ""`` and, matching transformers, acts as the always-matching + # fallback that applies the pattern to the whole key. + seen: list[str] = [] + for c in (base_dot + scope_dot, scope_dot): + if c not in seen: + seen.append(c) + return tuple(seen) + + def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[RenameRule]: """Drop child-model reverse renames when the child namespace already exists. @@ -294,6 +353,13 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename ``model.language_model.*`` reverse can also reach its parent VLM. In that model, ``model.language_model`` is already registered and applying the rule globally would capture both that namespace and siblings such as ``model.visual``. + + Only rules that are genuinely confined are skipped. A rule is confined when every + one of its ``scope_prefixes`` is non-empty, because :func:`_sub_scoped` then applies + it solely under that prefix and it cannot reach a sibling namespace. An *empty* + candidate (which ``scope_prefix == ""`` produces) matches every key, so such a rule + has the reach of an unscoped one and must still face this heuristic -- otherwise a + root-scoped nested-text rename would bypass the guard added for NVBug 6525534. """ named_modules = getattr(model, "named_modules", None) if not callable(named_modules): @@ -303,6 +369,11 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename probe_suffix = ".\x00modelopt_namespace_probe" kept: list[RenameRule] = [] for rule in rules: + # `all(...)` matters: an empty candidate matches every key, so a rule carrying one + # is not actually confined and still needs the check below. + if rule.scope_prefixes and all(rule.scope_prefixes): + kept.append(rule) + continue pattern = re.compile(rule.pattern) shadowed = False for module_name in module_names: @@ -375,9 +446,24 @@ def _build_reverse_rules(model) -> tuple[list[SplitRule], list[RenameRule], list for conv in conversions: rev = conv.reverse_transform() # hub<-in-memory; reversed name patterns + ops if isinstance(rev, WeightRenaming): + scope_prefixes = _scope_prefixes(rev) for pattern, repl in zip(_as_list(rev.source_patterns), _as_list(rev.target_patterns)): - weight_renamings.append(RenameRule(pattern=pattern, repl=repl)) + weight_renamings.append( + RenameRule(pattern=pattern, repl=repl, scope_prefixes=scope_prefixes) + ) elif isinstance(rev, WeightConverter): + # Converter-derived rules (expert leaf renames, dense split) are matched by + # module suffix, not by an anchored pattern, so they carry no scope and would + # reach identically-named modules in sibling namespaces. No current model + # scopes a WeightConverter -- transformers only scopes WeightRenaming / + # PrefixChange -- so rather than emit rules we cannot scope, refuse the + # conversion and let the caller fall back to in-memory names. That is a + # warning plus unchanged names, instead of a silently mis-named checkpoint. + if _scope_prefixes(rev): + raise QuantConversionUnsupportedError( + f"scoped WeightConverter (scope_prefix=" + f"{getattr(rev, 'scope_prefix', None)!r}) cannot be reversed scope-aware" + ) ops = list(rev.operations) if any(isinstance(op, SplitModulelist) for op in ops): # Expert converter: ModelOpt already un-stacked/un-fused experts to diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index 98bf6c250c2..7cbe311370c 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -316,6 +316,151 @@ def test_nested_text_prefix_reverse_still_applies_to_text_model(): assert mapper("model.layers.0") == "model.language_model.layers.0" +def test_scoped_submodel_prefix_change_does_not_capture_siblings(): + """A vision sub-model's ``PrefixChange`` must not prefix the whole VLM state dict. + + NVBug 6525511: ``LlavaForConditionalGeneration`` on transformers>=5.12 collects the + vision tower's own "add ``vision_model.``" prefix change. transformers scopes it to + ``model.vision_tower`` via ``scope_prefix`` and only matches keys under that prefix; + applying the raw pattern instead prefixes *every* key, so the export writes + ``vision_model.language_model.*`` / ``vision_model.lm_head.*`` and vLLM fails with + "There is no module or parameter named 'vision_model'". + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import PrefixChange + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.vision_tower = torch.nn.Module() + model.model.vision_tower.encoder = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + model.lm_head = torch.nn.Linear(2, 2, bias=False) + + prefix_change = PrefixChange(prefix_to_remove="vision_model") + prefix_change.scope_prefix = "model.vision_tower" + prefix_change.base_model_prefix = "model" + model._weight_conversions = [prefix_change] + + state_dict = { + "model.vision_tower.encoder.weight": torch.randn(2, 2), + "model.language_model.layers.0.weight": torch.randn(2, 2), + "lm_head.weight": torch.randn(2, 2), + } + reverted = revert_weight_conversion_quant_aware(model, state_dict) + + # Only the vision tower's own subtree gains the ``vision_model.`` segment. + assert set(reverted) == { + "model.vision_tower.vision_model.encoder.weight", + "model.language_model.layers.0.weight", + "lm_head.weight", + } + # Regression guard: nothing may be moved under a bogus top-level ``vision_model``. + assert not any(k.startswith("vision_model.") for k in reverted) + + +def test_scoped_rule_maps_config_module_names_consistently(): + """``build_reverse_name_mapper`` must apply the same scoping as the weight rename. + + Otherwise ``exclude_modules`` (which lists the BF16 vision tower) lands in a + different namespace than the weights and a deployment loader silently treats an + excluded layer as quantized. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import PrefixChange + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.vision_tower = torch.nn.Module() + model.model.vision_tower.encoder = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + + prefix_change = PrefixChange(prefix_to_remove="vision_model") + prefix_change.scope_prefix = "model.vision_tower" + prefix_change.base_model_prefix = "model" + model._weight_conversions = [prefix_change] + + mapper = build_reverse_name_mapper(model) + assert mapper is not None + assert mapper("model.vision_tower.encoder") == "model.vision_tower.vision_model.encoder" + # Sibling namespaces are untouched. + assert mapper("model.language_model.layers.0") == "model.language_model.layers.0" + # A trailing-wildcard exclude pattern tracks the same rename its weights got, so the + # excluded (BF16) vision tower still matches the exported tensor names. + assert mapper("model.vision_tower*") == "model.vision_tower.vision_model*" + + +def test_root_scoped_rule_still_faces_shadowing_guard(): + """A ``scope_prefix == ""`` rule has whole-key-space reach and must not bypass #2032. + + ``_scope_prefixes`` keeps an empty candidate for the root scope, which + ``_sub_scoped`` matches against every key -- so such a rule is as broad as an + unscoped one. Skipping the shadowing heuristic merely because ``scope_prefixes`` is a + non-empty *tuple* would reintroduce NVBug 6525534: the nested text model's + ``^model.language_model.`` reverse would be kept and rewrite ``model.visual.*``. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import WeightRenaming + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.visual = torch.nn.Module() + model.model.visual.patch_embed = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + + renaming = WeightRenaming( + source_patterns=r"^model.language_model.", + target_patterns=r"^model.(?!language_model.)", + ) + # Root scope: reaches every key, exactly like an unscoped rule. + renaming.scope_prefix = "" + renaming.base_model_prefix = "" + model._weight_conversions = [renaming] + + state_dict = { + "model.visual.patch_embed.weight": torch.randn(2, 2), + "model.language_model.layers.0.weight": torch.randn(2, 2), + } + reverted = revert_weight_conversion_quant_aware(model, state_dict) + + # The sibling vision namespace must be untouched (the #2032 guarantee). + assert set(reverted) == set(state_dict) + assert not any("language_model.visual" in k for k in reverted) + + +def test_scoped_weight_converter_is_refused(): + """A scoped ``WeightConverter`` must fall back rather than emit unscoped rules. + + Converter-derived rules (expert leaf renames, dense splits) match by module suffix, + so they cannot be confined to a sub-model subtree the way an anchored rename can. + No current architecture scopes a converter -- transformers only scopes + ``WeightRenaming``/``PrefixChange`` -- so if one ever appears, refusing the whole + conversion keeps the in-memory names (a warning) instead of silently rewriting an + identically-named module in a sibling namespace. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: optional dependency, guarded by the importorskip above. + from transformers.core_model_loading import Chunk, WeightConverter + + conv = WeightConverter( + source_patterns="mlp.gate_up_proj", + target_patterns=["mlp.gate_proj", "mlp.up_proj"], + operations=[Chunk(dim=0)], + ) + conv.scope_prefix = "model.language_model" + conv.base_model_prefix = "model" + model = types.SimpleNamespace(_weight_conversions=[conv]) + + sd = _nvfp4_linear("model.language_model.layers.0.mlp.gate_up_proj", 8, 16) + with pytest.raises(QuantConversionUnsupportedError, match="scoped WeightConverter"): + revert_weight_conversion_quant_aware(model, sd) + + def test_split_collision_raises(): """A split whose target key already exists must fail instead of overwriting.""" sd = _nvfp4_linear("m.gate_up_proj", 8, 16) From 0a34b32c2d2f30413ef9b7870480182d5d3008a5 Mon Sep 17 00:00:00 2001 From: Zhiyu Date: Fri, 7 Aug 2026 13:19:47 -0700 Subject: [PATCH 03/11] fix(export): save models with legacy list-style _tied_weights_keys (NVBug 6518665) (#2071) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes [NVBug 6518665](https://nvbugspro.nvidia.com/bug/6518665) / [OMNIML-5583](https://jirasw.nvidia.com/browse/OMNIML-5583): PTQ of `Step-3.7-Flash` (`--recipe general/ptq/nvfp4_mlp_only-kv_fp8`, transformers 5.5.4) calibrates fine and then dies while writing the checkpoint: ``` File "modelopt/torch/export/unified_export_hf.py", line 1510, in export_hf_checkpoint model.save_pretrained( File "modelopt/torch/opt/plugins/transformers.py", line 155, in _save_pretrained_with_checks File "transformers/modeling_utils.py", line 3352, in save_pretrained state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save) File "transformers/modeling_utils.py", line 338, in _get_tied_weight_keys tied_weight_keys.extend([f"{name}.{k}" if name else k for k in tied.keys()]) AttributeError: 'list' object has no attribute 'keys' ``` **Root cause is version drift in `transformers`, not in quantization.** transformers 5.0 changed `_tied_weights_keys` from `list[str]` to a `{target: source}` dict, and `save_pretrained` → `_get_tied_weight_keys` calls `.keys()` on every submodule's declaration without a type check. `stepfun-ai/Step-3.7-Flash`'s remote code still uses the 4.x list format — `_tied_weights_keys = ["lm_head.weight"]` on `Step3p7TextModel`, `Step3p7Model` and `Step3p7ForConditionalGeneration`. The load-time consumer (`get_expanded_tied_weights_keys`) returns early when `config.tie_word_embeddings` is false, which it is for this checkpoint, so the model loads and calibrates normally and only fails at the end of the run, after the expensive part. The same crash reproduces with a plain `AutoModel.from_pretrained(..., trust_remote_code=True).save_pretrained(...)`, and transformers 5.12 tolerates `None` there but still not a list, so it is not fixed upstream either. ### Fix `_save_pretrained_with_checks` — the entry point every ModelOpt HF save routes through (unified export, `ModelOptHFTrainer`, user-called `save_pretrained`) — now normalizes a list-style declaration to the equivalent dict for the duration of the save and restores the original attribute afterwards. Mapping each entry to itself preserves the legacy semantics: those list entries were exactly the dedup patterns `_get_tied_weight_keys` is expected to return. No-op on `transformers<5`, where the list format is native, and no-op for models that already declare a dict. The restore tracks whether the instance owned the attribute, so for the usual class-level declaration nothing is left shadowing it. ### Usage ```bash python hf_ptq.py --model /local/Step-3.7-Flash --recipe general/ptq/nvfp4_mlp_only-kv_fp8 \ --dataset /local/cnn_dailymail --calib_size 32 --export_path /local/Step-3.7-Flash-nvfp4 --trust_remote_code ``` ### Testing Two new tests in `tests/unit/torch/opt/plugins/test_transformers_save_load.py`: - `test_save_pretrained_with_legacy_tied_weights_keys` (parametrized over `tie_word_embeddings`) — a quantized tiny-Llama declaring list-style keys saves and round-trips; without the fix it fails with the exact reported `AttributeError` at `modeling_utils.py:338`. It asserts the saved `model.safetensors` keys directly, so no weight can be silently dropped by the save-time dedup: `model.embed_tokens.weight` is always present, and `lm_head.weight` is present iff the weights are not actually tied (when they are, transformers drops the alias and re-ties it on load — the transformers 4.x behavior). - `test_legacy_tied_weights_keys_as_dict_restores_class_attribute` — the shim leaves no instance attribute shadowing a class-level declaration (skipped on `transformers<5`). Ran locally against both transformers majors: - transformers 5.5.4 (the version in the bug report) + torch 2.11: `tests/unit/torch/opt/plugins/` 27 passed, `tests/unit/torch/export/` 124 passed. - transformers 4.57.6: `tests/unit/torch/opt/plugins/` 26 passed (the shim no-ops, new dict-restore test skips). Not yet re-run end-to-end on the real Step-3.7-Flash checkpoint (1.4 TB / 8×B200); QA can re-run the reported command against this branch. ### Known limitation The shim covers the save side only. A legacy-list model that *also* sets `tie_word_embeddings=True` still crashes on the **load** side, in `get_expanded_tied_weights_keys` during `PreTrainedModel.post_init` — before ModelOpt has a model object to patch. Such a checkpoint cannot be loaded by transformers 5 at all, with or without ModelOpt, so it is out of scope here. `Step-3.7-Flash` is unaffected (`tie_word_embeddings` is unset, i.e. false). ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information The upstream-correct fix is for StepFun to migrate their modeling code to the transformers 5 dict format; this shim unblocks every not-yet-migrated `trust_remote_code` checkpoint in the meantime. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility with Transformers 5 when exporting models with legacy tied-weight declarations. * Preserved tied-weight settings after saving, including class-level declarations. * Ensured saved models reload with equivalent state and outputs. * **Tests** * Added coverage for quantized models, tied and untied embeddings, and legacy tied-weight formats. --------- Signed-off-by: Zhiyu Cheng Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/torch/opt/plugins/transformers.py | 48 ++++++++++++-- .../plugins/test_transformers_save_load.py | 62 +++++++++++++++++++ 3 files changed, 106 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1bc8ebad9fe..7d97f71fbcd 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -83,6 +83,7 @@ Changelog - Fix ``examples/vllm_serve`` serving shared experts uncalibrated: their ``gate_proj``/``up_proj`` quantizer keys were not merged into ``gate_up_proj`` on reload, so they matched no module and were dropped. - Fix Qwen3-VL MoE PTQ failing on ``transformers>=5.12`` with ``AttributeError: 'QuantQwen3VLMoeTextExperts' object has no attribute 'hidden_size'`` (NVBug 6518551). transformers 5.12 moved ``Qwen3VLMoeTextExperts`` onto the standard ``@use_experts_implementation`` fused layout (``hidden_size``/``expert_dim`` renamed to ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` transposed to ``(num_experts, 2*intermediate_dim, hidden_dim)``, two ``F.linear`` calls per expert), but the legacy ``_QuantQwen3VLMoeTextExperts`` wrapper stayed statically registered and shadowed on-the-fly detection. The new layout is now left to ``register_fused_experts_on_the_fly``, which claims it with the generic ``_QuantFusedExperts``; the legacy wrapper is still registered on ``transformers<5.12``, whose ``torch.bmm``-based forward the generic wrapper cannot intercept. - Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors. +- Fix HF checkpoint export failing with ``AttributeError: 'list' object has no attribute 'keys'`` for models whose modeling code still declares tied weights in the ``transformers<5`` list format (NVBug 6518665, observed on ``stepfun-ai/Step-3.7-Flash``). transformers 5.0 changed ``_tied_weights_keys`` to a ``{target: source}`` dict and ``save_pretrained`` calls ``.keys()`` on every submodule's declaration without a type check, so such models — common among ``trust_remote_code`` checkpoints — load fine but die at the end of PTQ, after calibration. ModelOpt's ``save_pretrained`` patch now normalizes a list-style declaration to the equivalent dict for the duration of the save (each entry mapped to itself, which is what the legacy list meant) and restores the original attribute afterwards. - Fix unified HF export of multimodal models whose vision tower carries its own ``PrefixChange`` conversion (``LlavaForConditionalGeneration`` on ``transformers>=5.12`` — NVBug 6525511). transformers collects conversion mappings recursively and scopes each sub-model's transforms to that sub-module via ``scope_prefix``, matching only keys under that prefix. ModelOpt's quant-aware reverse conversion read the raw patterns and ignored ``scope_prefix``, so the vision tower's "add a ``vision_model.`` prefix" rule was applied to *every* key in the state dict: an exported llava-1.5-13b checkpoint had all 758 tensors moved under a bogus top-level ``vision_model.`` namespace (``vision_model.language_model.*``, ``vision_model.lm_head.*``), and vLLM rejected it with ``ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration``. Reverse rename rules now carry their scope and are applied only to keys under it, matching transformers' own ``WeightTransform._scoped_match`` semantics. ``Gemma3ForConditionalGeneration`` was affected identically and is fixed by the same change. 0.45 (2026-07-02) diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index f72715d410c..ce963c0df82 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -55,6 +55,10 @@ "ModelOptTrainerArguments", ] +# transformers 5.0 changed `_tied_weights_keys` from a list to a {target: source} dict and added +# the `load_config` parameter to `_load_state_dict_into_zero3_model`. +_TRANSFORMERS_GE_5_0 = Version(transformers.__version__) >= Version("5.0") + def is_liger_available(): try: @@ -147,20 +151,54 @@ def _new_from_config(cls, /, config, **kwargs): return model +@contextmanager +def _legacy_tied_weights_keys_as_dict(model: nn.Module): + """Temporarily normalize legacy list-style ``_tied_weights_keys`` to dict-style. + + transformers 5.0 changed ``_tied_weights_keys`` from ``list[str]`` to a + ``{target: source}`` dict, and ``save_pretrained`` calls ``.keys()`` on the attribute of + every submodule (``_get_tied_weight_keys``) without a type check. Modeling code still on + the 4.x list format - common for ``trust_remote_code`` checkpoints, e.g. + ``stepfun-ai/Step-3.7-Flash`` - therefore loads fine but dies on save with + ``AttributeError: 'list' object has no attribute 'keys'``. + + Mapping each entry to itself preserves the legacy semantics: the list entries were exactly + the dedup patterns ``_get_tied_weight_keys`` is expected to return. The original attribute + is restored on exit so the shim stays invisible to the rest of the model's lifetime. + """ + if not _TRANSFORMERS_GE_5_0: + yield + return + + patched = [] + try: + for module in model.modules(): + tied = getattr(module, "_tied_weights_keys", None) + if isinstance(tied, (list, tuple, set)): + # The attribute is usually a class attribute; remember whether this instance + # had its own so the restore does not leave a shadowing copy behind. + patched.append((module, tied, "_tied_weights_keys" in module.__dict__)) + module._tied_weights_keys = {key: key for key in tied} + yield + finally: + for module, tied, had_own_attr in patched: + if had_own_attr: + module._tied_weights_keys = tied + else: + del module._tied_weights_keys + + def _save_pretrained_with_checks(self, save_directory, *args, **kwargs): if getattr(self, "_tp_size", None) is not None and ModeloptStateManager.is_converted(self): raise NotImplementedError( "ModelOpt does not support saving tensor parallel sharded Huggingface transformer models yet. " ) - return _new_save_pretrained(self, save_directory, *args, **kwargs) + with _legacy_tied_weights_keys_as_dict(self): + return _new_save_pretrained(self, save_directory, *args, **kwargs) # [Fix for huggingface bug] deepspeed zero3 training backend only loads params into the model from # state_dict, but not buffers. So lets explicitly load the buffers into the model from state_dict. -# The `load_config` parameter was added to `_load_state_dict_into_zero3_model` in transformers 5.0. -_TRANSFORMERS_GE_5_0 = Version(transformers.__version__) >= Version("5.0") - - def _load_params_and_buffers_into_zero3_model(model_to_load, state_dict, load_config=None): buffer_names = [name for name, _ in model_to_load.named_buffers()] buffer_state_dict = {k: v for k, v in state_dict.items() if k in buffer_names} diff --git a/tests/unit/torch/opt/plugins/test_transformers_save_load.py b/tests/unit/torch/opt/plugins/test_transformers_save_load.py index fced5734e4f..e8b3ed22151 100644 --- a/tests/unit/torch/opt/plugins/test_transformers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_transformers_save_load.py @@ -18,13 +18,20 @@ import pytest import torch +import torch.nn as nn from _test_utils.torch.opt.utils import apply_mode_with_sampling from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, tf_modelopt_state_and_output_tester, ) +from safetensors.torch import load_file from transformers import AutoConfig, AutoModelForCausalLM, LlamaForCausalLM +from modelopt.torch.opt.plugins.transformers import ( + _TRANSFORMERS_GE_5_0, + _legacy_tied_weights_keys_as_dict, +) + @pytest.mark.parametrize("model_cls", [LlamaForCausalLM, AutoModelForCausalLM]) def test_causal_lm_save_restore(tmp_path, model_cls): @@ -40,6 +47,61 @@ def test_causal_lm_save_restore(tmp_path, model_cls): tf_modelopt_state_and_output_tester(model_ref, model_test) +@pytest.mark.parametrize("tie_word_embeddings", [False, True]) +def test_save_pretrained_with_legacy_tied_weights_keys(tmp_path, tie_word_embeddings): + """A model declaring 4.x list-style `_tied_weights_keys` must still save (nvbug 6518665). + + transformers>=5 expects a `{target: source}` dict there and calls `.keys()` on it while + saving, which crashes for `trust_remote_code` modeling code that has not migrated yet. + + Both tying configurations are covered because normalizing the list to `{key: key}` feeds + those keys to the save-time dedup as patterns: an untied weight must survive it, and a + genuinely tied one must still be deduped down to its canonical name. + """ + tiny_llama_dir = create_tiny_llama_dir( + tmp_path, hidden_size=128, dtype=torch.float32, tie_word_embeddings=tie_word_embeddings + ) + model = AutoModelForCausalLM.from_pretrained(tiny_llama_dir) + model = apply_mode_with_sampling(model, ["quantize"]) + + model._tied_weights_keys = ["lm_head.weight"] + model.model._tied_weights_keys = ["embed_tokens.weight"] + + save_dir = tiny_llama_dir / "legacy_tied_keys_model" + model.save_pretrained(save_dir) + + # The declarations the model owned before the save are restored verbatim. + assert model._tied_weights_keys == ["lm_head.weight"] + assert model.model._tied_weights_keys == ["embed_tokens.weight"] + + # No weight is silently dropped: `lm_head.weight` is written out unless it really does + # share storage with the embedding, in which case transformers re-ties it on load. + saved_keys = set(load_file(save_dir / "model.safetensors")) + assert "model.embed_tokens.weight" in saved_keys + assert ("lm_head.weight" in saved_keys) is not tie_word_embeddings + + model_test = AutoModelForCausalLM.from_pretrained(save_dir) + tf_modelopt_state_and_output_tester(model, model_test) + + +@pytest.mark.skipif(not _TRANSFORMERS_GE_5_0, reason="list-style keys are native to transformers 4") +def test_legacy_tied_weights_keys_as_dict_restores_class_attribute(): + """The shim must not leave an instance attribute shadowing the class declaration.""" + + class _LegacyChild(nn.Module): + # How remote-code models declare it: on the class, not the instance. + _tied_weights_keys = ["lm_head.weight"] + + parent = nn.Module() + parent.child = _LegacyChild() + + with _legacy_tied_weights_keys_as_dict(parent): + assert parent.child._tied_weights_keys == {"lm_head.weight": "lm_head.weight"} + + assert _LegacyChild._tied_weights_keys == ["lm_head.weight"] + assert "_tied_weights_keys" not in parent.child.__dict__ + + def test_causal_lm_from_config(tmp_path): """Test loading a model using from_config after applying optimizations""" tiny_llama_dir = create_tiny_llama_dir(tmp_path, hidden_size=128, dtype=torch.float32) From 453f706ac39a359217a5246a18d85daf5ac0066d Mon Sep 17 00:00:00 2001 From: kinjalpatel27 <31936134+kinjalpatel27@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:53:11 -0700 Subject: [PATCH 04/11] [6562078]: fix calibration for vLLM 0.26.0 (#2093) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes calibration failures when running fake-quantization with vLLM 0.26.0. Two root causes are addressed: 1. **`finish_requests` must be called explicitly after `add_requests`.** In vLLM 0.26.0 the scheduler calls `finish_requests` *before* `add_requests` inside `execute_model`, so request IDs are not registered yet and cleanup never runs. The calibration loop now calls `finish_requests` directly after each batch using `dataclasses.replace`. Wrapped in `try/finally` with an inner `try/except` so it always runs and never masks the original exception. A warning is emitted when `finish_requests` is absent so the regression is self-diagnosing on future vLLM API changes. 2. **`NewRequestData` gained a `prefill_token_ids` field.** vLLM 0.26.0 added this required argument; the calibration helper now passes it. Additional: - Dockerfile updated to vLLM 0.26.0 with `USER vllm` (non-root). - README updated to include vLLM 0.26.0 in tested versions. ### Usage ```bash cd examples/vllm_serve QUANT_CFG=FP8_DEFAULT_CFG QUANT_CALIB_SIZE=8 CALIB_BATCH_SIZE=1 \ python3 vllm_serve_fakequant.py Qwen/Qwen1.5-MoE-A2.7B-Chat -tp 1 \ --host 0.0.0.0 --port 8000 ``` ### Testing Tested end-to-end FQ calibration with vLLM 0.26.0 using the Docker image built from `examples/vllm_serve/Dockerfile`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A (examples change only) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information Changes are confined to `examples/vllm_serve/` and do not affect the core library. ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Added compatibility with vLLM 0.26.0 in the serving example. * Improved calibration request handling and cleanup after model execution. * Enabled the serving container to run with a non-root user. * **Bug Fixes** * Calibration cleanup failures no longer obscure the original model execution error. * Added warnings when calibration cleanup cannot be completed. * **Documentation** * Updated the serving example documentation to list vLLM 0.26.0 among tested versions. --------- Signed-off-by: Kinjal Patel Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/vllm_serve/Dockerfile | 10 ++++--- examples/vllm_serve/README.md | 2 +- examples/vllm_serve/vllm_ptq_utils.py | 39 ++++++++++++++++++++++++--- 3 files changed, 43 insertions(+), 8 deletions(-) diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 7213c6fc430..3406c62ba71 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -1,4 +1,4 @@ -FROM vllm/vllm-openai:v0.20.0 +FROM vllm/vllm-openai:v0.26.0 # Set environment variables ENV PIP_NO_CACHE_DIR=off \ @@ -25,12 +25,16 @@ RUN cd Model-Optimizer && \ # Llama4 requires this RUN pip install flash-attn==2.7.4.post1 --no-build-isolation -# Pre-compile CUDA extensions to avoid compilation time during runtime +# Pre-compile CUDA extensions into a world-accessible directory so the vllm +# user can use the cache at runtime. +ENV TORCH_EXTENSIONS_DIR=/workspace/torch_extensions RUN python3 -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()" || true -# Allow users to run without root +# Allow the non-root vllm user to access the workspace RUN chmod -R 777 /workspace +USER vllm + # Override the ENTRYPOINT from the base image to allow flexible usage ENTRYPOINT [] diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 858243686d0..75bcf37089d 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,7 +4,7 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research. -The general fakequant example is tested with vLLM 0.9.0 and 0.19.1. The compact +The general fakequant example is tested with vLLM 0.9.0, 0.19.1, and 0.26.0. The compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. ## Prepare environment diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 88b31d54a70..709d6532fb3 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -14,6 +14,7 @@ # limitations under the License. import dataclasses +import warnings from collections.abc import Callable from typing import Any @@ -66,6 +67,7 @@ def calibrate_loop(model: Any) -> None: NewRequestData, req_id=req_id, prompt_token_ids=input_ids_list, + prefill_token_ids=input_ids_list, mm_kwargs=[], mm_hashes=[], mm_positions=[], @@ -95,10 +97,39 @@ def calibrate_loop(model: Any) -> None: structured_output_request_ids={}, grammar_bitmask=None, ) - output = self.execute_model(scheduler_output) - if hasattr(self, "sample_tokens"): - if output is None: # TODO: make this default when vllm <= 0.11 is outdated - self.sample_tokens(None) + try: + output = self.execute_model(scheduler_output) + if hasattr(self, "sample_tokens"): + if output is None: # TODO: make this default when vllm <= 0.11 is outdated + self.sample_tokens(None) + finally: + # finish_requests runs before add_requests inside execute_model, so + # req IDs aren't registered yet at that point — call it directly after. + # Wrap in try/except so a cleanup error never masks the original exception. + try: + if hasattr(self.model_runner, "finish_requests"): + cleanup_output = _create_new_data_cls( + type(scheduler_output), + scheduled_new_reqs=[], + scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs, + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks, + finished_req_ids=set(num_scheduled_tokens.keys()), + free_encoder_mm_hashes=[], + kv_connector_metadata=None, + structured_output_request_ids={}, + grammar_bitmask=None, + ) + self.model_runner.finish_requests(cleanup_output) + else: + warnings.warn( + "model_runner.finish_requests not found; request state may leak during calibration." + ) + except Exception: + warnings.warn("Failed to clean up request state after calibration batch.") return calibrate_loop From 92c5ce446e1e082b6664452e250e6384c6668555 Mon Sep 17 00:00:00 2001 From: sychen52 <41452870+sychen52@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:34:18 -0700 Subject: [PATCH 05/11] [NVBUG: 6562021] Fix vLLM FlashAttention KV cache layout handling (#2084) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix The ModelOptSparseAttentionImpl.forward (FlashAttention path) assumes the paged KV cache has shape [2, num_blocks, page_size, num_kv_heads, head_dim], where dimension 0 is the K/V split. unbind(0) on this shape returns exactly 2 tensors. Newer vLLM versions changed the FlashAttention kv_cache layout to [num_blocks, 2, page_size, num_kv_heads, head_dim] (same as FlashInfer). unbind(0) now returns num_blocks tensors, causing the unpack error. Now, we detect the layout first and then unbind. ### Usage save as before. ### Testing unittest ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **Bug Fixes** * Improved FlashAttention compatibility with vLLM key/value cache layouts. * Automatically detects supported cache dimension ordering for accurate sparse-attention processing across legacy and newer layouts. * **Tests** * Expanded coverage for supported cache layouts, cache validation, and paged-cache scenarios. Signed-off-by: Shiyang Chen Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../attention_sparsity/plugins/vllm.py | 21 +++- .../test_sparse_attn_worker.py | 99 ++++++++++++++++--- .../attention_sparsity/test_vllm_plugin.py | 31 +++--- 3 files changed, 126 insertions(+), 25 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 00411eaa4ee..243db16a2bc 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -47,6 +47,19 @@ from modelopt.torch.kernels.quantization.attention.bmm2_qdq import fake_quant_v_onwrite +@functools.cache +def _flash_attention_kv_cache_layout() -> str: + """Return the installed vLLM backend's K/V packing contract.""" + cache_shape = FlashAttentionBackend.get_kv_cache_shape(3, 16, 1, 16) + if cache_shape == (2, 3, 16, 1, 16): + return "kv-first" + if cache_shape == (3, 2, 16, 1, 16): + return "blocks-first" + if cache_shape == (3, 1, 16, 32): + return "packed" + raise RuntimeError(f"Unsupported vLLM FlashAttention KV cache shape {cache_shape}") + + def _target_sparse_ratio_for_phase(target_sparse_ratio, phase: str) -> float: """Return target sparsity for a phase, defaulting old checkpoint metadata.""" if isinstance(target_sparse_ratio, float | int): @@ -514,7 +527,13 @@ def native_forward(): if resolved is None: return native_forward() - key_cache, value_cache = kv_cache.unbind(0) + cache_layout = _flash_attention_kv_cache_layout() + if cache_layout == "kv-first": + key_cache, value_cache = kv_cache.unbind(0) + elif cache_layout == "blocks-first": + key_cache, value_cache = kv_cache.unbind(1) + else: + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) is_decode_only = attn_metadata.max_query_len <= 1 common_kw = { "layer": layer, diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 05b3c6ace8d..578922db077 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -28,7 +28,7 @@ import vllm from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends import flashinfer as flashinfer_backend -from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl from vllm.v1.attention.backends.flashinfer import ( FlashInferBackend, FlashInferImpl, @@ -541,6 +541,76 @@ def _make_flash_attention_impl(*, sparse=False, quantized=False): return impl +def _flash_attention_kv_cache(num_blocks, page_size, num_kv_heads, head_size): + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "packed": + shape = [num_blocks, num_kv_heads, page_size, 2 * head_size] + else: + shape = [num_blocks, page_size, num_kv_heads, head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + return torch.zeros(shape, dtype=torch.float16) + + +@pytest.mark.parametrize( + ("layout", "backend_shape"), + [ + ("kv-first", (2, 3, 16, 1, 16)), + ("blocks-first", (3, 2, 16, 1, 16)), + ("packed", (3, 1, 16, 32)), + ], +) +def test_flash_attention_forward_follows_backend_kv_cache_layout( + monkeypatch, layout, backend_shape +): + impl = _make_flash_attention_impl(sparse=True) + if layout == "packed": + shape = [3, impl.num_kv_heads, 16, 2 * impl.head_size] + else: + shape = [3, 16, impl.num_kv_heads, impl.head_size] + shape.insert(0 if layout == "kv-first" else 1, 2) + monkeypatch.setattr( + FlashAttentionBackend, "get_kv_cache_shape", staticmethod(lambda *_args: backend_shape) + ) + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + kv_cache = torch.zeros(shape, dtype=torch.float16) + query = torch.zeros(4, impl.num_heads, impl.head_size, dtype=torch.float16) + metadata = _flash_attention_metadata(query.shape[0], 16) + captured = {} + + def fake_attention(query, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setattr(vllm_plugin, "triton_attention", fake_attention) + + try: + impl.forward( + layer=None, + query=query, + key=query, + value=query, + kv_cache=kv_cache, + attn_metadata=metadata, + output=torch.empty_like(query), + ) + finally: + vllm_plugin._flash_attention_kv_cache_layout.cache_clear() + + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + assert captured["k_cache"].shape == expected_key_cache.shape + assert captured["v_cache"].shape == expected_value_cache.shape + assert captured["k_cache"].stride() == expected_key_cache.stride() + assert captured["v_cache"].stride() == expected_value_cache.stride() + assert captured["k_cache"].data_ptr() == expected_key_cache.data_ptr() + assert captured["v_cache"].data_ptr() == expected_value_cache.data_ptr() + assert captured["page_size"] == 16 + + def _flash_attention_mixed_metadata(decode_len=1, prefill_len=17): query_lens = (decode_len, prefill_len) seq_lens = (16, 34) @@ -568,7 +638,7 @@ def test_flash_attention_mixed_batch_splits_decode_and_prefill(monkeypatch, quan prefill_tokens = 17 impl = _make_flash_attention_impl(sparse=True, quantized=quantized) query = torch.zeros(1 + prefill_tokens, 2, 64, dtype=torch.float16) - kv_cache = torch.zeros(2, 4, 16, 2, 64, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(4, 16, 2, 64) metadata = _flash_attention_mixed_metadata(decode_len=1, prefill_len=prefill_tokens) layer = SimpleNamespace( _query_quant_in_kernel=quantized, @@ -761,7 +831,7 @@ def test_forward_delegates_cascade_metadata_to_vllm(monkeypatch): """Cascade/prefix-cache metadata should use vLLM's native implementation.""" impl = _clone_sparse_impl(_make_old_impl()) q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) output = torch.empty_like(q) attn_metadata = type("AttnMetadata", (), {"use_cascade": True})() called = {} @@ -838,9 +908,7 @@ def test_forward_delegates_launches_without_effective_sparse_work( impl = _clone_sparse_impl(_make_old_impl()) impl.sparse_kw = sparse_kw q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros( - 2, 1, max_seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16 - ) + kv_cache = _flash_attention_kv_cache(1, max_seq_len, impl.num_kv_heads, impl.head_size) output = torch.empty_like(q) attn_metadata = _flash_attention_metadata(max_query_len, max_seq_len) called = {} @@ -903,7 +971,7 @@ def test_forward_resolves_calibrated_skip_softmax_threshold(monkeypatch): "target_sparse_ratio": {"prefill": 0.4, "decode": 0.6}, } q = torch.zeros(max_query_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, seq_len, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, seq_len, impl.num_kv_heads, impl.head_size) attn_metadata = _flash_attention_metadata(max_query_len, seq_len) captured = {} @@ -980,7 +1048,7 @@ def quantize_q(query): } q = torch.full((4, impl.num_heads, impl.head_size), 2.0, dtype=torch.float16) q[2:] = 10_000 - kv_cache = torch.zeros(2, 4, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(4, 16, impl.num_kv_heads, impl.head_size) metadata = SimpleNamespace( num_actual_tokens=q.shape[0], max_query_len=1, @@ -1023,8 +1091,15 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): "v_qdq_scale": 1.0, } key_cache, value_cache, block_table, seq_lens, decode_kw = calls["decode"] - assert key_cache.data_ptr() == kv_cache[0].data_ptr() - assert value_cache.data_ptr() == kv_cache[1].data_ptr() + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "packed": + expected_key_cache, expected_value_cache = kv_cache.transpose(1, 2).split( + impl.head_size, dim=-1 + ) + else: + expected_key_cache, expected_value_cache = kv_cache.unbind(0 if layout == "kv-first" else 1) + assert key_cache.data_ptr() == expected_key_cache.data_ptr() + assert value_cache.data_ptr() == expected_value_cache.data_ptr() assert block_table is metadata.block_table assert seq_lens is metadata.seq_lens assert calls["query"].shape[0] == metadata.seq_lens.shape[0] @@ -1048,7 +1123,7 @@ def test_quantized_skip_softmax_decode_stays_on_shared_kernel(monkeypatch): } impl.sparse_kw = {"skip_softmax_threshold": 0.001} q = torch.zeros(1, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) metadata = _flash_attention_metadata(1, 16) captured = {} @@ -1140,7 +1215,7 @@ def test_forward_allows_chunked_prefill_metadata(monkeypatch): q_len = 4 kv_len = 10 q = torch.zeros(q_len, impl.num_heads, impl.head_size, dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, impl.num_kv_heads, impl.head_size, dtype=torch.float16) + kv_cache = _flash_attention_kv_cache(1, 16, impl.num_kv_heads, impl.head_size) attn_metadata = _flash_attention_metadata(q_len, kv_len) captured = {} diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py index fa11b144354..d08a4072b4d 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py @@ -20,7 +20,7 @@ * ``query_start_loc`` -> ``b_start_loc`` / ``b_seq_len`` * ``seq_lens`` -> ``b_seq_len_k`` -* ``kv_cache.unbind(0)`` -> key_cache / value_cache (axis order) +* backend-declared K/V axis -> key_cache / value_cache * ``k_cache.shape[1]`` -> ``page_size`` Asserted against a contiguous reference call to the underlying Triton kernel. @@ -33,7 +33,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE -from modelopt.torch.sparsity.attention_sparsity.plugins.vllm import ModelOptSparseAttentionImpl +from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention import attention as triton_attention @@ -46,11 +46,19 @@ } +def _make_backend_paged_cache(k_cache, v_cache): + layout = vllm_plugin._flash_attention_kv_cache_layout() + if layout == "kv-first": + return torch.stack([k_cache, v_cache], dim=0) + if layout == "blocks-first": + return torch.stack([k_cache, v_cache], dim=1) + return torch.cat([k_cache, v_cache], dim=-1).transpose(1, 2) + + def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size): - """Scatter contiguous K/V into a paged KV cache stacked as [2, ...]. + """Scatter contiguous K/V into the installed vLLM paged-cache layout. - Returns a single ``kv_cache`` tensor (matching vLLM's layout that - ``ModelOptSparseAttentionImpl`` consumes via ``kv_cache.unbind(0)``). + Returns a single ``kv_cache`` tensor with the backend-declared K/V axis. """ batch = b_seq_len.shape[0] device, dtype = k.device, k.dtype @@ -76,14 +84,13 @@ def _make_paged_cache(k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page v_cache[g, :n] = v[start + ts : start + te] g += 1 - # Stack on a new leading axis so kv_cache.unbind(0) recovers (k_cache, v_cache). - kv_cache = torch.stack([k_cache, v_cache], dim=0) + kv_cache = _make_backend_paged_cache(k_cache, v_cache) return kv_cache, block_table def _make_impl(num_heads, head_dim, num_kv_heads): """Construct ModelOptSparseAttentionImpl with minimal valid kwargs.""" - return ModelOptSparseAttentionImpl( + return vllm_plugin.ModelOptSparseAttentionImpl( num_heads=num_heads, head_size=head_dim, scale=1.0 / (head_dim**0.5), @@ -132,7 +139,7 @@ def test_prefill_matches_contiguous(self): **_ACTIVE_PREFILL_SPARSE_KW, ) - # Build paged kv_cache shaped [2, num_blocks, page_size, num_kv_heads, head_dim]. + # Build the paged cache using the installed backend's K/V axis. kv_cache, block_table = _make_paged_cache( k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size ) @@ -174,7 +181,8 @@ def test_chunked_prefill_is_forwarded_to_kernel(self): block_table=torch.zeros(1, 1, device="cuda", dtype=torch.int32), ) q = torch.zeros(4, 2, 64, device="cuda", dtype=torch.float16) - kv_cache = torch.zeros(2, 1, 16, 2, 64, device="cuda", dtype=torch.float16) + k_cache = torch.zeros(1, 16, 2, 64, device="cuda", dtype=torch.float16) + kv_cache = _make_backend_paged_cache(k_cache, torch.zeros_like(k_cache)) out = impl.forward( layer=None, query=q, @@ -346,8 +354,7 @@ def test_page_size_inferred_from_k_cache(self): kv_cache, block_table = _make_paged_cache( k, v, b_start_loc, b_seq_len, num_kv_heads, head_dim, page_size ) - # Sanity: kv_cache axis 1 is page_size. - assert kv_cache.shape == (2, seq_len // page_size, page_size, num_kv_heads, head_dim) + assert kv_cache.shape[2] == page_size attn_metadata = SimpleNamespace( num_actual_tokens=seq_len, From 354a358c63a62a8399c629b38fa1c2389d21ca13 Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:22:23 -0700 Subject: [PATCH 06/11] [NVBug: 6563509] Drop Phi-3-vision / Phi-4-multimodal PTQ support (#2115) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Deprecation Resolves [NVBug 6563509](https://nvbugspro.nvidia.com/bug/6563509), where `hf_ptq.py` on Phi-4-multimodal-instruct died with `RuntimeError: Tensor.item() cannot be called on meta tensors`. The crash is real but not fixable on our side, and it is not the reason the model is unusable. Phi-4-multimodal's bundled remote code predates Transformers v5 and does not load on **any** version in our supported range (`transformers>=4.57,<5.15`): | Blocker | Where | |---|---| | `peft.get_peft_model` reads `prepare_inputs_for_generation`, gone since transformers 4.52 dropped `GenerationMixin` from `PreTrainedModel` | `modeling_phi4mm.py:1959` | | `_tied_weights_keys` declared as a list; Transformers 5.x calls `.keys()` on it in `post_init` | `modeling_phi4mm.py:1937` | | `int(torch.tensor(...))` in `__init__`, which cannot run on a meta device — the reported crash | `speech_conformer_encoder.py:1435` | The model card pins `transformers==4.48.2` / `peft==0.13.2`, so there is no overlap with our floor and nothing on our side can bridge it. The model is therefore dropped rather than worked around. **Phi-3-vision is dropped alongside it because it is the older, superseded model in the same family** — with its successor unsupportable there is no reason to keep carrying the predecessor. This is a product-scope call, not a separate compatibility finding: Phi-3-vision shares the list-valued `_tied_weights_keys` defect (`modeling_phi3_v.py:1214`) and so is likewise broken on Transformers 5.x, but it does **not** hit the `peft` blocker, and it was not re-verified on 4.57. Per the 0.46 changelog we have already bumped the floor to 4.57 and noted that "Transformers 4.x support will be dropped in a future release", so any remaining window closes on its own. Same reasoning already applied to VILA / NVILA in this release. **Removed** - the support-matrix row in `examples/hf_ptq/README.md` - `"Phi4MMForCausalLM": "phi4mm"` from `MODEL_NAME_TO_TYPE` - the multimodal-detection heuristics that only ever matched these two — `vision_lora`, `audio_processor`, `embd_layer.image_embd_layer`, and the `phi4mm` model-type check — in both `is_multimodal_model` and `_is_multimodal_config` - the `Phi3Image` / `PhiImage` exclusions in `is_embedding` - the phi4mm input-mode warning in `hf_ptq.py` - `modelopt_recipes/huggingface/phi4mm/` and its references in `modelopt_recipes/ptq.md` **Not changed:** the device-map sizing path (meta-device skeleton, `infer_auto_device_map`, and the `--gpu_max_mem_percentage` cap) keeps its original behavior. That cap is wanted exactly where it already fires — when the model is already offloading to CPU, where it costs little and the headroom is required. With the affected checkpoints removed, there is no supported model that trips the meta-device build, so there is nothing to work around here. Text-only **Phi-3/Phi-4** and **Phi-3.5-MoE** are natively supported by transformers and are untouched. On H200, `nvcr.io/nvidia/tensorrt-llm/release` (torch 2.12, transformers 5.5.4), against the real checkpoint: - **Version matrix** (vanilla transformers, no modelopt) — Phi-4-MM loads at 4.48.2 / 4.49.0 / 4.50.0 / 4.51.3 and fails at 4.53.3 / 4.56.2 / 4.57.1 (`AttributeError: 'Phi4MMModel' object has no attribute 'prepare_inputs_for_generation'`) and at 5.5.4 (meta-init, then tied-keys). This is what establishes that no supported version works. - `tests/examples/hf_ptq/test_example_utils.py` — 28 passed. - **Sweep**: `tests/examples/hf_ptq` + `tests/unit/torch/export` — failure set identical to the pre-change tree (GPU/model-dependent `test_vlm_ptq`, plus `test_quant_aware_conversion` scoped-mapping tests), so none are introduced here. - `pre-commit` clean on all changed files, including recipe validation. - Is this change backward compatible?: ❌ — PTQ for Phi-3-vision and Phi-4-multimodal is removed, along with the `huggingface/phi4mm/ptq/*` recipes. Phi-4-multimodal is already unloadable on every supported transformers version, so no working workflow regresses; Phi-3-vision is a deliberate scope removal as its superseded predecessor. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A — this is a deletion; the existing `test_get_model_*` / `test_resolve_init_config_*` tests are unchanged and still pass. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ — not yet run. Two related references were left in place deliberately; say the word and I'll fold them in: - `tests/examples/hf_ptq/test_deploy.py` still deploys the already-published `nvidia/Phi-4-multimodal-instruct-{NVFP4,FP8}` checkpoints. Those artifacts exist and serve fine; this PR only removes the ability to *produce* them. - `examples/torch_onnx/README.md` still lists Phi-4-multimodal-instruct. That is a separate ONNX pipeline that does not go through `get_model()` and was not tested here. Earlier revisions of this branch also reworked the device-map sizing so the meta-tensor crash could not occur. That was reverted in 701180ed6: the guard is correct as written, and every alternative either changed behavior for models that fit today or moved the guard somewhere it does not belong, for a crash that only ever affected the checkpoints this PR removes. --------- Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 3 +- examples/hf_ptq/README.md | 1 - examples/hf_ptq/example_utils.py | 6 ---- examples/hf_ptq/hf_ptq.py | 3 -- modelopt/torch/export/layer_utils.py | 7 +--- modelopt/torch/export/model_utils.py | 15 -------- .../huggingface/phi4mm/ptq/README.md | 13 ------- .../phi4mm/ptq/disabled_quantizers.yaml | 34 ------------------ .../phi4mm/ptq/nvfp4-kv_fp8_cast.yaml | 36 ------------------- modelopt_recipes/ptq.md | 6 ++-- 10 files changed, 5 insertions(+), 119 deletions(-) delete mode 100644 modelopt_recipes/huggingface/phi4mm/ptq/README.md delete mode 100644 modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml delete mode 100644 modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7d97f71fbcd..4df2af5608f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,7 +17,7 @@ Changelog - **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). - New ``sync_tied_input_amax`` helper max-merges per-side ``input_quantizer.amax`` across tied modules before export so single-backbone consumers that load one ``input_scale`` per parameter don't clip either side. - 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. + - 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 ``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 Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface//auto_quantize/``. @@ -60,6 +60,7 @@ Changelog - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. - Remove the deprecated ``examples/llm_qad`` Megatron-LM QAD example (deprecated in 0.45). Use the `megatron_bridge QAD example `_ instead, which provides a simpler Python-based interface and better model coverage. +- Dropped **Phi-4-multimodal** PTQ support in ``examples/hf_ptq`` (NVBug 6563509). Its bundled remote code predates Transformers v5 and no longer loads on any version ModelOpt supports (``transformers>=4.57``): it requires ``transformers<4.52`` because it reaches ``prepare_inputs_for_generation`` through ``peft``, which needs ``PreTrainedModel`` to still inherit ``GenerationMixin``, and it declares ``_tied_weights_keys`` as a list, which Transformers 5.x rejects. **Phi-3-vision** is dropped alongside it: it is the older, superseded model in the same family, so with its successor unsupportable there is no reason to keep carrying the predecessor. (Phi-3-vision shares the list-valued ``_tied_weights_keys`` defect and so is likewise broken on Transformers 5.x, though it does not hit the ``peft`` blocker.) The support-matrix row, the ``phi4mm`` model type, the multimodal-detection heuristics that only ever matched these two (``vision_lora`` / ``audio_processor`` / ``embd_layer.image_embd_layer``), the ``Phi3Image`` / ``PhiImage`` embedding-export exclusions, and the ``modelopt_recipes/huggingface/phi4mm/`` recipes have been removed. Text-only Phi-3/Phi-4 and Phi-3.5-MoE are unaffected. **Deprecations** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index a3967565ae0..8d6be9c3845 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -118,7 +118,6 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http | Whisper9 | ✅ | ❌ | ❌ | ❌ | - | | Nemotron-3 | ✅ | ❌ | ❌ | ❌ | ✅ | | Llava (VLM)11 | ✅ | ✅12 | ✅ | ✅ | - | -| Phi-3-vision, Phi-4-multimodal (VLM)11 | ✅ | ✅12 | ✅ | ✅ | ✅ | | Qwen2, 2.5-VL (VLM)11 | ✅ | ✅12 | ✅ | ✅ | ✅ | | Gemma 3 (VLM)11 | ✅ | - | - | - | - | | Nemotron VL (VLM)11,13 | ✅ | - | - | - | ✅ | diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d1031dd6084..b8eef827d95 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -168,12 +168,6 @@ def _is_multimodal_config(config): """Check if a config indicates a multimodal model (config-only version of is_multimodal_model).""" return ( hasattr(config, "vision_config") # Standard vision config (e.g., Qwen2.5-VL) - or getattr(config, "model_type", "") == "phi4mm" # Phi-4 multimodal - or hasattr(config, "vision_lora") # Vision LoRA configurations - or hasattr(config, "audio_processor") # Audio processing capabilities - or ( - hasattr(config, "embd_layer") and hasattr(config.embd_layer, "image_embd_layer") - ) # Image embedding layers or getattr(config, "is_encoder_decoder", False) # Encoder-decoder VL models or any( # Architecture-based detection for custom VL models (e.g., Nemotron-Parse) "conditionalgeneration" in arch.lower() for arch in getattr(config, "architectures", []) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 436c8867227..23265a47e70 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -699,9 +699,6 @@ def load_model(args: argparse.Namespace): # Left padding usually provides better calibration result. tokenizer.padding_side = "left" - if model_type == "phi4mm": - warnings.warn("Please set the default input_mode to InputMode.LANGUAGE before quantizing.") - return ( full_model, language_model, diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index d5f1fb2330d..de136fcd378 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -222,12 +222,7 @@ def is_conv(module: nn.Module) -> bool: def is_embedding(module: nn.Module) -> bool: """Returns whether the module is an embedding layer.""" module_type_name = type(module).__name__ - return ( - "Embedding" in module_type_name - and "Rotary" not in module_type_name - and "PhiImage" not in module_type_name - and "Phi3Image" not in module_type_name - ) + return "Embedding" in module_type_name and "Rotary" not in module_type_name def build_embedding_config(module: nn.Module, normalization_constant: float = 1) -> EmbeddingConfig: diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index 307ea9aac51..1729dbfffcf 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -44,7 +44,6 @@ "phi3small": "phi3small", "phi3": "phi3", "PhiMoEForCausalLM": "phi3", - "Phi4MMForCausalLM": "phi4mm", "phi": "phi", "TLGv4ForCausalLM": "phi", "MixtralForCausalLM": "llama", @@ -88,10 +87,6 @@ def is_multimodal_model(model): This function detects various multimodal model architectures by checking for: - Standard vision configurations (vision_config) - Language model attributes (language_model) - - Specific multimodal model types (phi4mm) - - Vision LoRA configurations - - Audio processing capabilities - - Image embedding layers - Nemotron-Parse conditional generation models Args: @@ -104,10 +99,6 @@ def is_multimodal_model(model): >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") >>> is_multimodal_model(model) True - - >>> model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-4-multimodal-instruct") - >>> is_multimodal_model(model) - True """ config = model.config @@ -118,12 +109,6 @@ def is_multimodal_model(model): return ( hasattr(config, "vision_config") # Standard vision config (e.g., Qwen2.5-VL) or hasattr(model, "language_model") # Language model attribute (e.g., LLaVA) - or getattr(config, "model_type", "") == "phi4mm" # Phi-4 multimodal - or hasattr(config, "vision_lora") # Vision LoRA configurations - or hasattr(config, "audio_processor") # Audio processing capabilities - or ( - hasattr(config, "embd_layer") and hasattr(config.embd_layer, "image_embd_layer") - ) # Image embedding layers or is_nemotron_parse # Nemotron-Parse conditional generation model ) diff --git a/modelopt_recipes/huggingface/phi4mm/ptq/README.md b/modelopt_recipes/huggingface/phi4mm/ptq/README.md deleted file mode 100644 index bedaf1fcb6b..00000000000 --- a/modelopt_recipes/huggingface/phi4mm/ptq/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# Phi-4-Multimodal PTQ recipes - -Phi-4-Multimodal is a multimodal model. Quantization should be applied only to -the language model; the speech, audio, image, and vision branches are kept in -full precision to avoid accuracy regressions on those modalities. - -| File | What's model-specific | -|------|-----------------------| -| `disabled_quantizers.yaml` | Reusable unit (`QuantizerCfgListConfig`). Merges the standard `default_disabled_quantizers` exclusions with Phi-4-MM ones (`*speech*`, `*audio*`, `*image*`, `*vision*`). Imported by recipes below as the single `disabled_quantizers` slot so they don't pull in two disabled-quantizer sets. | -| `nvfp4-kv_fp8_cast.yaml` | NVFP4 W4A4 model quantization + FP8 KV-cache cast (constant amax, no KV calibration). Identical numerics to the general `nvfp4` preset / `kv_fp8_cast` unit; what makes it model-specific is that it imports `disabled_quantizers.yaml` from this folder to skip the non-language branches. | - -Additional `-kv_fp8_cast.yaml` recipes can be generated for other formats -if needed; only `nvfp4-kv_fp8_cast.yaml` is shipped by default. diff --git a/modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml b/modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml deleted file mode 100644 index 1c6089f087f..00000000000 --- a/modelopt_recipes/huggingface/phi4mm/ptq/disabled_quantizers.yaml +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# QuantizerCfgList snippet of disabled quantizers for Phi-4-Multimodal. -# Splices in the standard `default_disabled_quantizers` exclusions and appends -# Phi-4-MM-specific ones so that only the language model is quantized; -# speech/audio/image/vision branches are skipped. Recipes that import this -# should NOT also import `default_disabled_quantizers`. - -# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig -imports: - default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers ---- - - $import: default_disabled_quantizers - - quantizer_name: '*speech*' - enable: false - - quantizer_name: '*audio*' - enable: false - - quantizer_name: '*image*' - enable: false - - quantizer_name: '*vision*' - enable: false diff --git a/modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml b/modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml deleted file mode 100644 index dfb1be1778d..00000000000 --- a/modelopt_recipes/huggingface/phi4mm/ptq/nvfp4-kv_fp8_cast.yaml +++ /dev/null @@ -1,36 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Phi-4-Multimodal-specific PTQ recipe for the `nvfp4` quantization format. -# Equivalent to the general `nvfp4` preset with quantization disabled -# on non-language branches. - -imports: - base_disable_all: configs/ptq/units/base_disable_all - w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 - disabled_quantizers: huggingface/phi4mm/ptq/disabled_quantizers - kv_fp8_cast: configs/ptq/units/kv_fp8_cast - -metadata: - recipe_type: ptq - description: 'Phi-4-Multimodal PTQ recipe (nvfp4): same numerics as the general nvfp4 preset, applied to the language model only (speech, audio, image, - and vision branches are skipped).' -quantize: - algorithm: max - quant_cfg: - - $import: base_disable_all - - $import: w4a4_nvfp4_nvfp4 - - $import: kv_fp8_cast - - $import: disabled_quantizers diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 22d6fbd5ff5..c1be8774976 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -232,7 +232,7 @@ that baseline. The deviations come in four kinds: |------|-------------------------------------|----------| | **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `minimax_m3_vl`, `qwen3_5`, `qwen3_5_moe`, `vit`, `nemotron_llama` | | **Algorithm override** | Same numerics & scope, but the *calibration algorithm* is tweaked because the default breaks or regresses | `gemma`, `gemma4`, `mpt` | -| **Extra exclusions** | Adds disabled-quantizer patterns so non-language branches stay full precision | `nemotron_vl`, `phi4mm`, `diffusion_gemma` | +| **Extra exclusions** | Adds disabled-quantizer patterns so non-language branches stay full precision | `nemotron_vl`, `diffusion_gemma` | | **Checkpoint mirror** | A mixed-precision map reproducing one published checkpoint exactly | `models/nvidia/Nemotron-3-*`, `models/nvidia/Mistral-Medium-3.5-128B-NVFP4` | The numerics and standard exclusions are still inherited from `configs/` @@ -312,7 +312,7 @@ These quantize the **same layers** as the general recipes; only the *Why special:* identical scope/numerics to a general scheme, but a general recipe's default algorithm would overflow or regress here. -### Extra exclusions — `nemotron_vl`, `phi4mm`, `diffusion_gemma` +### Extra exclusions — `nemotron_vl`, `diffusion_gemma` Each of these is **numerically identical** to a general recipe. What makes them special is a model-local `disabled_quantizers.yaml` unit that *extends* the @@ -322,8 +322,6 @@ standard exclusions so a model-specific branch stays in full precision: `nvfp4_default-kv_fp8_cast` numerics, adding `*vision*`, `*image*`, `*radio*`, `*visual*`, `*encoder*`, `*model_encoder*` so only the language decoder is quantized. -- **`phi4mm`** (Phi-4-Multimodal) — general `nvfp4_default-kv_fp8_cast` - numerics, adding `*speech*`, `*audio*`, `*image*`, `*vision*`. - **`diffusion_gemma`** (block-diffusion encoder-decoder text LLM on a Gemma4 MoE backbone) — general `nvfp4_experts_only-kv_fp8_cast` numerics, adding `*self_conditioning*`: the self-conditioning network is text-only and never From dc5dea33563c488ea9b3463249407f7806093cb8 Mon Sep 17 00:00:00 2001 From: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:49:06 -0700 Subject: [PATCH 07/11] [NVBug 6571812] Fix NemotronH dense MLP quantization recipes (#2133) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes NVBug 6571812 by enabling NVFP4 quantization for NemotronH dense MLP projections registered as `backbone.layers.*.mixer.up_proj` and `backbone.layers.*.mixer.down_proj`. The existing MLP-only recipes enable `*mlp*` quantizers after disabling all quantizers. NemotronH's dense projection names do not contain `mlp`, so no dense MLP weights were selected and HF export emitted `quant_algo: null`. This change: - adds a reusable `mixer_mlp_nvfp4` recipe unit for dynamic NVFP4 - applies it to the MLP-only and OMLP-only recipes and presets - adds equivalent static selectors to the MSE recipe - verifies that dense mixer projections match while Mamba `in_proj`/`out_proj` and `shared_experts` remain excluded - documents the fix in the 0.46 changelog ### Usage ```bash python examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 \ --recipe general/ptq/nvfp4_mlp_only-kv_fp8_cast \ --export_fmt hf \ --export_path /path/to/output ``` ### Testing - `python -m pytest tests/unit/recipe`: 231 passed - `pre-commit run`: all applicable hooks passed - Remote GB300 PTQ/export smoke test: all 34 expected dense MLP projections used NVFP4, with 0 missed or unexpected projections - TensorRT-LLM `1.3.0rc6` TP1 canary: loaded all 445 tensors, `/health` and `/v1/models` returned HTTP 200, and chat completion returned Paris - Slurm job `2900754`: `COMPLETED`, exit `0:0` ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: N/A ### Additional Information - NVBug 6571812 - Commit `e9b8f9ffc55e66cdbbacdbbb287e0dedf96ab082` is signed and includes the DCO sign-off. ## Summary by CodeRabbit - **New Features** - Added NVFP4 quantization support for dense mixer MLP projections. - Expanded NVFP4 MLP and OMLP preset recipes, including dynamic input and static-weight configurations. - Improved compatibility across supported quantization recipe combinations. - **Documentation** - Documented the new NVFP4 mixer MLP quantization option. - Added release notes covering calibration, grouped-MoE quantization, logging, integrations, and compatibility updates. - **Tests** - Added coverage validating NVFP4 recipe behavior and projection selection. Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../ptq/presets/model/nvfp4_mlp_only.yaml | 2 ++ .../ptq/presets/model/nvfp4_omlp_only.yaml | 2 ++ modelopt_recipes/configs/ptq/units/README.md | 1 + .../configs/ptq/units/mixer_mlp_nvfp4.yaml | 34 ++++++++++++++++++ .../general/ptq/nvfp4_mlp_only-kv_fp8.yaml | 2 ++ .../ptq/nvfp4_mlp_only-kv_fp8_cast.yaml | 2 ++ .../ptq/nvfp4_mlp_only-novit-kv_fp8.yaml | 2 ++ .../ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml | 12 +++++++ .../general/ptq/nvfp4_omlp_only-kv_fp8.yaml | 2 ++ .../ptq/nvfp4_omlp_only-kv_fp8_cast.yaml | 2 ++ tests/unit/recipe/test_loader.py | 35 +++++++++++++++++++ 12 files changed, 97 insertions(+) create mode 100644 modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4df2af5608f..6bbb95f8250 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -74,6 +74,7 @@ Changelog **Bug Fixes** +- Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` and ``mixer.down_proj``, which were missed by the previous ``*mlp*`` selector and produced checkpoints with a null ``quant_algo`` (NVBug 6571812). - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export. diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml index c5d36fd9236..c8c269c1286 100644 --- a/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_mlp_only.yaml @@ -21,6 +21,7 @@ imports: block_sparse_moe_nvfp4: configs/ptq/units/block_sparse_moe_nvfp4 experts_nvfp4: configs/ptq/units/experts_nvfp4 default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 algorithm: max @@ -32,6 +33,7 @@ quant_cfg: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - $import: block_sparse_moe_nvfp4 - $import: experts_nvfp4 - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml index 82bf401ea9f..cc908ced186 100644 --- a/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_omlp_only.yaml @@ -20,6 +20,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all block_sparse_moe_nvfp4: configs/ptq/units/block_sparse_moe_nvfp4 default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 algorithm: max @@ -37,5 +38,6 @@ quant_cfg: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - $import: block_sparse_moe_nvfp4 - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/units/README.md b/modelopt_recipes/configs/ptq/units/README.md index cd738f62626..db37b9222ca 100644 --- a/modelopt_recipes/configs/ptq/units/README.md +++ b/modelopt_recipes/configs/ptq/units/README.md @@ -30,4 +30,5 @@ recipes (under `general/` or `models/`) or presets (under `presets/`). | `w4a4_nvfp4_nvfp4.yaml` | NVFP4 weight + activation quantizer entries (W4A4); supported on Blackwell+ GPUs | | `block_sparse_moe_nvfp4.yaml` | NVFP4 W4A4 on `*block_sparse_moe*` weight/input quantizers | | `experts_nvfp4.yaml` | NVFP4 W4A4 on `*.experts.*` weight/input quantizers | +| `mixer_mlp_nvfp4.yaml` | NVFP4 W4A4 on dense `*.mixer.{up,down}_proj` weight/input quantizers | | `attention_qkv_fp8.yaml` | FP8 E4M3 on attention q/k/v bmm and softmax quantizers | diff --git a/modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml b/modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml new file mode 100644 index 00000000000..f38b19d67e0 --- /dev/null +++ b/modelopt_recipes/configs/ptq/units/mixer_mlp_nvfp4.yaml @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizerCfgList snippet that enables dynamic NVFP4 on dense MLP projections +# registered directly under a ``mixer`` module. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig +imports: + nvfp4: configs/numerics/nvfp4 +--- + - quantizer_name: '*.mixer.up_proj.weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.up_proj.input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.down_proj.weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.down_proj.input_quantizer' + cfg: + $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml index a4cf71a1dbd..4fd2e0a7558 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8: configs/ptq/units/kv_fp8 @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml index 225ecf7f086..a12951bda65 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-kv_fp8_cast.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8_cast: configs/ptq/units/kv_fp8_cast @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml index 7bbf21393f8..2d80d7a5701 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only-novit-kv_fp8.yaml @@ -33,6 +33,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8: configs/ptq/units/kv_fp8 @@ -52,6 +53,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml index 71c354ee1b1..18fed45d266 100644 --- a/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast.yaml @@ -43,6 +43,18 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - quantizer_name: '*.mixer.up_proj.weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*.mixer.up_proj.input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.mixer.down_proj.weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*.mixer.down_proj.input_quantizer' + cfg: + $import: nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4_static diff --git a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml index 5348e8c7123..41541e4b2e4 100644 --- a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8: configs/ptq/units/kv_fp8 @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml index ba9e1e1c27a..14da2d92d4a 100644 --- a/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_omlp_only-kv_fp8_cast.yaml @@ -18,6 +18,7 @@ imports: base_disable_all: configs/ptq/units/base_disable_all default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + mixer_mlp_nvfp4: configs/ptq/units/mixer_mlp_nvfp4 nvfp4: configs/numerics/nvfp4 kv_fp8_cast: configs/ptq/units/kv_fp8_cast @@ -36,6 +37,7 @@ quantize: - quantizer_name: '*mlp*input_quantizer' cfg: $import: nvfp4 + - $import: mixer_mlp_nvfp4 - quantizer_name: '*block_sparse_moe*weight_quantizer' cfg: $import: nvfp4 diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index e15b897a224..f88240b4dcf 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -213,6 +213,41 @@ def test_nvfp4_mlp_only_novit_recipe_disables_vision_quantizers(): assert {"*visual*", "*vision_tower*"} <= disabled_quantizers +@pytest.mark.parametrize( + "recipe_path", + [ + "general/ptq/nvfp4_mlp_only-kv_fp8", + "general/ptq/nvfp4_mlp_only-novit-kv_fp8", + "general/ptq/nvfp4_mlp_only-kv_fp8_cast", + "general/ptq/nvfp4_mlp_only_mse-kv_fp8_cast", + "general/ptq/nvfp4_omlp_only-kv_fp8", + "general/ptq/nvfp4_omlp_only-kv_fp8_cast", + ], +) +def test_nvfp4_mlp_only_recipes_match_nemotron_h_dense_mlp(recipe_path): + recipe = load_recipe(recipe_path) + enabled_patterns = [ + entry["quantizer_name"] + for entry in recipe.quantize.model_dump()["quant_cfg"] + if entry["enable"] + ] + + for quantizer_name in ( + "backbone.layers.0.mixer.up_proj.weight_quantizer", + "backbone.layers.0.mixer.up_proj.input_quantizer", + "backbone.layers.0.mixer.down_proj.weight_quantizer", + "backbone.layers.0.mixer.down_proj.input_quantizer", + ): + assert any(fnmatch(quantizer_name, pattern) for pattern in enabled_patterns) + + for quantizer_name in ( + "backbone.layers.0.mixer.in_proj.weight_quantizer", + "backbone.layers.0.mixer.out_proj.input_quantizer", + "backbone.layers.0.mixer.shared_experts.up_proj.weight_quantizer", + ): + assert not any(fnmatch(quantizer_name, pattern) for pattern in enabled_patterns) + + @pytest.mark.parametrize( "recipe_path", [ From 114961a1303c28367188f74af48054eba683270c Mon Sep 17 00:00:00 2001 From: Jenny Chen Date: Tue, 11 Aug 2026 14:05:46 -0400 Subject: [PATCH 08/11] Add Nemotron Lightning 3.5 NVFP4 recipe and QAD example (#2146) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: New example Add Nemotron Lightning 3.5 NVFP4 recipe and QAD example Also exclude MTP in default disabled quantizers ```python ``` Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A - Did you get Claude approval on this PR?: ✅ / ❌ / N/A - **New Features** - Added an end-to-end NVFP4 quantization, distillation, and export workflow for NVIDIA Nemotron 3.5 Lightning 30B-A3B. - Added a PTQ configuration supporting NVFP4 W4A16 quantization with optimized scaling and FP8 support for selected components. - **Bug Fixes** - Preserved custom model output locations when provided, while retaining the existing default path. - **Configuration** - Disabled quantization for MTP modules by default. - Updated an existing Nemotron workflow to use the aggressive NVFP4 quantization profile. --------- Signed-off-by: Jennifer Chen Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 2 +- .../units/default_disabled_quantizers.yaml | 2 + .../ptq/nvfp4_w4a16.yaml | 0 .../ptq/nvfp4-max-calib.yaml | 0 .../ptq/nvfp4-mse.yaml | 0 .../ptq/nvfp4-4o6.yaml | 0 .../ptq/w4a16_nvfp4_4o6.yaml | 76 +++++++++ modelopt_recipes/ptq.md | 10 +- .../common/megatron_lm/quantize/quantize.sh | 2 +- .../megatron_lm_ptq.yaml | 4 +- .../megatron_lm_ptq.yaml | 4 +- .../megatron_lm_qad.yaml | 151 ++++++++++++++++++ 12 files changed, 242 insertions(+), 9 deletions(-) rename modelopt_recipes/huggingface/models/nvidia/{Nemotron-3-Nano-4B => Nemotron-3-Nano-4B-BF16}/ptq/nvfp4_w4a16.yaml (100%) rename modelopt_recipes/huggingface/models/nvidia/{Nemotron-3-Super-120B-A12B => Nemotron-3-Super-120B-A12B-BF16}/ptq/nvfp4-max-calib.yaml (100%) rename modelopt_recipes/huggingface/models/nvidia/{Nemotron-3-Super-120B-A12B => Nemotron-3-Super-120B-A12B-BF16}/ptq/nvfp4-mse.yaml (100%) rename modelopt_recipes/huggingface/models/nvidia/{Nemotron-3-Ultra-550B-A55B => Nemotron-3-Ultra-550B-A55B-BF16}/ptq/nvfp4-4o6.yaml (100%) create mode 100644 modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6bbb95f8250..1168c0f2b16 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -105,7 +105,7 @@ Changelog - The PTQ example scripts ``examples/llm_ptq/hf_ptq.py``, ``examples/llm_ptq/multinode_ptq.py`` and ``examples/megatron_bridge/quantize.py`` now derive their ``--qformat`` / ``--kv_cache_qformat`` (``--quant_cfg`` / ``--kv_cache_quant`` for Megatron-Bridge) CLI vocabularies by discovering the YAML presets under ``modelopt_recipes/configs/ptq/presets/{model,kv}/`` rather than carrying hardcoded ``QUANT_CFG_CHOICES`` / ``KV_QUANT_CFG_CHOICES`` tables. The discovery helper, alias table and ready-built ``QUANT_CFG_CHOICES`` / ``KV_QUANT_CFG_CHOICES`` mappings now live in ``modelopt.recipe.presets`` and are shared by all three scripts. Presets are loaded eagerly into a plain dict at import. Adding a new preset YAML makes it available on the CLI of all three with no script change — note this means each script now accepts every preset under those directories, not just a previously curated subset. All previously-supported short names (``int8_sq``, ``nvfp4_awq``, ``fp8_pb_wo``, ``nvfp4_mse``, ``w4a8_awq``, ``nvfp4_local_hessian``, ``fp8_pc_pt``, ``int8_wo``) keep working via a small deprecation alias table; new formats should be exposed as preset YAMLs (or, longer term, as full ``--recipe`` recipes). - Add ``configs/ptq/presets/kv/fp8_cast.yaml`` and ``configs/ptq/presets/kv/nvfp4_cast.yaml``, promoting ``fp8_cast`` / ``nvfp4_cast`` to first-class KV presets composed from the existing ``kv_fp8_cast`` / ``kv_nvfp4_cast`` unit fragments. The previous runtime ``use_constant_amax`` post-edit in ``hf_ptq.py`` is removed; ``use_constant_amax: true`` now lives in the YAML and is therefore authoritative. **Custom (out-of-tree) recipes that target a cast KV format must set ``use_constant_amax: true`` themselves on the ``[kv]_bmm_quantizer`` config** — in-tree recipes already do via the ``kv_*_cast`` units. - Add FP8 KV-cache cast variants for the partial-NVFP4 and weight-only general PTQ recipes: ``general/ptq/nvfp4_mlp_only-kv_fp8_cast``, ``general/ptq/nvfp4_experts_only-kv_fp8_cast``, ``general/ptq/nvfp4_omlp_only-kv_fp8_cast``, and ``general/ptq/nvfp4_weight_only-kv_fp8_cast``. These compose the same model-quant configs as their ``-kv_fp8`` siblings with the ``kv_fp8_cast`` unit (constant-amax FP8 KV cache, no KV calibration forward pass). -- Add Nemotron-3-Super-120B-A12B PTQ recipes ``modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml`` (MSE-mixed) and ``nvfp4-max-calib.yaml`` (max-calib mixed): NVFP4 W4A4 routed experts + FP8 per-tensor shared experts / Mamba in/out_proj + FP8 KV cache. +- Add Nemotron-3-Super-120B-A12B PTQ recipes ``modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yaml`` (MSE-mixed) and ``nvfp4-max-calib.yaml`` (max-calib mixed): NVFP4 W4A4 routed experts + FP8 per-tensor shared experts / Mamba in/out_proj + FP8 KV cache. - Group layerwise calibration options under a nested ``LayerwiseConfig`` and add two knobs: ``get_qdq_activations_from_prev_layer`` (correct GPTQ-Hessian vs max-calib activation semantics — defaults to True for GPTQ, False for max/mse/local_hessian) and ``save_every`` (gate per-window ``next_inputs.pt`` activation-cache writes). Legacy bool ``layerwise`` and flat ``layerwise_checkpoint_dir`` keys still work; the bool form emits a ``DeprecationWarning``. - Add two layerwise-calibration memory optimizations: ``calib_mutates_weights`` (set False for amax-only algorithms — max/mse/local_hessian — to skip the per-layer weight checkpoint blob and in-memory writeback, persisting only quantizer state), and meta-device skip-layer placeholders (already-calibrated layers emit zero-filled ``meta`` tensors instead of real-device buffers, eliminating their activation memory — models with real-device inter-layer ops on the hidden state are unsupported). - Add ``examples/alpamayo`` showing FP8, NVFP4, and AutoQuantize (mixed-precision) quantization of the Alpamayo (formerly Alpamayo-R1) ~10B vision-language-action model, with a joint VLM + diffusion calibration loop and both fake-quant and ``--real-quant`` packed-checkpoint export. See `examples/alpamayo/README.md `_ for details. diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 5e48dc73b7e..3aadadd289c 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -38,6 +38,8 @@ enable: false - quantizer_name: '*router*' enable: false + - quantizer_name: 'mtp.*' + enable: false - quantizer_name: 'output.*' enable: false # Multimodal vision branch: keep the vision encoder (SigLIP / ViT) and any diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B/ptq/nvfp4_w4a16.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B/ptq/nvfp4_w4a16.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-max-calib.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-max-calib.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-max-calib.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-max-calib.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6.yaml similarity index 100% rename from modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6.yaml rename to modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6.yaml diff --git a/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml new file mode 100644 index 00000000000..ab6933007b8 --- /dev/null +++ b/modelopt_recipes/huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6.yaml @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Quantization config for nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4: +# NVFP4 weights use the four_over_six (4/6) per-block scale selection. +# Weight-only; activations BF16. MSE selects per-block M=6 (mult 1.0) vs +# M=4 (mult 1.5); the choice is baked into the amax by MSE. +# +# HF -> MCore name mapping for Nemotron-3.5-Lightning-30B-A3B (Mamba/Attn/MoE hybrid): +# - mixer.experts..{up,down}_proj -> mlp.experts.local_experts..linear_fc{1,2} +# - mixer.shared_experts.{up,down}_proj -> mlp.shared_experts.linear_fc{1,2} +# - mixer.in_proj / out_proj -> mixer.in_proj / out_proj (same name; W4A16) +# - lm_head -> output_layer (W4A16) +imports: + nvfp4_four_over_six: configs/numerics/nvfp4_four_over_six + fp8_default: configs/numerics/fp8 +metadata: + recipe_type: ptq + description: > + Lightning 3.5 W4A16 PTQ, NVFP4 four_over_six (4/6) weight scales. + Routed MoE experts, shared experts, and lm_head; FP8 Mamba in_proj/out_proj. + W4A16 use weight-only NVFP4 4/6 (static, MSE-selected). KV cache FP8. + Attention BF16. +quantize: + # 4/6: MSE selects per-block between M=6 (keep amax) and M=4 (amax x 6/4). + algorithm: + method: mse + fp8_scale_sweep: false + start_multiplier: 1.0 # M=6 (keep amax) + stop_multiplier: 1.5 # M=4 (amax x 6/4) + step_size: 0.5 # candidates [1.0, 1.5] + quant_cfg: + - quantizer_name: '*' + enable: false + # W4A16 NVFP4 4/6 weight-only (block 16, static, four_over_six: true). HF + MCore names. + - quantizer_name: '*mixer.experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: '*mixer.shared_experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: '*mlp.experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: '*mlp.shared_experts*weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + # FP8 Mamba projections. + - quantizer_name: '*mixer.in_proj*weight_quantizer' + enable: true + cfg: {$import: fp8_default} + - quantizer_name: '*mixer.in_proj*input_quantizer' + enable: true + cfg: {$import: fp8_default} + - quantizer_name: '*mixer.out_proj*weight_quantizer' + enable: true + cfg: {$import: fp8_default} + - quantizer_name: '*mixer.out_proj*input_quantizer' + enable: true + cfg: {$import: fp8_default} + # KV cache -> FP8. + - quantizer_name: '*[kv]_bmm_quantizer' + enable: true + cfg: + num_bits: e4m3 + # lm_head (output_layer) W4A16 4/6 weight-only. + - quantizer_name: 'output_layer.weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + - quantizer_name: 'lm_head.weight_quantizer' + enable: true + cfg: {$import: nvfp4_four_over_six} + # Keep the entire MTP subtree in BF16. This rule must remain last so it + # overrides the broad expert, Mamba, KV-cache, and output-layer selectors. + - quantizer_name: 'mtp.*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index c1be8774976..847117acc0c 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -342,7 +342,7 @@ checkpoint's** quant config verbatim: `nvidia/Mistral-Medium-3.5-128B-NVFP4`: decoder MLP layers 4–86 use NVFP4 W4A4, edge MLP layers 0–3 and 87 use FP8 W8A8, and all attention projections and the KV cache use FP8. It uses max calibration. -- **`Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse`** mirrors +- **`Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse`** mirrors `nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4` exactly — a hybrid **Mamba-MoE** with a hand-mapped, **per-component** precision scheme: - MoE routed experts → NVFP4 W4A4, `group_size 16`, **static** weight scales @@ -353,13 +353,17 @@ checkpoint's** quant config verbatim: `nvfp4-mse.yaml` uses MSE calibration with an FP8-scale sweep (matches the release); `nvfp4-max-calib.yaml` is the identical layer map under plain `max` calibration, kept for comparison. -- **`Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6`** follows the same Super-style +- **`Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6`** follows the same Super-style component map (routed experts NVFP4 W4A4 block-16; shared experts + Mamba `in/out_proj` + KV cache FP8; everything else BF16), but the routed-expert weights use **Four-over-Six (4/6)** NVFP4: an MSE search picks each weight's amax multiplier from `[1.0, 1.5]` (M=6 vs. M=4). Activations stay dynamic NVFP4 (not MSE-calibrated). -- **`Nemotron-3-Nano-4B/ptq/nvfp4_w4a16`** mirrors the GGUF **Q4_K_M** bit +- **`Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6`** applies + Four-over-Six NVFP4 W4A16 to routed experts, shared experts, and the language + model head; Mamba `in/out_proj` weights and inputs plus the KV cache use FP8, + while attention remains BF16. +- **`Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16`** mirrors the GGUF **Q4_K_M** bit allocation of the Nemotron-H hybrid, mapped onto NVFP4/FP8 **per layer**: Q4_K/Q5_0 linears → NVFP4 W4A4 (attention q/k/v/o kept uniform so export can fuse them), the Q6_K MLP `down_proj` layers → FP8 W8A8, embeddings → NVFP4 diff --git a/tools/launcher/common/megatron_lm/quantize/quantize.sh b/tools/launcher/common/megatron_lm/quantize/quantize.sh index 083ef7399f0..c6b72094c28 100755 --- a/tools/launcher/common/megatron_lm/quantize/quantize.sh +++ b/tools/launcher/common/megatron_lm/quantize/quantize.sh @@ -34,7 +34,7 @@ if [[ -z ${HF_MODEL_CKPT} ]]; then fi # Persist PTQ ckpt + HF export under /cicd ($SLURM_JOB_DIR/cicd) so later # experiments can re-use them. -export MLM_MODEL_SAVE="/cicd/megatron-lm/${MLM_MODEL_CFG}" +export MLM_MODEL_SAVE="${MLM_MODEL_SAVE:-/cicd/megatron-lm/${MLM_MODEL_CFG}}" # If QUANT_CFG is a recipe path, collapse to a flat tag (strip dirs + .yaml/.yml). _QUANT_CFG_TAG="$(basename "${QUANT_CFG}")" _QUANT_CFG_TAG="${_QUANT_CFG_TAG%.yaml}" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml index 05f6d986f43..6ec49dd7a0a 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_lm_ptq.yaml @@ -29,7 +29,7 @@ pipeline: - --calib-size 32 environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 # MMLU + Export run as separate tasks; quantize.sh does quantize only. - RUN_MMLU: "false" @@ -52,7 +52,7 @@ pipeline: script: common/megatron_lm/export/export.sh environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 - TP: "1" - PP: "4" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml index d297a1c4b50..dfe363da634 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16/megatron_lm_ptq.yaml @@ -30,7 +30,7 @@ pipeline: - --calib-size 32 environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6 + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6 - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 # MMLU + Export run as separate tasks; quantize.sh does quantize only. - RUN_MMLU: "false" @@ -53,7 +53,7 @@ pipeline: script: common/megatron_lm/export/export.sh environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 - - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B/ptq/nvfp4-4o6 + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3-Ultra-550B-A55B-BF16/ptq/nvfp4-4o6 - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 - TP: "1" - PP: "12" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml new file mode 100644 index 00000000000..b9201c27784 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml @@ -0,0 +1,151 @@ +# NVIDIA Nemotron 3.5 Lightning 30B-A3B NVFP4 quantization-aware distillation (QAD). +# +# The pipeline converts the HuggingFace BF16 model to an MCore teacher, +# quantizes a separate MCore student, distills the student for 400 iterations, +# and exports the resulting checkpoint. The training task uses an explicit +# Nemotron-Post-Training-Dataset-v2 chat shard so Hugging Face Datasets does not +# prepare the repository's other large splits. +# +# PTQ topology: 1 B200 node x 4 GPUs, TP=1, PP=1, CP=1, EP=4, ETP=1. +# QAD topology: 2 B200 nodes x 4 GPUs, TP=2, PP=1, CP=1, EP=4, ETP=1. +# With micro-batch-size=1 and global-batch-size=16, train-samples=6400 produces +# 400 iterations. To use another sequence length, change both --seq-length and +# --max-position-embeddings. Our final QAD run used a sequence length of 524,288. +# +# Requirements: +# - The BF16 model is mounted at /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16. +# - HF_TOKEN can access the gated nvidia/Nemotron-Post-Training-Dataset-v2 dataset. +# +# Usage from tools/launcher: +# source .env-slurm +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/megatron_lm_qad.yaml --yes + +job_name: Nemotron-3.5-Lightning-30B-A3B_QAD_32k_400iter +pipeline: + allow_to_fail: false + skip: false + note: "NVFP4 TEGroupedMLP QAD at 32K for 400 iterations on one explicit Nemotron post-training chat shard" + + # Import the BF16 Hugging Face checkpoint as the MCore teacher checkpoint. + task_0: + script: common/megatron_bridge/import/import.sh + environment: + - HF_MODEL_ID: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - OUTPUT_DIR: /cicd/megatron-lm-bf16/nvidia + - TORCH_DTYPE: bfloat16 + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 4 + + # Quantize the teacher checkpoint into the NVFP4 student checkpoint. + task_1: + script: common/megatron_lm/quantize/quantize.sh + args: + - --seq-length 32768 --max-position-embeddings 32768 + - --calib-size 768 + - --skip-generate + - --export-default-te-spec + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6 + - MLM_MODEL_CKPT: /cicd/megatron-lm-bf16/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-MCore + - MLM_MODEL_SAVE: /cicd/megatron-lm/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-W4A16 + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - RUN_MMLU: "false" + - RUN_EXPORT: "false" + - DP: "1" + - CP: "1" + - TP: "1" + - PP: "1" + - EP: "4" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + + # Distill the quantized student from the BF16 teacher on chat-template data. + task_2: + script: common/megatron_lm/train/sft.sh + args: + # Data + - --seq-length 32768 --max-position-embeddings 32768 + - --micro-batch-size 1 --global-batch-size 16 + - --train-samples 6400 + - --lr-decay-samples 6400 + - --lr-warmup-samples 0 + - --split 99,1,0 + - --finetune-data-split chat + - --finetune-data-files data/chat-00000-of-00012.parquet + # QAD + - --modelopt-enabled + - --export-default-te-spec + - --export-kd-teacher-load /cicd/megatron-lm-bf16/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-MCore + - --attention-dropout 0.0 --hidden-dropout 0.0 + - --no-check-for-nan-in-loss-and-grad + - --recompute-granularity selective + - --recompute-modules layernorm moe + - --sequence-parallel + - --ckpt-fully-parallel-load --ckpt-fully-parallel-save + # Optimizer + - --lr 5.0e-6 + - --lr-decay-style constant + - --clip-grad 1.0 --weight-decay 0.0 + - --adam-beta1 0.9 --adam-beta2 0.95 + - --init-method-std 0.010 + - --use-distributed-optimizer + # Evaluation and checkpoints + - --eval-iters 2 --eval-interval 25 + - --save-interval 50 --log-interval 10 + - --dist-ckpt-strictness log_all + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - MLM_MODEL_CKPT: /cicd/megatron-lm/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-W4A16 + - MLM_MODEL_SAVE: /cicd/megatron-lm-qad/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - MLM_TRAIN_SCRIPT: finetune + - DATASET: nvidia/Nemotron-Post-Training-Dataset-v2 + - DP: "1" + - CP: "1" + - TP: "2" + - PP: "1" + - EP: "4" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 2 + ntasks_per_node: 4 + gpus_per_node: 4 + + task_3: + script: common/megatron_lm/export/export.sh + args: + - --export-default-te-spec + - --dist-ckpt-strictness log_all + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - MLM_MODEL_CKPT: /cicd/megatron-lm-qad/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - EXPORT_DIR: /cicd/export/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16_NVFP4_QAD + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - DP: "1" + - CP: "1" + - TP: "1" + - PP: "4" + - EP: "1" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06.00 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 From d9d1bf5cb4c9cb0aefac35ca8026268c2af907aa Mon Sep 17 00:00:00 2001 From: sugunav14 <178320438+sugunav14@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:51:02 -0700 Subject: [PATCH 09/11] Bug fix: 6542481 (#2064) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Fixes `AssertionError: Model already has modelopt state!` when exporting a QLoRA checkpoint (NVBug 6542481). The QLoRA output is adapter-only, so `from_pretrained` resolves the quantized base model and already restores the ModelOpt state; `export.py` then restored a second time. Fixing that exposed two more breakages on the same path, also fixed here: - `_restore_qtensor_wrappers` missed every module — PEFT renames the compressed linears to `.base_layer`, so no weight got re-wrapped and the packed NVFP4 weight hit a shape error. - `postprocess_state_dict` dropped `weight_scale_2` (missing from the QLoRA rename map), leaving the exported checkpoint impossible to dequantize. No API change — `examples/llm_qat/export.py --pyt_ckpt_path --export_path ` now completes on the documented quantize → train → export flow. Reproduced in the reported environment (TRT-LLM 1.3.0rc22, transformers 5.5.4, NVFP4). - Added the missing export step to `test_qwen3_qlora_nvfp4` and a unit test for the QLoRA `base_layer` rename; both fail without the fix. - Exported base model is byte-identical to a plain PTQ export; dequantized NVFP4 weights match the bf16 original (worst rel. error 0.10). - No regressions: `tests/gpu/torch/export/test_export.py` (49 passed), save/load plugin tests. - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ❌ — can add if wanted - Did you get Claude approval on this PR?: ❌ — not run yet Fixes NVBug 6542481. * **Bug Fixes** * Improved QLoRA checkpoint export and restoration across supported model configurations. * Preserved secondary weight-scale information and other deployment tensors in exported checkpoints. * Corrected handling of quantized base-layer weights after adapter reparenting. * Prevented duplicate state restoration when checkpoints already include the required model state. * Removed internal adapter prefixes and quantizer details from exported state data. * **Tests** * Added validation for packed weights, quantization metadata, required scales, and removal of embedded adapter layers. * Added regression coverage for QLoRA state processing and quantized weight restoration. --------- Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + examples/llm_qat/export.py | 9 ++- modelopt/torch/export/quant_utils.py | 24 +++---- modelopt/torch/opt/plugins/transformers.py | 38 ++++++++--- tests/examples/llm_qat/test_llm_qat.py | 68 ++++++++++++++++++- tests/gpu/torch/export/test_export.py | 35 ++++++++++ .../torch/opt/plugins/test_hf_patching.py | 63 +++++++++++++++++ 7 files changed, 211 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1168c0f2b16..eff4cea75bf 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -87,6 +87,7 @@ Changelog - Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors. - Fix HF checkpoint export failing with ``AttributeError: 'list' object has no attribute 'keys'`` for models whose modeling code still declares tied weights in the ``transformers<5`` list format (NVBug 6518665, observed on ``stepfun-ai/Step-3.7-Flash``). transformers 5.0 changed ``_tied_weights_keys`` to a ``{target: source}`` dict and ``save_pretrained`` calls ``.keys()`` on every submodule's declaration without a type check, so such models — common among ``trust_remote_code`` checkpoints — load fine but die at the end of PTQ, after calibration. ModelOpt's ``save_pretrained`` patch now normalizes a list-style declaration to the equivalent dict for the duration of the save (each entry mapped to itself, which is what the legacy list meant) and restores the original attribute afterwards. - Fix unified HF export of multimodal models whose vision tower carries its own ``PrefixChange`` conversion (``LlavaForConditionalGeneration`` on ``transformers>=5.12`` — NVBug 6525511). transformers collects conversion mappings recursively and scopes each sub-model's transforms to that sub-module via ``scope_prefix``, matching only keys under that prefix. ModelOpt's quant-aware reverse conversion read the raw patterns and ignored ``scope_prefix``, so the vision tower's "add a ``vision_model.`` prefix" rule was applied to *every* key in the state dict: an exported llava-1.5-13b checkpoint had all 758 tensors moved under a bogus top-level ``vision_model.`` namespace (``vision_model.language_model.*``, ``vision_model.lm_head.*``), and vLLM rejected it with ``ValueError: There is no module or parameter named 'vision_model' in LlavaForConditionalGeneration``. Reverse rename rules now carry their scope and are applied only to keys under it, matching transformers' own ``WeightTransform._scoped_match`` semantics. ``Gemma3ForConditionalGeneration`` was affected identically and is fixed by the same change. +- Fix QLoRA export in ``examples/llm_qat/export.py`` failing with ``AssertionError: Model already has modelopt state!`` (NVBug 6542481). The QLoRA training output is an adapter-only checkpoint, so ``from_pretrained`` resolves the quantized base model from ``adapter_config.json`` and ``enable_huggingface_checkpointing`` already restores its ModelOpt state; the export then restored a second time. It now restores only when the loaded model is not already converted. Two further breakages on the same path are also fixed: ``_restore_qtensor_wrappers`` matched no modules because PEFT re-parents the quantized linear as ``.base_layer`` while ``q_tensor_state`` is keyed by the name it was saved with (the packed NVFP4 weight then reached ``F.linear`` and raised a shape error), and ``postprocess_state_dict`` silently dropped every ``base_layer.*`` key missing from a hand-maintained rename map — losing the NVFP4 ``weight_scale_2`` global scale and any linear ``bias`` (Qwen2-style q/k/v biases), and leaving ``base_layer`` in the exported AWQ ``pre_quant_scale`` key. The rename is now a generic ``.base_layer.`` strip. 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py index f48e85c3ee4..afe2bd4d1cf 100644 --- a/examples/llm_qat/export.py +++ b/examples/llm_qat/export.py @@ -23,7 +23,7 @@ import modelopt.torch.opt as mto from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.unified_export_hf import _export_transformers_checkpoint -from modelopt.torch.opt.conversion import restore_from_modelopt_state +from modelopt.torch.opt.conversion import ModeloptStateManager, restore_from_modelopt_state from modelopt.torch.quantization.utils import set_quantizer_state_dict from modelopt.torch.utils import print_rank_0 @@ -48,8 +48,11 @@ def get_model( # Load model model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=device_map) - # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this - if hasattr(model, "peft_config"): + # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this. + # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state. + # Skipping is safe only because QATTrainer writes modelopt_state_train.pth at trainer init, + # from that same base state. + if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model): modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth") restore_from_modelopt_state(model, modelopt_state) print_rank_0("Restored modelopt state") diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..cc894d0ffd5 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -959,6 +959,14 @@ def from_quantized_weight( raise NotImplementedError(f"quantization format {quantization} not supported") +def _strip_base_layer(key: str, is_modelopt_qlora: bool) -> str: + """Drop the `base_layer` component PEFT inserts, which deployment does not expect. + + Stripping generically means new key types (bias, scales) need no enumeration here. + """ + return key.replace(".base_layer.", ".") if is_modelopt_qlora else key + + def postprocess_state_dict( state_dict: dict, maxbound: float, @@ -991,16 +999,8 @@ def postprocess_state_dict( "weight_shape", ] - # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment - if is_modelopt_qlora: - replacements.update( - { - "base_layer.weight": "weight", - "base_layer.input_scale": "input_scale", - "base_layer.weight_scale": "weight_scale", - } - ) - skip_keys.append("base_layer") + def _export_key(key: str) -> str: + return _strip_base_layer(key, is_modelopt_qlora) post_state_dict = {} @@ -1012,7 +1012,7 @@ def postprocess_state_dict( # Skip keys not related to quantizers if all(skip_key not in key for skip_key in skip_keys): - post_state_dict[key] = value + post_state_dict[_export_key(key)] = value continue # Apply replacements if the key matches any suffix in the replacements dict @@ -1033,7 +1033,7 @@ def postprocess_state_dict( logger.warning( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) - post_state_dict[prefix + new_suffix] = value + post_state_dict[_export_key(prefix + new_suffix)] = value break # Squeeze scales with a leading dimension of 1 diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index ce963c0df82..a291b5abf36 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -114,16 +114,34 @@ def _restore_qtensor_wrappers(model, model_path): q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {}) if not q_tensor_state: continue - for name, module in model.named_modules(): - if ( - isinstance(module, RealQuantLinear) - and name in q_tensor_state - and not isinstance(module.weight, QTensorWrapper) - ): - module._parameters["weight"] = QTensorWrapper( - qtensor=module.weight.data, - metadata=q_tensor_state[name]["metadata"], - ) + # PEFT nests the quantized linear as `.base_layer`, and either the saved keys or the + # live names may carry that suffix. Normalize both so the lookup works in either direction. + q_tensor_state = {k.removesuffix(".base_layer"): v for k, v in q_tensor_state.items()} + + pending = [ + (name, module) + for name, module in model.named_modules() + if isinstance(module, RealQuantLinear) and not isinstance(module.weight, QTensorWrapper) + ] + matched = 0 + for name, module in pending: + key = name.removesuffix(".base_layer") + if key not in q_tensor_state: + continue + module._parameters["weight"] = QTensorWrapper( + qtensor=module.weight.data, + metadata=q_tensor_state[key]["metadata"], + ) + matched += 1 + + # A total miss means some wrapper renamed the modules. Warn instead of letting it surface + # as an opaque shape error at dequantization. + if pending and not matched: + warnings.warn( + f"Found {len(q_tensor_state)} compressed weight(s) in {modelopt_state_path} but " + f"re-wrapped none of the {len(pending)} candidate module(s); their names may have " + "been remapped. The model will likely fail when the packed weights are used." + ) def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs): diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py index f87117501f5..4cd67d7f905 100644 --- a/tests/examples/llm_qat/test_llm_qat.py +++ b/tests/examples/llm_qat/test_llm_qat.py @@ -14,8 +14,12 @@ # limitations under the License. +import json + import pytest +import torch from _test_utils.examples.run_command import run_example_command +from safetensors.torch import load_file # Mapping from backend name to accelerate config file BACKEND_CONFIGS = { @@ -86,6 +90,18 @@ def _run_train(config: str, extra_cmd_args: list[str], backend: str = "fsdp2", c setup_free_port=True, ) + +def _run_export(ckpt_dir: str, export_dir: str): + run_example_command( + [ + "python", "export.py", + "--pyt_ckpt_path", ckpt_dir, + "--export_path", export_dir, + ], + "llm_qat", + ) + + def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path): """Test dataset_utils.py standalone CLI pre-tokenization.""" cache_dir = tmp_path / "dataset_cache" @@ -152,18 +168,43 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): ) # Step 2: LoRA QAT + lora_qat_output_dir = tmp_path / "lora_qat" _run_train( "configs/train/qat_nvfp4.yaml", [ "--model_name_or_path", str(ptq_output_dir), "--do_train", "True", "--lora", "True", - "--output_dir", str(tmp_path / "lora_qat"), + "--output_dir", str(lora_qat_output_dir), ], backend="fsdp2", cache_dir=cache_dir, ) + # Step 3: Export. This checkpoint is fake-quantized, so the calibrated amaxes rather than + # packed weights are what must survive the load. + export_dir = tmp_path / "lora_qat_export" + _run_export(str(lora_qat_output_dir), str(export_dir)) + + base_model_dir = export_dir / "base_model" + with open(base_model_dir / "hf_quant_config.json") as f: + assert json.load(f)["quantization"]["quant_algo"] == "NVFP4" + + base_weights = load_file(base_model_dir / "model.safetensors") + assert not any("base_layer" in k or k.endswith("_amax") for k in base_weights) + + # LoRA freezes the base model, so a direct PTQ export is a trusted oracle for every calibrated + # value. This catches scales that keep their key but were reset to defaults. + ptq_export_dir = tmp_path / "ptq_export" + _run_export(str(ptq_output_dir), str(ptq_export_dir)) + reference = load_file(ptq_export_dir / "model.safetensors") + + scales = [k for k in reference if k.endswith(("_scale", "_scale_2"))] + assert scales, "no NVFP4 scales in the reference PTQ export" + for key in scales: + assert key in base_weights, f"{key} missing from the LoRA-QAT export" + assert torch.equal(base_weights[key], reference[key]), f"{key} does not match PTQ export" + @pytest.mark.parametrize("backend", [ "fsdp2", @@ -219,14 +260,37 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path): ) # Step 2: QLoRA training + qlora_output_dir = tmp_path / "qlora" _run_train( "configs/train/qlora_nvfp4.yaml", [ "--model_name_or_path", str(ptq_output_dir), "--do_train", "True", "--lora", "True", - "--output_dir", str(tmp_path / "qlora"), + "--output_dir", str(qlora_output_dir), ], backend="ddp", cache_dir=cache_dir, ) + + # Step 3: Export the QLoRA checkpoint for deployment + export_dir = tmp_path / "qlora_export" + _run_export(str(qlora_output_dir), str(export_dir)) + + # The base model is exported compressed; the adapters stay at the top level. + base_model_dir = export_dir / "base_model" + assert (export_dir / "adapter_model.safetensors").is_file() + assert (base_model_dir / "hf_quant_config.json").is_file() + + with open(base_model_dir / "hf_quant_config.json") as f: + assert json.load(f)["quantization"]["quant_algo"] == "NVFP4" + + # NVFP4 needs the packed weight and *both* scales to be dequantizable downstream. + base_weights = load_file(base_model_dir / "model.safetensors") + packed_weights = [k for k, v in base_weights.items() if k.endswith(".weight") and v.dtype == torch.uint8] + assert packed_weights, "no NVFP4-packed weights found in the exported base model" + for key in packed_weights: + prefix = key.removesuffix(".weight") + assert f"{prefix}.weight_scale" in base_weights + assert f"{prefix}.weight_scale_2" in base_weights + assert not any("base_layer" in k or "lora" in k for k in base_weights) diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index cac0a9a9aef..55137a64639 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -259,6 +259,41 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta assert processed_state_dict == expected_state_dict +def test_postprocess_state_dict_qlora_strips_base_layer(): + """Every QLoRA `base_layer.*` tensor needed for deployment must survive the rename. + + Dropping the NVFP4 global scale or a bias yields an undeployable checkpoint. + """ + state_dict = { + "layer1.base_layer.weight": torch.ones(4, 2, dtype=torch.uint8), + "layer1.base_layer.weight_scale": torch.ones(4, 1), + "layer1.base_layer.weight_scale_2": torch.tensor([0.5]), + "layer1.base_layer.input_scale": torch.tensor([0.25]), + "layer1.base_layer.bias": torch.arange(4.0), + "layer1.base_layer.input_quantizer._pre_quant_scale": torch.ones(2), + # Quantizer internals must still be dropped. + "layer1.base_layer.weight_quantizer._amax": torch.tensor([1.0]), + "layer1.base_layer.input_quantizer._amax": torch.tensor([1.0]), + "layer1.base_layer.weight_quantizer._scale": torch.ones(4, 1), + "layer1.base_layer.weight_quantizer._double_scale": torch.tensor([0.5]), + } + + processed_state_dict = postprocess_state_dict( + state_dict, 448.0, QUANTIZATION_NONE, is_modelopt_qlora=True + ) + + assert set(processed_state_dict) == { + "layer1.weight", + "layer1.weight_scale", + "layer1.weight_scale_2", + "layer1.input_scale", + "layer1.bias", + "layer1.pre_quant_scale", + } + assert torch.equal(processed_state_dict["layer1.weight_scale_2"], torch.tensor([0.5])) + assert torch.equal(processed_state_dict["layer1.bias"], torch.arange(4.0)) + + @pytest.mark.parametrize( ("config", "expected"), [ diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py index 8a44ad23c76..0476f44c56c 100644 --- a/tests/unit/torch/opt/plugins/test_hf_patching.py +++ b/tests/unit/torch/opt/plugins/test_hf_patching.py @@ -14,6 +14,8 @@ # limitations under the License. import pytest +import torch +import torch.nn as nn from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, get_tiny_qwen3, @@ -23,6 +25,9 @@ import modelopt.torch.distill as mtd import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.transformers import _restore_qtensor_wrappers +from modelopt.torch.quantization.qtensor import QTensorWrapper @pytest.mark.parametrize( @@ -54,3 +59,61 @@ def test_nested_model_save_restore(tmp_path, model_cls, teacher_model_type): tf_output_tester(model, model_test) # KD state is not saved and it should be empty assert not mto.ModeloptStateManager(model_test).has_state + + +class _LoraLike(nn.Module): + """Stand-in for peft's `lora.Linear`, which nests the original module under `base_layer`.""" + + def __init__(self, base_layer): + super().__init__() + self.base_layer = base_layer + + +def _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer=False): + model = nn.Sequential() + model.fc = nn.Linear(64, 32) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(2, 64))) + mtq.compress(model) + assert isinstance(model.fc.weight, QTensorWrapper) + + state = mto.modelopt_state(model) + if state_keyed_with_base_layer: + # Compressing after the adapters are attached saves the keys with the peft suffix. + for _, mode_config in state["modelopt_state_dict"]: + q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {}) + for key in list(q_tensor_state): + q_tensor_state[f"{key}.base_layer"] = q_tensor_state.pop(key) + torch.save(state, tmp_path / "modelopt_state.pth") + + # transformers>=5 assigns a plain Parameter holding the packed data, dropping the wrapper. + packed = model.fc.weight.data.clone() + del model.fc._parameters["weight"] + model.fc._parameters["weight"] = nn.Parameter(packed, requires_grad=False) + assert not isinstance(model.fc.weight, QTensorWrapper) + return model + + +@pytest.mark.parametrize("wrap_in_lora", [False, True]) +@pytest.mark.parametrize("state_keyed_with_base_layer", [False, True]) +def test_restore_qtensor_wrappers(tmp_path, wrap_in_lora, state_keyed_with_base_layer): + """Either side may carry the `.base_layer` suffix, so the lookup must work in both directions.""" + model = _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer) + if wrap_in_lora: + model.fc = _LoraLike(model.fc) + + _restore_qtensor_wrappers(model, str(tmp_path)) + + linear = model.fc.base_layer if wrap_in_lora else model.fc + assert isinstance(linear.weight, QTensorWrapper) + assert linear.weight.metadata["shape"] == torch.Size([32, 64]) + + +def test_restore_qtensor_wrappers_warns_when_nothing_matches(tmp_path): + """A total miss must be loud -- it otherwise surfaces as an opaque shape error at dequant.""" + model = _compressed_model_and_state_dir(tmp_path) + model.fc = _LoraLike(_LoraLike(model.fc)) # a nesting the lookup does not know about + + with pytest.warns(UserWarning, match="re-wrapped none"): + _restore_qtensor_wrappers(model, str(tmp_path)) + + assert not isinstance(model.fc.base_layer.base_layer.weight, QTensorWrapper) From 903b95dc78cadca232c4f1f0211216e75c30cd02 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:05:56 +0530 Subject: [PATCH 10/11] Minitron pruning fixes for Nemotron-3.5-Lightning-30B-A3B and Deepseek (#2159) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix + new feature Two model families that could not be pruned end-to-end now can: - **Nemotron-3.5-Lightning** (`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`) — a native `NemotronHForCausalLM` that ships without remote code and carries MTP heads. Fixes a calibration crash and HF-export failures on the modern Megatron-Bridge / transformers stack. - **DeepSeek-V3** — fixes an MLA Q-LoRA crash during calibration, and adds a `candidate_filter` search option to `mcore_minitron` so its MoE-FFN dimensions stay prunable while remaining representable in HF. Also makes a rank-local failure under pipeline parallelism fail fast instead of stalling. #### 1. Nemotron Lightning: prune + HF export (`examples/megatron_bridge/prune_minitron.py`) 1. **MTP calibration crash.** On newer Megatron-LM, `mtp_process` is derived from the hybrid *pattern*, not from `mtp_num_layers`. Setting only `mtp_num_layers=0` in the calibration provider overrides was insufficient: the provider's `finalize()` re-appended the MTP suffix to `hybrid_layer_pattern` (because `mtp_hybrid_override_pattern` was still set and `mtp_use_repeated_layer=True`), so `mtp_process=True` while `mtp_num_layers=0` and the calibration forward hit `assert self.config.mtp_num_layers > 0`. Fix: also clear `mtp_hybrid_override_pattern` in the calibration overrides so MTP is fully disabled (MTP heads are dropped from the pruned model, as before). 2. **HF export via a config-only bridge (hybrid models only).** The old export built a *dummy* HF model to obtain the bridge, then streamed weights. This breaks on native NemotronH because (a) native `NemotronHConfig` makes `hybrid_override_pattern` a read-only property, and (b) transformers 5.12 saves the input embedding under a different key than the bridge mapping expects (`backbone.embedding` vs `backbone.embeddings`); the mismatch made `build_conversion_tasks` drop the embedding task on its owning rank, leaving an owner-less PP placeholder that crashed `save_hf_weights` with `Object must exist on at least one PP rank`. Fix: stream weights through a **config-only** bridge (`AutoBridge.from_hf_config(hf_cfg).save_hf_pretrained(...)`), available since Megatron-Bridge 0.5.0 (nemo:26.06). A config-only bridge has `hf_keys=None`, so the embedding task is never dropped, no dummy model is built, and the output uses the canonical HF key names. This is **restricted to hybrid providers**, which are the only models that need it; non-hybrids keep the dummy-model path that CI has always exercised. Writing the source artifacts is now rank-0-only. Every rank used to write the source `config.json`, which races with the pruned `config.json` that `save_hf_pretrained` writes from rank 0 alone: a late write from another rank leaves a checkpoint whose config does not match its weights. This reproduced intermittently on both Qwen3 and NemotronH before the fix, and 3/3 clean runs after. `save_hf_pretrained` takes no `trust_remote_code` argument — it reads the flag **off the bridge** to fetch the source checkpoint's artifacts, and `from_hf_config` cannot infer it because `AutoConfig.from_pretrained` consumes the kwarg rather than storing it on the config. So the flag is set explicitly on the bridge instance; otherwise remote-code models would silently lose it. 3. **Config write-back correctness:** - `hybrid_override_pattern` is only written for older remote-code configs that lack `layer_types`; native configs carry the cadence in `layer_types` (read-only `hybrid_override_pattern` is skipped). - `n_shared_experts` is preserved (a fixed count) instead of being re-derived by `moe_shared_expert_intermediate_size // moe_ffn_hidden_size`, which is DeepSeek-style logic that would corrupt NemotronH's count. Non-hybrids, VLMs, and Megatron-Bridge builds without config-only export keep the dummy-model path, with a `warn_rank_0` when a hybrid has to fall back. The README's `transformers<5` workaround is **removed**: it existed because the dummy-model path broke on transformers 5, and the config-only path handles NemotronH on every supported container. #### 2. `candidate_filter` for `mcore_minitron` (`modelopt/torch/prune/plugins/mcore_minitron.py`) DeepSeek-style MoE configs have no explicit shared-expert-size field: they size the shared expert as `n_shared_experts * moe_intermediate_size`, where `moe_intermediate_size` is the (also prunable) **routed** expert size. So only candidates with `moe_shared_expert_intermediate_size % moe_ffn_hidden_size == 0` can be written back to HF at all. Candidates come from a Cartesian `product()` of independent per-hparam choice lists, so no per-hparam restriction can express a constraint *between* two hparams. New optional `candidate_filter` search-config key (default `None`, so existing behaviour is unchanged): a callable that rejects candidate configs before the metric computation, making the search cheaper rather than more expensive. It receives **every** supported hparam, with non-searched ones filled in from the model config, so a filter still works when one of its hparams was skipped or had a single choice. Rejected candidates are not cached, so — like `score_func`, whose cached scores are reused without re-validation — the filter is assumed unchanged when resuming from a `checkpoint`. `prune_minitron.py` wires this up for DeepSeek-style configs, so **both** `moe_ffn_hidden_size` and `moe_shared_expert_intermediate_size` stay prunable (the search then only picks shared sizes that are a multiple of the routed one). A `--prune_export_config` that violates the constraint never reaches the filter, so the export path now raises `ValueError` instead of writing a checkpoint whose config disagrees with its weights. #### 3. MLA Q-LoRA pruning (`modelopt/torch/prune/plugins/mcore_minitron.py`) Pruning any MLA model with `q_lora_rank` set died during calibration with `AttributeError: 'tuple' object has no attribute 'view'`. `hidden_size` importance estimation blanket-patches every `TELayerNormColumnParallelLinear` with `return_layernorm_output=True` to capture post-layernorm activations. When `q_lora_rank` is set, MCore builds `linear_q_up_proj` as a `TELayerNormColumnParallelLinear` — the Q-LoRA layernorm is fused into it, which is why `q_layernorm` is `IdentityOp` — so it was patched too, even though its layernorm is over the **latent rank**, not `hidden_size`. TE then returns `((out, ln_out), bias)` and MCore's `q, _ = self.linear_q_up_proj(...)` leaves `q` a tuple. Isolated by probing the module before and after dynamic conversion: | Setup | `linear_q_up_proj` returns | Forward | | --- | --- | --- | | Before conversion | `tuple(Tensor, NoneType)` | — | | After conversion, no hooks | `tuple(Tensor, NoneType)` | OK | | After conversion **+ importance hooks** | `tuple(tuple(Tensor, Tensor), NoneType)` | AttributeError | So conversion is innocent; registering the importance hooks is the trigger. Fix: exclude MLA's Q/KV up-projections from both the patch and unpatch loops. `test_mcore_mla_pruning` did not catch this because it builds MLA without `q_lora_rank`, where MCore uses a plain `linear_q_proj` and nothing is patched. #### 4. Fail fast instead of stalling on a rank-local error under PP (`modelopt/torch/utils/distributed.py`) A rank raising inside a distributed entrypoint left the whole job stalled until the process group timed out, with **no diagnostic output at all**: the failing rank blocked in `cleanup()`'s barrier while its peers blocked in `recv_from_prev_pipeline_rank_`, and Python only prints a traceback once the enclosing `finally` returns. A crash on one rank was indistinguishable from a slow job. - `dist.cleanup()` skips the barrier when unwinding from an exception. - New `dist.abort()` prints the traceback, flushes and exits immediately. Skipping the barrier alone is **not** enough — a stack dump showed the failing rank then blocking in `destroy_process_group` for the same reason — so the error path must not tear the process group down at all. `SystemExit` is re-raised rather than aborted, so an intentional exit (e.g. the `--score_lower_bound` gate) keeps its exit code and prints no traceback. Kept out of `cleanup()` so no library caller gets a surprise process exit. - Called from the entrypoints that wrap `main()` in `try/finally`: the five `examples/megatron_bridge` scripts. Measured on a 2-GPU PP run whose rank 0 raises during calibration: **10 min timeout kill with no visible error → 31s, exit 1, real traceback.** This is a latent, pre-existing issue (the `try/finally` predates this PR); it only surfaces on a failing PP run, which is why CI never hit it. ### Usage ```bash torchrun --nproc_per_node 4 examples/megatron_bridge/prune_minitron.py \ --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 \ --pp_size 4 \ --prune_target_active_params 3e9 \ --output_hf_path /path/to/Nemotron-3.5-Lightning-30B-A3B-Pruned-A3.0B ``` ### Testing - **End-to-end on nemo:26.08.rc6** (4× GB300, transformers 5.12.1, Megatron-Bridge with config-only export): pruning + export complete (`EXIT=0`, "Saved pruned model … Done!"). The exported checkpoint has canonical **plural** `backbone.embeddings.weight` keys, **0 MTP tensors**, and a config that reloads correctly (`num_hidden_layers=52` from `layers_block_type`, `n_shared_experts=1`, `num_nextn_predict_layers=0`, pruned `hidden_size`/`mamba_*`/MoE dims, reconstructed `hybrid_override_pattern`).
Pruning search log (--prune_target_active_params 3e9) ```text Top 10 Candidates with Scores ┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━┓ ┃ # ┃ export_config ┃ active_params ┃ params ┃ score ┃ ┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━┩ │ 1 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.49B │ 0.5406 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3584} │ │ │ │ │ 2 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 56, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2427 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 3 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 21.61B │ 0.2643 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 4 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 19.28B │ 0.4552 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3712} │ │ │ │ │ 5 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 104, │ 3.00B │ 22.28B │ 0.5860 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 6 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 21.99B │ 0.2294 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3328} │ │ │ │ │ 7 │ {'num_layers': 48, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.5231 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 8 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 96, │ 3.00B │ 21.81B │ 0.5042 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3584} │ │ │ │ │ 9 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2462 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ │ 10 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 20.70B │ 0.5685 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ └────┴───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┴───────────────┴────────┴────────┘ ╭──────────────────────────────────────────────────────────────────────── Best Subnet ─────────────────────────────────────────────────────────────────────────╮ │ export_config {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 104, 'moe_ffn_hidden_size': 1856, │ │ 'moe_shared_expert_intermediate_size': 3072} │ │ active_params 3.00B │ │ params 22.28B │ │ score 0.5860 │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭────────────────────────────────────────────────────── Pruned Model Stats ───────────────────────────────────────────────────────╮ │ Total Parameters 22.28B │ │ Active Parameters 3.00B │ │ Memory (BF16, seq_length=8192, batch_size=8) weights: 42489.7 MB, kv_cache: 384.0 MB, mamba_state: 190.5 MB, Total: 43064.2 MB │ ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ```
- **`tests/examples/megatron_bridge/test_prune_minitron.py`** — `nemotron_h` now exports to HF and reloads (previously it stopped at a Megatron checkpoint, since the dummy-model path needed `transformers<5`), plus an `n_shared_experts` config assertion; the dead `megatron_format` branch is gone. It runs on the CI container: **verified on nemo:26.06.01 (transformers 5.8.1) and nemo:26.08.rc6.** - **`tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py`** — the `nas_memory_mb` search test now passes a `candidate_filter` and asserts the exact number of rejected candidates (256 of the 512-combo grid) plus the surviving candidates' validity; its `expected_top_k` goldens are regenerated accordingly. Because `moe_shared_expert_intermediate_size` is in that test's skip list, this also covers the model-config fallback for hparams that are not in the search space. Verified on 2 GPUs, on both the CI container (nemo:26.06.01) and nemo:26.08.rc6: | Test | Result | | --- | --- | | `test_prune_minitron[qwen3]` | PASSED on 26.06.01 and 26.08.rc6 | | `test_prune_minitron[deepseek_v3]` | PASSED (52s) — MLA Q-LoRA + `candidate_filter` end-to-end | | `test_prune_minitron[nemotron_h]` | PASSED on 26.06.01 (58s) and 26.08.rc6 (61s) | | `test_mcore_mamba_hybrid_pruning_nas_memory_mb` | PASSED | | `test_mcore_mamba_hybrid_pruning_nas_params` | PASSED (unchanged sibling, run to check the regenerated goldens did not disturb it) | | 2-GPU PP run failing on rank 0 | fails in 31s with a real traceback (was a 10 min stall) | ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — `candidate_filter` defaults to `None` (existing searches unchanged), and the config-only export is limited to hybrid providers on nemo:26.08+, so dense / MoE / VLM exports keep the path they use today. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: ✅ ### Additional Information Enables the Prune + Distill workflow for `NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` (native, no-remote-code `NemotronHForCausalLM` with MTP heads). Pruning-time MTP support was scoped and intentionally deferred — MTP heads are dropped and can be re-derived via a short SFT with `mtp_num_layers=1` on the pruned+distilled model. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/megatron_bridge/README.md | 3 - examples/megatron_bridge/distill.py | 2 + .../export_distilled_megatron_to_hf.py | 2 + .../export_quantized_megatron_to_hf.py | 2 + examples/megatron_bridge/prune_minitron.py | 106 ++++++++++++++---- examples/megatron_bridge/quantize.py | 2 + .../torch/prune/plugins/mcore_minitron.py | 55 +++++++-- modelopt/torch/utils/distributed.py | 34 +++++- .../megatron_bridge/test_prune_minitron.py | 53 +++++---- .../test_mcore_mamba_minitron_pruning.py | 32 +++++- 10 files changed, 228 insertions(+), 63 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 2fae7fe3545..030edaf8fd5 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -365,9 +365,6 @@ torchrun --nproc_per_node 1 prune_minitron.py --help > [!NOTE] > Multi-token-prediction (MTP) heads (e.g. Qwen3.5) are not pruned yet — they are dropped for the prune run and the saved checkpoint has no MTP. Autoregressive inference is unaffected; for speculative decoding, run a short MTP SFT on the pruned model. -> [!NOTE] -> If pruning a Nemotron model and you want to save the pruned model back in HF format, please downgrade to `transformers<5` via `python -m pip install "transformers<5"` before pruning. - ### Vision-Language Models (VLMs) For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `prune_minitron.py` automatically prunes only the **language model** and leaves the vision tower intact, then saves the full VLM back. All the pruning modes above (parameter count, active parameter count, memory footprint, and manual `export_config`) work unchanged, with two VLM-specific caveats: diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 16dd37d5f8d..b35369c40f3 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -480,5 +480,7 @@ def _restore_student_hook(model_chunks): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index 1ba76ea7ae4..d5e95ff9d18 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -288,5 +288,7 @@ def main(args: argparse.Namespace): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 17db5e6da34..e4e3703d8a5 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -164,5 +164,7 @@ def main(args: argparse.Namespace): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 6f2b7c3829f..7a9b4b8bd56 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -105,6 +105,24 @@ def _hf_config_has_mtp(hf_cfg) -> bool: ) +# HF names the shared expert size with or without the ``moe_`` prefix depending on the model +# (e.g. Qwen3.5-MoE uses ``shared_expert_intermediate_size``). +_SHARED_EXPERT_SIZE_FIELDS = ( + "moe_shared_expert_intermediate_size", + "shared_expert_intermediate_size", +) + + +def _is_deepseek_style_moe(text_cfg) -> bool: + """Whether the shared expert is sized as ``n_shared_experts * moe_intermediate_size``. + + Such configs can only represent a shared expert size that is a multiple of the routed one. + """ + return getattr(text_cfg, "n_shared_experts", None) is not None and not any( + hasattr(text_cfg, field) for field in _SHARED_EXPERT_SIZE_FIELDS + ) + + def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--hf_model_name_or_path", type=str, required=True) @@ -400,7 +418,9 @@ 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 + # MTP is not supported during calibration; drop it + "mtp_num_layers": 0, + "mtp_hybrid_override_pattern": None, }, init_model_parallel=True, moe_grouped_gemm=not args.no_moe_grouped_gemm, @@ -542,6 +562,17 @@ def score_func(m): pruning_config["max_width_pruning"] = args.max_width_pruning pruning_config["max_depth_pruning"] = args.max_depth_pruning pruning_config["hparams_to_skip"] = args.hparams_to_skip + # DeepSeek-style MoE configs size the shared expert as n_shared_experts * moe_intermediate_size, + # so only candidates whose shared size is a multiple of the routed one can be saved to HF. + src_hf_cfg = bridge.hf_pretrained.config + if _is_deepseek_style_moe(getattr(src_hf_cfg, "text_config", src_hf_cfg)): + warn_rank_0( + "DeepSeek-style MoE config detected: restricting the search to candidates whose " + "moe_shared_expert_intermediate_size is a multiple of moe_ffn_hidden_size." + ) + pruning_config["candidate_filter"] = lambda cfg: ( + cfg["moe_shared_expert_intermediate_size"] % cfg["moe_ffn_hidden_size"] == 0 + ) pruning_config["top_k"] = args.top_k # memory_mb constraint requires batch_size and seq_length pruning_config["batch_size"] = args.inference_batch_size @@ -588,10 +619,11 @@ def score_func(m): else: print_rank_0(f"Saving pruned model to {args.output_hf_path} in HF checkpoint format") - # [WAR] Save the pruned HF model by hand until Megatron-Bridge natively supports it. - # TODO: Replace this whole block with ``AutoBridge.from_auto_config(...).save_hf_weights(...)`` - # once the Megatron-Bridge fix ships (nemo:26.08). - bridge.hf_pretrained.save_artifacts(args.output_hf_path) + # Build the pruned HF config field-by-field from the pruned Megatron config, then stream weights. + # Rank 0 only: a late write from another rank would leave config.json stale. + if dist.is_master(): + bridge.hf_pretrained.save_artifacts(args.output_hf_path) + dist.barrier() hf_cfg = AutoConfig.from_pretrained( args.output_hf_path, trust_remote_code=args.trust_remote_code ) @@ -610,12 +642,7 @@ 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 - # 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", - ): + for shared_expert_field in _SHARED_EXPERT_SIZE_FIELDS: if hasattr(text_cfg, shared_expert_field): setattr( text_cfg, shared_expert_field, mcore_cfg.moe_shared_expert_intermediate_size @@ -624,7 +651,16 @@ def score_func(m): 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"): + # n_shared_experts is a fixed count; only DeepSeek-style configs record the pruned shared + # expert size through it. candidate_filter keeps the search divisible, so only a manual + # --prune_export_config can violate this. + if _is_deepseek_style_moe(text_cfg): + if mcore_cfg.moe_shared_expert_intermediate_size % mcore_cfg.moe_ffn_hidden_size: + raise ValueError( + f"{mcore_cfg.moe_shared_expert_intermediate_size=} must be a multiple of " + f"{mcore_cfg.moe_ffn_hidden_size=} for this config, which stores the shared " + "expert size as n_shared_experts * moe_intermediate_size. " + ) text_cfg.n_shared_experts = ( mcore_cfg.moe_shared_expert_intermediate_size // mcore_cfg.moe_ffn_hidden_size ) @@ -658,8 +694,11 @@ def score_func(m): "distillation cannot recover this vision-path change -- consider full VLM " "training/distillation instead of LM-only to recover vision quality." ) - if isinstance(provider, _HYBRID_PROVIDER_TYPES) and hasattr( - text_cfg, "hybrid_override_pattern" + # Only older remote-code configs need this; native configs carry the cadence in layer_types. + if ( + isinstance(provider, _HYBRID_PROVIDER_TYPES) + and not hasattr(text_cfg, "layer_types") + and hasattr(text_cfg, "hybrid_override_pattern") ): # MCore's pattern can carry an MTP suffix (``/...``) and PP boundaries (``|``) which we need to remove text_cfg.hybrid_override_pattern = "".join( @@ -671,15 +710,36 @@ def score_func(m): 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 - 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( - args.output_hf_path, trust_remote_code=args.trust_remote_code + # Config-only bridge (hf_keys=None) keeps the embedding task when transformers' saved key + # differs from the bridge mapping (NemotronH's backbone.embedding vs ...embeddings). + use_config_only_export = ( + hasattr(AutoBridge, "from_hf_config") + and isinstance(provider, _HYBRID_PROVIDER_TYPES) + and not is_vlm ) - pruned_bridge.save_hf_weights(model, args.output_hf_path) + if use_config_only_export: + pruned_bridge = AutoBridge.from_hf_config(hf_cfg) + # save_hf_pretrained reads trust_remote_code off the bridge to fetch source artifacts; + # from_hf_config can't infer it since AutoConfig consumes the kwarg. + pruned_bridge.trust_remote_code = args.trust_remote_code + pruned_bridge.save_hf_pretrained( + model, args.output_hf_path, source_path=args.hf_model_name_or_path + ) + else: + if isinstance(provider, _HYBRID_PROVIDER_TYPES) and not is_vlm: + warn_rank_0( + "Megatron-Bridge lacks config-only HF export; falling back to the dummy-model " + "path, which cannot round-trip a pruned native NemotronH config. Use " + "transformers<5 or a newer Megatron-Bridge if the save fails." + ) + 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( + args.output_hf_path, trust_remote_code=args.trust_remote_code + ) + pruned_bridge.save_hf_weights(model, args.output_hf_path) copy_hf_ckpt_remote_code(args.hf_model_name_or_path, args.output_hf_path) print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format") @@ -704,5 +764,7 @@ def score_func(m): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 6355e60e435..1e4ee79f574 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -454,5 +454,7 @@ def forward_loop(_model=None): args = get_args() try: main(args) + except BaseException: + dist.abort() # peers may be stuck in a collective this rank will never reach finally: dist.cleanup() diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 28684fa3a4a..ad12fc7ebff 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -47,6 +47,7 @@ gather_from_tensor_model_parallel_region, reduce_from_tensor_model_parallel_region, ) +from megatron.core.transformer.multi_latent_attention import MLASelfAttention from pydantic import create_model from rich.console import Console from rich.markup import escape as rich_escape @@ -255,6 +256,11 @@ 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). + - `candidate_filter`: Callable rejecting candidate configs the caller cannot use, e.g. ones + their checkpoint format cannot represent (default: None). Receives every supported hparam, + with non-searched ones filled in from the model config (unset ones are omitted, so a + filter using them raises `KeyError`). Like `score_func`, it is assumed + unchanged when resuming from a `checkpoint`, since rejected candidates are not cached. - `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. @@ -280,6 +286,7 @@ def default_search_config(self) -> SearchConfig: "max_width_pruning": 0.40, "max_depth_pruning": 0.20, "hparams_to_skip": None, + "candidate_filter": None, "top_k": 10, # Memory footprint config (only used with memory_mb constraint) "seq_length": 4096, @@ -521,6 +528,7 @@ def search_best_arch_by_metrics(self) -> dict: max_width_pruning = self.config["max_width_pruning"] max_depth_pruning = self.config["max_depth_pruning"] hparams_to_skip = self.config["hparams_to_skip"] + candidate_filter = self.config["candidate_filter"] top_k = self.config["top_k"] constraints_str = ", ".join(f"{self._fmt_metric(v, k)} {k}" for k, v in max_metrics.items()) print_rank_0(f"\nSearching for the best pruned architecture under {constraints_str}...") @@ -549,12 +557,24 @@ def search_best_arch_by_metrics(self) -> dict: max_depth_pruning, hparams_to_skip, ) + # Only place to reject invalid hparam *combinations*; unsearched ones come from config. + base_config = { + hp: getattr(self.model.config, hp) + for hp in SUPPORTED_HPARAMS + if getattr(self.model.config, hp, None) is not None + } selected = [] + num_filtered = 0 for ss_config in tqdm( search_space_configs, desc="Finding all candidates fitting the constraints...", disable=not dist.is_master(), ): + if candidate_filter is not None and not candidate_filter( + {**base_config, **ss_config} + ): + num_filtered += 1 + continue candidate_metrics = self._compute_candidate_metrics(ss_config, max_num_layers) if all(candidate_metrics[k] <= max_metrics[k] for k in active_metric_keys): selected.append( @@ -562,7 +582,11 @@ def search_best_arch_by_metrics(self) -> dict: ss_config, {k: candidate_metrics[k] for k in active_metric_keys}, None ) ) - assert len(selected) > 0, "No subnets found fitting the constraints!" + if num_filtered: + print_rank_0(f"Rejected {num_filtered} candidates via candidate_filter.") + assert len(selected) > 0, "No subnets found fitting the constraints!" + ( + f" candidate_filter rejected all {num_filtered} candidates." if num_filtered else "" + ) print_rank_0(f"Found {len(selected)} candidates fitting the constraints!") self.all_candidates_per_constraint[constraints_cache_key] = sorted( selected, key=lambda x: x.metrics[primary_key], reverse=True @@ -952,6 +976,25 @@ def restore(self) -> RestoreEntrypoint: return restore_mcore_minitron +def _fused_ln_linears_over_hidden_size(module: nn.Module): + """Yield fused-layernorm linears whose layernorm is over ``hidden_size``. + + MLA's Q/KV up-projections are ``TELayerNormColumnParallelLinear`` too, but their layernorm is + over the latent rank and MCore unpacks their output as ``(out, bias)``. Setting + ``return_layernorm_output`` there makes it ``((out, ln_out), bias)`` and breaks the forward. + """ + up_projs = { + id(sub) + for m in module.modules() + if isinstance(m, MLASelfAttention) + for attr in ("linear_q_up_proj", "linear_kv_up_proj") + if (sub := getattr(m, attr, None)) is not None + } + for m in module.modules(): + if isinstance(m, TELayerNormColumnParallelLinear) and id(m) not in up_projs: + yield m + + class ImportanceEstimatorRegistry: """Register importance estimators and forward hooks for all supported modules in the model. @@ -1034,9 +1077,8 @@ def cleanup(self) -> None: self._hooks.clear() # Unpatch return_layernorm_output on fused TELayerNormColumnParallelLinear modules - for m in self.model.modules(): - if isinstance(m, TELayerNormColumnParallelLinear): - m.return_layernorm_output = False + for m in _fused_ln_linears_over_hidden_size(self.model): + m.return_layernorm_output = False def get_layer_scores(self) -> dict[int, torch.Tensor]: """Get the layer scores (1-indexed) from the model. @@ -1166,9 +1208,8 @@ def _estimate_hidden_size_importance(mod): # Layernorms are fused into TELayerNormColumnParallelLinear. We temporarily # patch return_layernorm_output=True so TE's fused kernel returns the layernorm output. # For MoE layers, pre_mlp_layernorm is a separate TENorm — use a regular forward hook. - for m in module.modules(): - if isinstance(m, TELayerNormColumnParallelLinear): - m.return_layernorm_output = True + for m in _fused_ln_linears_over_hidden_size(module): + m.return_layernorm_output = True for layer in module.decoder.layers: if isinstance(layer, _DynamicTransformerLayer): diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 12287865b7f..245ca81de79 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -18,7 +18,9 @@ import functools import io import os +import sys import time +import traceback from collections.abc import Callable from contextlib import suppress from datetime import timedelta @@ -212,13 +214,39 @@ def setup(timeout: timedelta | None = None): def cleanup(): - """Cleans up the distributed environment.""" + """Cleans up the distributed environment. + + The barrier is skipped when unwinding from an error, since peers may be blocked in a collective + this rank will never reach. ``SystemExit`` is treated as a clean exit (every rank reaches it). + That is not sufficient on its own -- ``destroy_process_group`` below blocks for the same reason + -- so error paths must call :func:`abort` before reaching this ``finally``. + """ if is_initialized(): - with suppress(Exception): - barrier() + exc = sys.exc_info()[1] + if exc is None or isinstance(exc, SystemExit): + with suppress(Exception): + barrier() torch.distributed.destroy_process_group() +def abort(exit_code: int = 1) -> None: + """Print the active exception and exit this rank immediately. + + Call from an ``except`` block in a distributed entrypoint. Both a barrier and + ``destroy_process_group`` stall when peers are blocked in a collective this rank will never + reach, and the traceback only prints once the enclosing ``finally`` returns -- so the run looks + hung rather than failed. Exiting lets the launcher (e.g. torchrun) terminate the peers. + ``SystemExit`` is re-raised instead, since every rank reaches an intentional exit. + """ + exc = sys.exc_info()[1] + if isinstance(exc, SystemExit): + raise exc + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(exit_code) + + def is_fsdp2_model(model) -> bool: """Return True if any submodule of ``model`` has been wrapped with FSDP2 ``fully_shard``.""" return any(isinstance(m, FSDPModule) for m in model.modules()) diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 32464814fef..78b30637468 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -19,6 +19,7 @@ from _test_utils.examples.megatron_bridge import qwen35_moe_bridge_supported from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.transformers_models import ( + create_tiny_deepseek_v3_dir, create_tiny_gemma3vl_dir, create_tiny_nemotron_h_dir, create_tiny_qwen3_5_moe_vl_dir, @@ -28,20 +29,18 @@ @pytest.mark.parametrize( - ("create_teacher", "megatron_format"), + ("create_teacher", "expected_pruned_config"), [ - # Dense Qwen3 LM, exported back to HF (reloadable to verify the pruned param count). + # Dense Qwen3 LM. 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. - # MTP heads are enabled so the run covers dropping them during calibration and the - # hybrid pattern MCore builds for them. + # NemotronH (nemotron-3-nano): Mamba + attention + MoE hybrid. + # MTP heads are enabled so the run covers dropping them during calibration. pytest.param( lambda tmp_path, num_gpus: create_tiny_nemotron_h_dir( tmp_path, @@ -50,26 +49,39 @@ num_nextn_predict_layers=1, mtp_hybrid_override_pattern="*E", ), - True, + {"n_shared_experts": 1}, id="nemotron_h", ), + # DeepSeek-V3: MLA (Q-LoRA) + MoE sizing the shared expert as + # n_shared_experts * moe_intermediate_size, so it covers candidate_filter end-to-end: + # without the filter the search picks a shared size that is not a multiple of the routed + # one and the reload below fails on the resulting shape mismatch. + pytest.param( + # n_group=1 so num_moe_experts stays divisible by it after expert pruning. + lambda tmp_path, num_gpus: create_tiny_deepseek_v3_dir( + tmp_path, + with_tokenizer=True, + return_model=True, + num_hidden_layers=num_gpus, + n_group=1, + topk_group=1, + ), + {"n_shared_experts": 1}, + id="deepseek_v3", + ), ], ) -def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): +def test_prune_minitron(tmp_path, num_gpus, create_teacher, expected_pruned_config): 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_path = tmp_path / "pruned" - output_kwarg = ( - {"output_megatron_path": pruned_path} - if megatron_format - else {"output_hf_path": pruned_path} - ) # TODO: Dont enable grouped GEMM for MoE models until nemo:26.08 container prune_command_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py", "--no_moe_grouped_gemm"], hf_model_name_or_path=teacher_hf_path, + output_hf_path=pruned_path, pp_size=num_gpus, calib_dataset_name="cnn_dailymail", calib_num_samples=8, @@ -80,17 +92,14 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): ss_channel_divisor=4, hparams_to_skip="num_attention_heads", top_k=1, - **output_kwarg, ) run_example_command(prune_command_parts, example_path="megatron_bridge") - if megatron_format: - # HF reload of a pruned NemotronH needs transformers<5; just verify the Megatron checkpoint. - assert (pruned_path / "latest_checkpointed_iteration.txt").exists() - else: - assert (pruned_path / "config.json").exists() - pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path) - assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params + 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 + for field, expected in expected_pruned_config.items(): + assert getattr(pruned_model.config, field) == expected @pytest.mark.parametrize( diff --git a/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py b/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py index 69d8c7c31ec..79f33b806f9 100644 --- a/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py +++ b/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py @@ -446,13 +446,32 @@ def _test_mcore_mamba_hybrid_pruning_nas_memory_mb(rank, size, ckpt_dir): ) memory_threshold = baseline_memory_mb * 0.7 + def _shared_is_multiple_of_routed(ss_cfg): + return ss_cfg["moe_shared_expert_intermediate_size"] % ss_cfg["moe_ffn_hidden_size"] == 0 + constraints = {"memory_mb": memory_threshold} config = { **_base_nas_config(ckpt_dir), "seq_length": sequence_length, "batch_size": 1, + # moe_shared_expert_intermediate_size is skipped, so the filter reads it from the model + # config: only routed sizes that divide it survive. + "candidate_filter": _shared_is_multiple_of_routed, } - model, searcher_state = prune_minitron(model, constraints, config, _NAS_CHANNEL_DIVISOR) + stdout_capture = io.StringIO() + with contextlib.redirect_stdout(stdout_capture): + model, searcher_state = prune_minitron(model, constraints, config, _NAS_CHANNEL_DIVISOR) + + shared_size = _NAS_MODEL_KWARGS["moe_shared_expert_intermediate_size"] + (candidates,) = searcher_state["all_candidates_per_constraint"].values() + assert all(shared_size % c.ss_config["moe_ffn_hidden_size"] == 0 for c in candidates) + if rank == 0: + # Half of the 512-combo search space: moe_ffn_hidden_size is [12, 16] and only 16 divides the + # (skipped, so unpruned) moe_shared_expert_intermediate_size of 16. + output = stdout_capture.getvalue() + match = re.search(r"Rejected (\d+) candidates", output) + assert match, f"candidate_filter rejected nothing:\n{output}" + assert int(match.group(1)) == 256, output pruned_params, _ = mcore_param_count( model.config, @@ -465,17 +484,18 @@ def _test_mcore_mamba_hybrid_pruning_nas_memory_mb(rank, size, ckpt_dir): sorted_layers = _get_sorted_layers(searcher_state) # fmt: off if sorted_layers == [1, 4, 3, 2]: + # All moe_ffn_hidden_size 16: the 12s do not divide the unpruned shared size of 16. expected_top_k = [ - [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 6, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 24}, {"memory_mb": 0.0226287841796875}, 114], # noqa: E501 [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 12, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 32}, {"memory_mb": 0.022613525390625}, 124], # noqa: E501 [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 6, "mamba_head_dim": 16, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 32}, {"memory_mb": 0.022556304931640625}, 126], # noqa: E501 - [{"num_layers": 4, "hidden_size": 16, "mamba_num_heads": 6, "mamba_head_dim": 12, "num_moe_experts": 7, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 24}, {"memory_mb": 0.022541046142578125}, 113], # noqa: E501 - [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 5, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 112], # noqa: E501 [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 5, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 116], # noqa: E501 - [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 6, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 113], # noqa: E501 [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 6, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 117], # noqa: E501 - [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 7, "moe_ffn_hidden_size": 12, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 114], # noqa: E501 [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 7, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 118], # noqa: E501 + [{"num_layers": 3, "hidden_size": 16, "mamba_num_heads": 8, "mamba_head_dim": 16, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 20}, {"memory_mb": 0.0225067138671875}, 119], # noqa: E501 + [{"num_layers": 4, "hidden_size": 16, "mamba_num_heads": 6, "mamba_head_dim": 12, "num_moe_experts": 5, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 28}, {"memory_mb": 0.022480010986328125}, 119], # noqa: E501 + [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 12, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 28}, {"memory_mb": 0.022430419921875}, 120], # noqa: E501 + [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 6, "mamba_head_dim": 16, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 28}, {"memory_mb": 0.022373199462890625}, 122], # noqa: E501 + [{"num_layers": 4, "hidden_size": 12, "mamba_num_heads": 8, "mamba_head_dim": 12, "num_moe_experts": 8, "moe_ffn_hidden_size": 16, "ffn_hidden_size": 24}, {"memory_mb": 0.022247314453125}, 116], # noqa: E501 ] else: raise RuntimeError(f"FIXME: Non deterministic test, assertions may fail: {sorted_layers=}") From b0a43476e057b34726aceeb6f04636552d968c85 Mon Sep 17 00:00:00 2001 From: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:51:16 +0800 Subject: [PATCH 11/11] fix(megatron): restore untied output_layer quantization under Megatron-Bridge (#2112) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Type of change:** Bug fix **Overview:** A quantized `output_layer` (lm_head) is **silently exported as BF16** when the model is built through Megatron-Bridge — no error, no warning, just an unquantized output layer in the exported checkpoint. This restores it. **1. The extra-state callbacks are never registered.** ```python if name.endswith("output_layer") and not getattr( getattr(module, "weight_quantizer", None), "is_enabled", False): continue ``` This hook also runs *before* `QuantModule` replacement (e.g. on restore), when `weight_quantizer` does not exist yet — so the check always skipped, and `output_layer` never received the ModelOpt extra-state callbacks. Its quantizer state (promotion to `StaticBlockScaleQuantizer`, `_amax`, `_global_amax`) was therefore never saved or restored. **2. Tiedness is read from a Megatron-LM global that Megatron-Bridge does not have.** ```python from megatron.training import get_args as _mlm_get_args _untied = bool(getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False)) except Exception as e: warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") ``` Megatron-Bridge has no `megatron.training` global args store, so the import raises, the `except` path treats the layer as tied, and quantization is dropped. The two combine into a silent failure: the save path writes nothing for the output layer, and the load path has nothing to restore. Resolve tiedness from the model instead of from a framework global: * `_resolve_output_layer_untied()` walks `named_modules()` for the framework-agnostic `share_embeddings_and_output_weights` flag and records it on the model config as `modelopt_output_layer_untied`. * `sharded_state_dict` prefers that flag and falls back to `get_args()` only when it is absent — **Megatron-LM behaviour is unchanged**. * Callback registration falls back to the tiedness flag when `weight_quantizer` does not exist yet, so an untied `output_layer` is registered. * When the load plan is built, the per-block NVFP4 scale buffers are materialized so the loader has somewhere to write. A Megatron dist-checkpoint load **silently skips any key the model does not advertise**, which is why the missing buffers produced no error. Divisibility is checked against `weight.shape[-1]` (the axis `_process_quantizer_amax` later views over), and a warning is emitted rather than leaving the buffers unallocated. `modelopt/torch/distill/plugins/megatron.py` gains a matching one-shot promotion on the first forward, for the single module the restore path does not promote. Measured on Nemotron-Nano-3: `already_sbsq=460 converted=1 disabled=0` — exactly `output_layer`. Verified end-to-end on Nemotron-Nano-3, W4A16 NVFP4 `four_over_six`, 4x GB200: * exported checkpoint carries `lm_head.weight` (U8 `[131072, 1344]`) plus `weight_scale` and `weight_scale_2` * matches a known-good reference export exactly: 52 shards, 18487 keys, 72 exclusions * converts to compressed-tensors and serves correctly on stock vLLM 0.26.0 at TP=2 * during QAD, `disabled=0` (was `disabled=1` before the fix) independently confirms the output layer stays quantized through distillation Without the fix the same pipeline produces a BF16 `lm_head` with no diagnostic of any kind. - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes. Megatron-LM keeps its existing `get_args()` path; the new model-derived flag is only consulted first, and only affects cases that are currently broken. - **Did you write any new necessary tests?**: Yes — unit coverage for both tiedness resolvers (`_resolve_output_layer_untied`, `_output_layer_untied`). Also verified end-to-end on Nemotron-Nano-3.5: byte-identical export, and QAD matching the reference iteration-10 KD loss at identical LR. - **Did you add or update any necessary documentation?**: No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No `quant_module_get_extra_state` now returns `{}` for a module with nothing quantized, so registering the extra-state callbacks unconditionally does not give unquantized layers non-empty extra state. An earlier revision of this PR gated registration on output-layer tiedness and did change the state-dict schema for untied-but-disabled `output_layer`; that revision has been replaced. 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Improved quantized Megatron model state handling when quantization is disabled. * Corrected output-layer weight-tying detection across model configurations, including vision and teacher components. * Improved compatibility with Megatron-LM settings when determining output-layer behavior. * Ensured sharded model state is generated consistently using the resolved weight-tying configuration. * Added reliable handling and caching of weight-tying information for consistent model state output. --------- Signed-off-by: James Shen Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../torch/quantization/plugins/megatron.py | 77 +++++++++--- .../quantization/plugins/test_megatron.py | 118 ++++++++++++++++++ 2 files changed, 180 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..d272fa400fe 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -133,6 +133,18 @@ def quant_module_get_extra_state(self) -> dict: QuantModule's extra_state with QuantModule.get_extra_state() which avoids the need to store the full module name. """ + # ``GPTModel.sharded_state_dict`` pops ``output_layer._extra_state`` and asserts it carries no + # data ("Expected output layer extra state to be empty", mcore models/gpt/gpt_model.py), so an + # output_layer with nothing quantized must contribute none. Scoped to output_layer: for every + # other module this quantizer_state is the only record that its quantizers were disabled (e.g. + # by auto_quantize or disable_quantizer), and dropping it would restore them enabled. + if ( + getattr(self, "_modelopt_output_layer", False) + and not isinstance(self, RealQuantLinear) + and not any(isinstance(m, TensorQuantizer) and m.is_enabled for m in self.modules()) + ): + return {} + extra_state = {} quantizer_state = {} @@ -248,6 +260,48 @@ def _incompatible_method(self, *args, **kwargs): return _incompatible_method +def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None: + """Whether ``output_layer`` weights are untied from the input embeddings, or None if unknown.""" + # named_modules() yields the root first, so the root's own flag wins when it has one. + for name, module in model.named_modules(): + # Skip subtrees that do not own the language model's output_layer: the vision tower (never + # quantized here) and a distillation teacher, which may be tied differently from the + # student it is wrapped with. + if "vision_model" in name or "_teacher_model" in name: + continue + shared = getattr(module, "share_embeddings_and_output_weights", None) + if shared is not None: + return not bool(shared) + return None + + +def _output_layer_untied(config) -> bool: + """Whether ``output_layer`` is untied, for use from ``sharded_state_dict``. + + Prefers the flag recorded by ``megatron_replace_quant_module_hook`` (the only source available + under Megatron-Bridge, which has no global args store), then Megatron-LM's + ``--untie-embeddings-and-output-weights``. + """ + untied = getattr(config, "modelopt_output_layer_untied", None) + if untied is not None: + return untied + try: + from megatron.training import get_args as _mlm_get_args + + return bool(getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False)) + except (ImportError, AssertionError) as e: + # ImportError: no megatron.training. AssertionError: get_args() before initialize_megatron. + # Warn once per config rather than on every save and every load. + if not getattr(config, "_modelopt_warned_output_layer_untied", False): + warn_rank_0( + f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}. " + "Treating output_layer as tied; its quantizer state will not be saved or " + "restored. If output_layer is in fact untied, it will be exported unquantized." + ) + config._modelopt_warned_output_layer_untied = True + return False + + def megatron_replace_quant_module_hook(model: torch.nn.Module): """Configure Megatron-Core model quantization support. @@ -260,6 +314,7 @@ def megatron_replace_quant_module_hook(model: torch.nn.Module): typing-matching the QuantModuleRegistry. 3. For Attention modules, we configure them to use core_attention path for KV cache quantization. """ + untied = _resolve_output_layer_untied(model) def _configure_attention_for_kv_cache_quant(module: Attention): """Configure Attention module for KV cache quantization compatibility.""" @@ -287,11 +342,8 @@ def _configure_attention_for_kv_cache_quant(module: Attention): def _register_extra_state_callbacks(model: torch.nn.Module): for name, module in model.named_modules(): if type(module) in QuantModuleRegistry: - # Skip output_layer w/o enabled weight_quantizer - if name.endswith("output_layer") and not getattr( - getattr(module, "weight_quantizer", None), "is_enabled", False - ): - continue + if name.endswith("output_layer"): + module._modelopt_output_layer = True register_modelopt_extra_state_callbacks( module, quant_module_get_extra_state, @@ -307,6 +359,10 @@ def _register_extra_state_callbacks(model: torch.nn.Module): if "vision_model" not in name: # We only enable hetereogenous_dist_checkpoint for language model, vision model is not quantized module.config.hetereogenous_dist_checkpoint = True + # Read back via ``self.config`` in sharded_state_dict; output_layer shares its + # parent MegatronModule's config object. The teacher may be tied differently. + if untied is not None and "_teacher_model" not in name: + module.config.modelopt_output_layer_untied = untied _register_extra_state_callbacks(module) @@ -374,16 +430,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # output_layer.input_quantizer._amax but TP-only does not. This lead to # state_dict mismatch. if prefix.endswith("output_layer."): - try: - from megatron.training import get_args as _mlm_get_args - - _untied = bool( - getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) - ) - except Exception as e: - warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") - _untied = False - if not _untied: + if not _output_layer_untied(self.config): return super().sharded_state_dict(prefix, sharded_offsets, metadata) quantizer_state_dict = {} diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 36f80787931..504e9cce965 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -14,9 +14,13 @@ # limitations under the License. import copy +import sys +import types from contextlib import nullcontext from functools import partial from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch import pytest import torch @@ -49,6 +53,7 @@ get_tensor_model_parallel_group, ) from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer import MegatronModule, TransformerConfig from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP from megatron.core.transformer.moe.router import TopKRouter @@ -58,8 +63,11 @@ from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.megatron import ( + _output_layer_untied, _QuantTEMCoreRowParallelLinear, + _resolve_output_layer_untied, get_mcore_layerwise_calibration_layers, + megatron_replace_quant_module_hook, ) from modelopt.torch.quantization.utils import is_quantized_linear from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -1301,3 +1309,113 @@ def test_homogeneous_sharded_state_dict_te_spec(dist_workers, tmp_path): {"transformer_impl": "transformer_engine"}, ), ) + + +def test_resolve_output_layer_untied(): + """The tiedness signal is read off the model, not from Megatron-LM global args.""" + + class _Flagged(torch.nn.Module): + def __init__(self, shared): + super().__init__() + self.share_embeddings_and_output_weights = shared + + # No signal anywhere -> unknown. + assert _resolve_output_layer_untied(torch.nn.Module()) is None + + # The root's own flag wins over any subtree. + root = _Flagged(False) + root.inner = _Flagged(True) + assert _resolve_output_layer_untied(root) is True + + # Otherwise fall back to a subtree scan. + root = torch.nn.Module() + root.language_model = _Flagged(True) + assert _resolve_output_layer_untied(root) is False + + # Subtrees that do not own the language model's output_layer are skipped: the vision tower + # and a distillation teacher, either of which may be tied differently from the student. + root = torch.nn.Module() + root.vision_model = _Flagged(True) + root._teacher_model = _Flagged(True) + root.language_model = _Flagged(False) + assert _resolve_output_layer_untied(root) is True + + +@pytest.mark.parametrize("mlm_untied", [True, False]) +def test_output_layer_untied_falls_back_to_megatron_lm_args(mlm_untied): + """With no model-derived flag, the answer comes from Megatron-LM's args.""" + + class _Config: + pass + + fake_training = types.ModuleType("megatron.training") + fake_training.get_args = lambda: SimpleNamespace(untie_embeddings_and_output_weights=mlm_untied) + + config = _Config() + with patch.dict(sys.modules, {"megatron.training": fake_training}): + assert _output_layer_untied(config) is mlm_untied + + # The model-derived flag takes precedence over the args fallback. + config.modelopt_output_layer_untied = not mlm_untied + with patch.dict(sys.modules, {"megatron.training": fake_training}): + assert _output_layer_untied(config) is (not mlm_untied) + + +def test_output_layer_untied_warns_once_when_args_unavailable(): + """Without either signal the layer is treated as tied, and the warning is not repeated.""" + + class _Config: + pass + + broken = types.ModuleType("megatron.training") # no get_args attribute + + config = _Config() + with ( + patch.dict(sys.modules, {"megatron.training": broken}), + patch("modelopt.torch.quantization.plugins.megatron.warn_rank_0") as warn, + ): + assert _output_layer_untied(config) is False + assert _output_layer_untied(config) is False + assert warn.call_count == 1 + + +def test_output_layer_untied_warns_when_args_uninitialized(): + """Megatron-LM importable but not initialized: treated as tied, warned once.""" + + class _Config: + pass + + def _uninitialized(): + raise AssertionError("args is not initialized.") + + fake_training = types.ModuleType("megatron.training") + fake_training.get_args = _uninitialized + + config = _Config() + with ( + patch.dict(sys.modules, {"megatron.training": fake_training}), + patch("modelopt.torch.quantization.plugins.megatron.warn_rank_0") as warn, + ): + assert _output_layer_untied(config) is False + assert _output_layer_untied(config) is False + assert warn.call_count == 1 + + +def test_output_layer_untied_not_stamped_onto_teacher_config(): + """A distillation teacher keeps its own tiedness; the student's answer must not leak in.""" + + def _config(): + return TransformerConfig(num_layers=1, hidden_size=8, num_attention_heads=1) + + class _Tiny(MegatronModule): + def __init__(self, config, shared): + super().__init__(config) + self.share_embeddings_and_output_weights = shared + + student = _Tiny(_config(), shared=False) # untied + student._teacher_model = _Tiny(_config(), shared=True) # tied -- must not be overwritten + + megatron_replace_quant_module_hook(student) + + assert student.config.modelopt_output_layer_untied is True + assert not hasattr(student._teacher_model.config, "modelopt_output_layer_untied")