Enable Yarn RoPE in minitron pruning for gpt-oss

Signed-off-by: Keval Morabia <[email protected]>

Author: kevalmorabia97

CHANGELOG.rst

@@ -10,7 +10,7 @@ Model Optimizer Changelog (Linux)
 
 **New Features**
 
-- Add MoE (e.g. Qwen3-30B-A3B) pruning support for ``num_moe_experts``, ``moe_ffn_hidden_size`` and ``moe_shared_expert_intermediate_size`` parameters in Minitron pruning (``mcore_minitron``).
+- Add MoE (e.g. Qwen3-30B-A3B, gpt-oss-20b) pruning support for ``num_moe_experts``, ``moe_ffn_hidden_size`` and ``moe_shared_expert_intermediate_size`` parameters in Minitron pruning (``mcore_minitron``).
 - Add ``specdec_bench`` example to benchmark speculative decoding performance. See `examples/specdec_bench/README.md <https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/specdec_bench#speculative-decoding-benchmark>`_ for more details.
 
 0.39 (2025-11-14)

examples/megatron-lm/README.md

@@ -24,6 +24,7 @@
 | `Qwen/Qwen3-{0.6B, 8B}` | ✅ | **Online** | | ✅ | ✅ |
 | `deepseek-ai/DeepSeek-R1` | ✅ | **Online** | | | |
 | `meta-llama/Llama-{3.1-8B, 3.1-405B, 3.2-1B}-Instruct` | ✅ | **Online** | | ✅ | ✅ |
+| `openai/gpt-oss-{20b, 120b}` | ✅ | **Online** | | ✅ | ✅ |
 
 ## Getting Started in a Local Environment
 

examples/pruning/README.md

@@ -6,7 +6,7 @@ Pruning can involve removal (prune) of Linear and Conv layers, and Transformer a
 
 This section focuses on applying Model Optimizer's state-of-the-art complementary pruning modes to enable you to search for the best subnet architecture from your provided base model:
 
-1. [Minitron](https://arxiv.org/pdf/2408.11796): A pruning method developed by NVIDIA Research for pruning GPT, Mamba and Hybrid Transformer Mamba models in NVIDIA NeMo or Megatron-LM framework. It uses the activation magnitudes to prune the embedding hidden size; mlp ffn hidden size; transformer attention heads and GQA query groups; mamba heads and head dimension; MoE number of experts, ffn hidden size, and shared expert intermediate size; and number of layers of the model.
+1. [Minitron](https://arxiv.org/pdf/2408.11796): A pruning method developed by NVIDIA Research for pruning GPT (and later extended to Mamba, MoE, and Hybrid Transformer Mamba) models in NVIDIA Megatron-LM or NeMo framework. It uses the activation magnitudes to prune the embedding hidden size; mlp ffn hidden size; transformer attention heads and GQA query groups; mamba heads and head dimension; MoE number of experts, ffn hidden size, and shared expert intermediate size; and number of layers of the model.
 1. FastNAS: A pruning method recommended for Computer Vision models. Given a pretrained model, FastNAS finds the subnet which maximizes the score function while meeting the given constraints.
 1. GradNAS: A light-weight pruning method recommended for language models like Hugging Face BERT, GPT-J. It uses the gradient information to prune the model's linear layers and attention heads to meet the given constraints.
 
@@ -89,7 +89,7 @@ If your model parameters are already sorted, you can skip the sorting step by se
 
 | **Algorithm** | **Model** | **Pruning Constraints** |
 | :---: | :---: | :---: |
-| Minitron | Megatron-core / NeMo based GPT / Mamba / MoE / Hybrid Models<sup>1</sup> | Export config with width (`hidden_size`, `ffn_hidden_size`, `num_attention_heads`, `num_query_groups`, `mamba_num_heads`, `mamba_head_dim`, `num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`) and/or depth (`num_layers`) values |
+| Minitron | Megatron-core / NeMo based GPT / Mamba / MoE / Hybrid LLM Models<sup>1</sup> | Export config with width (`hidden_size`, `ffn_hidden_size`, `num_attention_heads`, `num_query_groups`, `mamba_num_heads`, `mamba_head_dim`, `num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`) and/or depth (`num_layers`) values |
 | FastNAS | Computer Vision models | flops, parameters |
 | GradNAS | HuggingFace BERT, GPT-J | flops, parameters |
 

modelopt/torch/nas/plugins/megatron.py

@@ -24,6 +24,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.models.gpt import GPTModel
 from megatron.core.parallel_state import (
     get_data_parallel_group,
@@ -154,22 +155,54 @@ def _setup(self):
         )
 
 
+# Embedding DynamicModule ##########################################################################
 @DMRegistry.register(
-    {VocabParallelEmbedding: "megatron.core.tensor_parallel.layers.VocabParallelEmbedding"}
+    {
+        VocabParallelEmbedding: "megatron.core.tensor_parallel.layers.VocabParallelEmbedding",
+        nn.Embedding: "nn.Embedding",
+    },
 )
-class _DynamicVocabParallelEmbedding(DynamicModule):
-    """A VocabParallelEmbedding layer with dynamic hyperparams."""
+class _DynamicEmbedding(DynamicModule):
+    """A Embedding layer with dynamic hyperparams."""
 
     def _setup(self):
         self._register_hparam("embedding_dim", TracedHp(list(range(1, self.embedding_dim + 1))))
         self._register_dynamic_attribute("weight", self._get_weight)
 
     @staticmethod
-    def _get_weight(mod: "_DynamicVocabParallelEmbedding", weight: torch.Tensor) -> torch.Tensor:
+    def _get_weight(mod: "_DynamicEmbedding", weight: torch.Tensor) -> torch.Tensor:
         """Return the weight tensor of the embedding layer."""
         return get_sliced_tensor(mod, weight, None, "embedding_dim")
 
 
+@DMRegistry.register(
+    {
+        LanguageModelEmbedding: "megatron.core.models.common.embeddings.language_model_embedding.LanguageModelEmbedding"
+    }
+)
+class _DynamicLanguageModelEmbedding(DynamicModule):
+    """A LanguageModelEmbedding layer with dynamic hyperparams."""
+
+    def _setup(self):
+        DMRegistry.convert(self.word_embeddings)
+        hp_hidden_size = self.word_embeddings.get_hparam("embedding_dim")
+        if hasattr(self, "position_embeddings") and self.position_embeddings is not None:
+            DMRegistry.convert(self.position_embeddings)
+            self.position_embeddings.embedding_dim = hp_hidden_size
+        if hasattr(self, "tokentype_embeddings") and self.tokentype_embeddings is not None:
+            DMRegistry.convert(self.tokentype_embeddings)
+            self.tokentype_embeddings.embedding_dim = hp_hidden_size
+
+    def export(self) -> torch.nn.Module:
+        self.word_embeddings.export()
+        if hasattr(self, "position_embeddings") and self.position_embeddings is not None:
+            self.position_embeddings.export()
+        if hasattr(self, "tokentype_embeddings") and self.tokentype_embeddings is not None:
+            self.tokentype_embeddings.export()
+        return super().export()
+
+
+# Normalization DynamicModule ######################################################################
 @DMRegistry.register({FusedLayerNorm: "megatron.core.fusions.fused_layer_norm.FusedLayerNorm"})
 class _DynamicFusedLayerNorm(_DynamicLayerNorm):
     """A FusedLayerNorm layer with dynamic hyperparams."""
@@ -211,8 +244,8 @@ def _setup(self):
             self.hparam_name = "moe_ffn_hidden_size"
         else:
             self.hparam_name = "ffn_hidden_size"
-        self.linear_fc1 = DMRegistry.convert(self.linear_fc1)
-        self.linear_fc2 = DMRegistry.convert(self.linear_fc2)
+        DMRegistry.convert(self.linear_fc1)
+        DMRegistry.convert(self.linear_fc2)
 
         ffn_hidden_size = TracedHp(list(range(1, self.config.ffn_hidden_size + 1)))
         fc1_output_size = (
@@ -537,7 +570,7 @@ def _setup(self):
                 {"_setup": lambda self: None},
             )
 
-            self.core_attention = _DynamicDotProductAttention.convert(self.core_attention)
+            _DynamicDotProductAttention.convert(self.core_attention)
             self.core_attention._register_dynamic_attribute(
                 "hidden_size_per_partition",
                 lambda mod, val: self.config.kv_channels * self.num_attention_heads_per_partition,
@@ -559,7 +592,7 @@ def _setup(self):
                 {"_setup": lambda self: None},
             )
 
-            self.core_attention = _DynamicTEDotProductAttention.convert(self.core_attention)
+            _DynamicTEDotProductAttention.convert(self.core_attention)
             self.core_attention._register_dynamic_attribute(
                 "num_attention_heads", lambda mod, val: self.num_attention_heads_per_partition
             )
@@ -571,10 +604,10 @@ def _setup(self):
             )
 
         # Convert the fused qkv and output projection linear layer to dynamic module
-        self.linear_qkv = _DynamicQKVColumnParallelLinear.convert(
+        _DynamicQKVColumnParallelLinear.convert(
             self.linear_qkv, num_heads_per_group, num_query_groups
         )
-        self.linear_proj = _DynamicProjRowParallelLinear.convert(
+        _DynamicProjRowParallelLinear.convert(
             self.linear_proj, num_heads_per_group, num_query_groups
         )
 
@@ -699,10 +732,10 @@ def _setup(self):
         self._register_hparam("num_local_experts", num_moe_experts)
 
         # Convert local_experts list and each individual expert MLP to dynamic modules
-        self.local_experts = DynamicModuleList.convert(self.local_experts)
+        DynamicModuleList.convert(self.local_experts)
         self.local_experts.depth = num_moe_experts  # Reuse same hparam for depth
         for i in range(len(self.local_experts)):
-            self.local_experts[i] = DMRegistry.convert(self.local_experts[i])
+            DMRegistry.convert(self.local_experts[i])
 
         # Track forward activations for importance estimation.
         # _activations name is needed for get_activations_and_layer_scores to save scores for re-running pruning.
@@ -777,8 +810,8 @@ def _setup(self):
         # Convert to dynamic modules
         # Reuse _DynamicSequentialMLP's num_moe_experts hparam for _DynamicTopKRouter's hparam so
         #   importance estimator and depth hparam is retained.
-        self.router = DMRegistry.convert(self.router)
-        self.experts = DMRegistry.convert(self.experts)
+        DMRegistry.convert(self.router)
+        DMRegistry.convert(self.experts)
         num_moe_experts_hp = self.experts.get_hparam("num_local_experts")
 
         # NOTE: Use num_moe_experts hparam name in top-level module to match TransformerConfig's name
@@ -789,7 +822,7 @@ def _setup(self):
         )
         self.router.num_experts = num_moe_experts_hp
         if self.use_shared_expert:
-            self.shared_experts = DMRegistry.convert(self.shared_experts)
+            DMRegistry.convert(self.shared_experts)
 
     def forward(self, *args, **kwargs):
         """Forward pass for the MoE layer."""
@@ -898,12 +931,12 @@ def _setup(self):
         # Convert the layernorms, self-attention, and mlp/moe layers to dynamic modules
         # NOTE: Mamba stack layers have either Attention or MLP, not both unlike GPT models
         if isinstance(self.self_attention, SelfAttention):
-            self.input_layernorm = DMRegistry.convert(self.input_layernorm)
-            self.self_attention = DMRegistry.convert(self.self_attention)
+            DMRegistry.convert(self.input_layernorm)
+            DMRegistry.convert(self.self_attention)
 
         if isinstance(self.mlp, (MLP, MoELayer)):
-            self.pre_mlp_layernorm = DMRegistry.convert(self.pre_mlp_layernorm)
-            self.mlp = DMRegistry.convert(self.mlp)
+            DMRegistry.convert(self.pre_mlp_layernorm)
+            DMRegistry.convert(self.mlp)
 
         # Register forward hook to collect activations for importance estimation
         self._setup_mixin()
@@ -1168,23 +1201,23 @@ def _setup(self):
         self._register_dynamic_attribute("headdim", lambda mod, val: self.mamba_head_dim)
 
         # Convert to dynamic modules
-        self.in_proj = DMRegistry.convert(self.in_proj)
+        DMRegistry.convert(self.in_proj)
         self.in_proj.output_size = build_concat_hp(
             [d_inner, d_inner, bc, mamba_num_heads]
         )  # z, x, B, C, dt
 
         conv_dim = build_concat_hp([d_inner, bc])  # z, B, C
-        self.conv1d = DMRegistry.convert(self.conv1d)
+        DMRegistry.convert(self.conv1d)
         self.conv1d.in_channels = conv_dim
         self.conv1d.out_channels = conv_dim
         ks = self.conv1d.get_hparam("kernel_size")
         ks.choices = [ks.original]
 
         if self.rmsnorm:
-            self.norm = DMRegistry.convert(self.norm)
+            DMRegistry.convert(self.norm)
             self.norm.hidden_size = d_inner
 
-        self.out_proj = DMRegistry.convert(self.out_proj)
+        DMRegistry.convert(self.out_proj)
         self.out_proj.input_size = d_inner
 
         # Register dynamic attributes for Mamba-specific parameters
@@ -1310,8 +1343,8 @@ class _DynamicMambaLayer(DynamicModule, MambaTransformerLayerMixin):
 
     def _setup(self):
         # Convert to dynamic module
-        self.mixer = DMRegistry.convert(self.mixer)
-        self.norm = DMRegistry.convert(self.norm)
+        DMRegistry.convert(self.mixer)
+        DMRegistry.convert(self.norm)
         self._setup_mixin()
 
     def set_hidden_size_hp(self, hidden_size: TracedHp) -> None:
@@ -1336,7 +1369,7 @@ def modify(
             ("mamba_head_dim", mamba_head_dim_divisor),
         ]:
             hp = self.mixer.get_hparam(hp_name)
-            choices = {int(make_divisible(c, divisor)) for c in hp.choices}  # type: ignore[arg-type]
+            choices = {int(make_divisible(c, divisor)) for c in hp.choices}
             hp.choices = list(set(hp.choices) & choices | {hp.original})
 
     def export(self):
@@ -1376,24 +1409,21 @@ def _setup(self):
         assert self.config.expert_model_parallel_size == 1, "Expert parallel is not supported."
         assert self.pre_process == is_pipeline_first_stage()
         assert self.post_process == is_pipeline_last_stage()
-        assert self.position_embedding_type in ["rope", "none"], (
-            f"Only rope position embedding is supported, got {self.position_embedding_type}."
-        )
 
         # Register num_layers hparam for depth pruning
         self._register_hparam("num_layers", TracedHp(list(range(1, self.config.num_layers + 1))))
 
         # Convert layers to dynamic modules and set the shared hidden_size hparam
         if is_pipeline_first_stage():
-            self.embedding.word_embeddings = DMRegistry.convert(self.embedding.word_embeddings)
+            DMRegistry.convert(self.embedding)
             hidden_size = self.embedding.word_embeddings.get_hparam("embedding_dim")
         else:
             hidden_size = None
         hidden_size = dist.broadcast(hidden_size, src=0)
         self._register_hparam("hidden_size", hidden_size)
 
         for i in range(len(self.decoder.layers)):
-            self.decoder.layers[i] = DMRegistry.convert(self.decoder.layers[i])
+            DMRegistry.convert(self.decoder.layers[i])
             self.decoder.layers[i].set_hidden_size_hp(hidden_size)
 
         # NOTE: GPTModel has final_layernorm, MambaModel has final_norm
@@ -1409,7 +1439,7 @@ def _setup(self):
                 DMRegistry.convert(getattr(self.decoder, self.final_norm_attr_name)),
             )
             getattr(self.decoder, self.final_norm_attr_name).num_features = hidden_size
-            self.output_layer = DMRegistry.convert(self.output_layer)
+            DMRegistry.convert(self.output_layer)
             self.output_layer.input_size = hidden_size
             self.output_layer.get_hparam("output_size").choices = [self.output_layer.output_size]
 
@@ -1548,7 +1578,7 @@ def export(self) -> torch.nn.Module:
             handle.remove()
         self._export_drop_layers()
         if is_pipeline_first_stage():
-            self.embedding.word_embeddings.export()
+            self.embedding.export()
         for layer in self.decoder.layers:
             layer.export()
         if is_pipeline_last_stage():
@@ -1587,9 +1617,7 @@ def set_activations_and_layer_scores(
         rank = get_pipeline_model_parallel_rank()
         pp_size = get_pipeline_model_parallel_world_size()
         assert len(activations_per_rank) == pp_size, (
-            len(activations_per_rank),
-            activations_per_rank,
-            pp_size,
+            f"Expected same PP size for stored pruning scores ({len(activations_per_rank)}) as current ({pp_size})!"
         )
         for layer in self.decoder.layers:
             layer._scores = layer_scores[layer.layer_number]
@@ -1611,14 +1639,14 @@ def drop_mcore_language_model_layers(model: nn.Module, *, layers_to_drop: list[i
     )
 
     supported_model_types = tuple(SUPPORTED_MODELS.keys())
-    for m in model.modules():
+    for n, m in model.named_modules():
         if isinstance(m, supported_model_types):
             model = m
             break
     assert isinstance(model, supported_model_types), (
         f"Model should have one of {supported_model_types} submodule, got {model}"
     )
-    print_rank_0(f"Dropping layers {layers_to_drop} from {type(model)}.")
+    print_rank_0(f"Dropping layers {layers_to_drop} from {n} ({type(model)}).")
 
     # get the number of layers remaining in each pp rank
     layers_remaining_per_pp = torch.zeros(

tests/_test_utils/torch/megatron/models.py

@@ -138,6 +138,7 @@ def get_mcore_gpt_model(
     ffn_hidden_size: int | None = 128,
     max_sequence_length: int = 16,
     vocab_size: int = 64,
+    position_embedding_type: str = "rope",
     activation_func: str = "swiglu",
     normalization: str = "LayerNorm",
     transformer_impl: str = "modelopt" if HAS_TE else "local",
@@ -191,9 +192,21 @@ def squared_relu(x):
         moe_router_dtype="fp32",
         moe_ffn_hidden_size=moe_ffn_hidden_size,
         moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size,
+        moe_router_enable_expert_bias=True,
+        moe_router_score_function="sigmoid",
         num_moe_experts=num_moe_experts,
     )
 
+    if position_embedding_type == "yarn":  # gpt-oss like model
+        warn("Yarn RoPE config format will change soon. This is a temporary workaround")
+        config.yarn_rotary_scaling_factor = 32.0
+        config.yarn_original_max_position_embeddings = 4096
+        config.yarn_beta_fast = 32.0
+        config.yarn_beta_slow = 1.0
+        config.yarn_mscale = 1.0
+        config.yarn_mscale_all_dim = 0.0
+        config.yarn_correction_range_round_to_int = False
+
     if transformer_impl == "local":
         assert HAS_APEX, "Apex not installed"
         transformer_layer_spec = get_gpt_layer_local_spec(
@@ -224,7 +237,7 @@ def squared_relu(x):
         pre_process=is_pipeline_first_stage(),
         post_process=is_pipeline_last_stage(),
         share_embeddings_and_output_weights=False,
-        position_embedding_type="rope",
+        position_embedding_type=position_embedding_type,
     )
     return model.to(torch.bfloat16) if bf16 else model
 

tests/gpu/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py

@@ -29,8 +29,8 @@
     run_mcore_inference_with_dummy_input,
 )
 from _test_utils.torch.misc import set_seed
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.parallel_state import destroy_model_parallel
-from megatron.core.tensor_parallel.layers import VocabParallelEmbedding
 from megatron.core.transformer.attention import SelfAttention
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.mlp import MLP
@@ -41,6 +41,8 @@
 from modelopt.torch.nas.modules import DynamicModuleList
 from modelopt.torch.nas.plugins.megatron import (
     _DynamicColumnParallelLinear,
+    _DynamicEmbedding,
+    _DynamicLanguageModelEmbedding,
     _DynamicMCoreLanguageModel,
     _DynamicMLP,
     _DynamicMoELayer,
@@ -51,7 +53,6 @@
     _DynamicSequentialMLP,
     _DynamicTopKRouter,
     _DynamicTransformerLayer,
-    _DynamicVocabParallelEmbedding,
     expand_head_indices,
 )
 from modelopt.torch.opt.utils import named_dynamic_modules, search_space_size
@@ -92,8 +93,9 @@ def _test_gpt_search_space(
 
     assert isinstance(model, _DynamicMCoreLanguageModel)
     for m in model.modules():
-        if isinstance(m, VocabParallelEmbedding):
-            assert isinstance(m, _DynamicVocabParallelEmbedding)
+        if isinstance(m, LanguageModelEmbedding):
+            assert isinstance(m, _DynamicLanguageModelEmbedding)
+            assert isinstance(m.word_embeddings, _DynamicEmbedding)
         elif isinstance(m, TransformerLayer):
             assert isinstance(m, _DynamicTransformerLayer)
         elif isinstance(m, MLP):