From cda764e113820e00f3f489f3179786d09eb85ac7 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Mon, 30 Mar 2026 16:37:41 -0700
Subject: [PATCH 01/18] copy gpt cfg and builder

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/__init__.py |   3 +
 megatron/training/models/gpt.py      | 384 +++++++++++++++++++++++++++
 2 files changed, 387 insertions(+)
 create mode 100644 megatron/training/models/gpt.py

diff --git a/megatron/training/models/__init__.py b/megatron/training/models/__init__.py
index 62b4a18a4dd..39e52b0a371 100644
--- a/megatron/training/models/__init__.py
+++ b/megatron/training/models/__init__.py
@@ -6,6 +6,7 @@
     unimodal_build_distributed_models,
 )
 from megatron.training.models.hybrid import HybridModelBuilder, HybridModelConfig
+from megatron.training.models.gpt import GPTModelBuilder, GPTModelConfig
 
 MambaModelConfig = HybridModelConfig
 MambaModelBuilder = HybridModelBuilder
@@ -21,4 +22,6 @@
     "HybridModelBuilder",
     "MambaModelConfig",
     "MambaModelBuilder",
+    "GPTModelConfig",
+    "GPTModelBuilder"
 ]
diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
new file mode 100644
index 00000000000..131f1b08114
--- /dev/null
+++ b/megatron/training/models/gpt.py
@@ -0,0 +1,384 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+
+import inspect
+import logging
+from typing import Any, Callable, ClassVar, Literal, override
+
+import torch
+from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
+from megatron.core.enums import ModelType
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.pipeline_parallel.utils import (
+    is_pp_first_stage,
+    is_pp_last_stage,
+    is_vp_first_stage,
+    is_vp_last_stage,
+)
+from megatron.core.post_training.modelopt.gpt.model_specs import get_gpt_modelopt_spec
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.module import Float16Module, MegatronModule
+from megatron.core.transformer.dot_product_attention import DotProductAttention as MCoreDotProductAttention
+from megatron.core.transformer.enums import AttnBackend
+
+from megatron.training.models.base import ModelConfig, ModelBuilder, compose_hooks
+from megatron.training.vocab_utils import calculate_padded_vocab_size
+from megatron.training.models.base import unimodal_build_distributed_models
+
+from megatron.core.transformer.transformer_config import  TransformerConfig
+
+
+logger = logging.getLogger(__name__)
+
+
+from dataclasses import dataclass
+
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_local_spec,
+    get_gpt_layer_with_transformer_engine_spec,
+)
+
+
+
+def transformer_engine_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
+    """Create a Transformer Engine layer specification based on the provided config."""
+    if "use_te_op_fuser" in inspect.signature(get_gpt_layer_with_transformer_engine_spec).parameters:
+        kwargs = {"use_te_op_fuser": config.use_transformer_engine_op_fuser}
+    else:
+        kwargs = {}
+    return get_gpt_layer_with_transformer_engine_spec(
+        num_experts=config.transformer.num_moe_experts,
+        moe_grouped_gemm=config.transformer.moe_grouped_gemm,
+        qk_layernorm=config.transformer.qk_layernorm,
+        fp8=bool(config.transformer.num_moe_experts and (config.transformer.fp8 is not None)),
+        **kwargs,
+    )
+
+
+def local_layer_spec(config: TransformerConfig) -> ModuleSpec:
+    """Create a local layer specification without Transformer Engine.
+
+    Args:
+        config: GPT configuration object
+
+    Returns:
+        ModuleSpec: Module specification for local implementation layers
+    """
+    return get_gpt_layer_local_spec(
+        num_experts=config.num_moe_experts,
+        moe_grouped_gemm=config.moe_grouped_gemm,
+        qk_layernorm=config.qk_layernorm,
+        normalization=config.normalization,
+    )
+
+
+def modelopt_transformer_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
+    """Layer specification for quantization with ModelOpt."""
+    # arbitrary attention mask is used for speculative decoding training
+    # When context parallel > 1, only causal mask type is supported
+    from megatron.core import parallel_state
+
+    use_arbitrary_attention_mask = (
+        config.use_arbitrary_attention_mask
+        if config.use_arbitrary_attention_mask is not None
+        else parallel_state.get_context_parallel_world_size() == 1
+    )
+    return get_gpt_modelopt_spec(
+        config=config.transformer,
+        local_core_attention=False,
+        remap_te_layernorm=True,
+        real_quant_cfg="None",
+        use_arbitrary_attention_mask=use_arbitrary_attention_mask,
+    )
+
+
+def default_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
+    """Determine the most appropriate layer specification based on availability."""
+    if config.restore_modelopt_state:
+        return modelopt_transformer_layer_spec(config)
+    # TODO (@maanug): migrate this spec from Megatron-Bridge
+    # elif config.use_transformer_engine_full_layer_spec:
+    #     return transformer_engine_full_layer_spec(config.transformer)
+    else:
+        return transformer_engine_layer_spec(config)
+
+
+@dataclass(kw_only=True)
+class GPTModelConfig(ModelConfig):
+    """Configuration for a Megatron Core GPT model.
+
+    This is purely a configuration object. All model construction
+    logic lives in ``GPTModelBuilder``.
+
+    Contains a ``TransformerConfig`` alongside GPT-specific parameters. Attributes
+    on the embedded ``transformer`` config are accessible directly on this object
+    via ``__getattr__``/``__setattr__`` proxying.
+
+    Note:
+        ``vocab_size`` must be set before passing this config to ``GPTModelBuilder``.
+    """
+
+    builder: ClassVar[str] = "megatron.training.models.gpt.GPTModelBuilder"
+    transformer: TransformerConfig
+    transformer_layer_spec: ModuleSpec | Callable[["GPTModelConfig"], ModuleSpec] = default_layer_spec
+
+    ### vocab padding related ###
+    vocab_size: int | None = None
+    """This represents the unpadded vocab size. The padded vocab size is
+    automatically calculated in the GPTModelBuilder.
+    """
+    make_vocab_size_divisible_by: int = 128
+    should_pad_vocab: bool = False
+    """Set if the tokenizer provides the vocab size. In this case, the vocab size will be padded.
+    Controls whether vocab size should be padded for tensor parallelism.
+    """
+
+    ### GPT Model initialization ###
+    seq_length: int = 1024
+    fp16_lm_cross_entropy: bool = False
+    parallel_output: bool = True
+    share_embeddings_and_output_weights: bool = False
+    position_embedding_type: Literal["learned_absolute", "rope", "mrope", "yarn", "none"] = "learned_absolute"
+    rotary_percent: float = 1.0
+    rotary_base: int = 10000
+    rope_scaling: bool = False
+    rope_scaling_factor: float = 8.0
+    scatter_embedding_sequence_parallel: bool = True
+    seq_len_interpolation_factor: float | None = None
+
+    tp_comm_overlap_cfg: str | dict[str, Any] | None = None
+    """Config file when tp_comm_overlap is enabled."""
+
+    ### settings for default layer spec options ###
+    use_transformer_engine_op_fuser: bool = False
+    use_arbitrary_attention_mask: bool | None = None
+
+    @override
+    def __getattr__(self, name: str, /) -> Any:
+        # __getattr__ is only called when normal attribute lookup has already failed,
+        # so use object.__getattribute__ to fetch `transformer` without recursing.
+        try:
+            transformer = object.__getattribute__(self, "transformer")
+        except AttributeError:
+            raise AttributeError(f"MambaModelConfig has no attribute '{name}'")
+        if hasattr(transformer, name):
+            return getattr(transformer, name)
+        raise AttributeError(f"Neither MambaModelConfig nor TransformerConfig has any attribute '{name}'.")
+
+    @override
+    def __setattr__(self, name: str, value: Any, /) -> None:
+        # Use object.__getattribute__ to avoid triggering __getattr__ while
+        # `transformer` may not yet exist (e.g. during dataclass __init__).
+        try:
+            transformer = object.__getattribute__(self, "transformer")
+        except AttributeError:
+            # `transformer` not yet initialised; store the attribute on self.
+            super().__setattr__(name, value)
+            return
+        if hasattr(transformer, name):
+            setattr(transformer, name, value)
+        else:
+            super().__setattr__(name, value)
+
+    def finalize(self) -> None:
+        """One time validation to run once config is ready to be used by builder."""
+
+        if hasattr(self.transformer, "finalize") and callable(self.transformer.finalize):
+            self.transformer.finalize()
+
+        if self.transformer.cuda_graph_impl != "none":
+            assert self.transformer.use_te_rng_tracker, (
+                "Transformer engine's RNG tracker is required for cudagraphs, it can be "
+                "enabled with use_te_rng_tracker=True'."
+            )
+
+        vp_size = self.transformer.virtual_pipeline_model_parallel_size
+        is_pipeline_asymmetric = (
+            self.transformer.account_for_embedding_in_pipeline_split
+            or self.transformer.account_for_loss_in_pipeline_split
+        )
+        is_pipeline_asymmetric |= (
+            self.transformer.num_layers_in_first_pipeline_stage or self.transformer.num_layers_in_last_pipeline_stage
+        ) is not None
+        is_flexible_pp_layout = is_pipeline_asymmetric or (self.transformer.pipeline_model_parallel_layout is not None)
+        if vp_size and not is_flexible_pp_layout:
+            p_size = self.transformer.pipeline_model_parallel_size
+            assert (self.transformer.num_layers // p_size) % vp_size == 0, (
+                "Make sure the number of model chunks is the same across all pipeline stages."
+            )
+
+
+class GPTModelBuilder(ModelBuilder[GPTModel, GPTModelConfig]):
+    """Builder to construct Megatron Core GPT models.
+
+    Example:
+        >>> transformer_cfg = TransformerConfig(num_layers=32, hidden_size=4096, ...)
+        >>> model_cfg = GPTModelConfig(transformer=transformer_cfg, vocab_size=32000, seq_length=2048, ...)
+        >>>
+        >>> # Single stage (e.g. inference)
+        >>> model = GPTModelBuilder(model_cfg).build_model(pg_collection)
+        >>>
+        >>> # Distributed training
+        >>> models = GPTModelBuilder(model_cfg).build_distributed_models(pg_collection)
+    """
+
+    def __init__(self, model_config: GPTModelConfig):
+        super().__init__(model_config)
+
+    def build_model(
+        self,
+        pg_collection: ProcessGroupCollection,
+        pre_process: bool | None = None,
+        post_process: bool | None = None,
+        vp_stage: int | None = None,
+    ) -> GPTModel:
+        """Build a single ``MCoreGPTModel`` stage.
+
+        Args:
+            pg_collection: Process groups for distributed training
+            pre_process: Include embedding layer
+            post_process: Include output layer
+            vp_stage: Virtual pipeline stage
+
+        Returns:
+            The constructed model
+
+        Note:
+            Virtual pipeline model parallelism is not supported for Mamba models.
+        """
+        transformer_layer_spec = self._model_config.transformer_layer_spec
+        if not isinstance(transformer_layer_spec, ModuleSpec):
+            # Check if the transformer_layer_spec function accepts vp_stage parameter
+            if "vp_stage" in inspect.signature(transformer_layer_spec).parameters:
+                transformer_layer_spec = transformer_layer_spec(self._model_config, vp_stage=vp_stage)
+            else:
+                transformer_layer_spec = transformer_layer_spec(self._model_config)
+
+        assert self._model_config.vocab_size is not None, "vocab_size must be configured before calling build_model()"
+        if self._model_config.should_pad_vocab:
+            padded_vocab_size = calculate_padded_vocab_size(
+                self._model_config.vocab_size,
+                self._model_config.make_vocab_size_divisible_by,
+                self._model_config.transformer.tensor_model_parallel_size,
+            )
+        else:
+            padded_vocab_size = self._model_config.vocab_size
+
+        mtp_spec = mtp_block_spec(self._model_config, transformer_layer_spec, vp_stage=vp_stage)
+
+        # override spec with local backend if configured
+        if self._model_config.attention_backend == AttnBackend.local:
+            if hasattr(transformer_layer_spec, "submodules"):
+                transformer_layer_spec.submodules.self_attention.submodules.core_attention = MCoreDotProductAttention
+
+        # Determine pre/post flags if not provided using vp + pp stage
+        vp_size = self._model_config.virtual_pipeline_model_parallel_size
+        if pre_process is None:
+            pre_process = is_vp_first_stage(vp_stage=vp_stage, vp_size=vp_size) and is_pp_first_stage(pg_collection.pp)
+        if post_process is None:
+            post_process = is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size) and is_pp_last_stage(pg_collection.pp)
+
+        model = GPTModel(
+            config=self._model_config.transformer,
+            transformer_layer_spec=transformer_layer_spec,
+            mtp_block_spec=mtp_spec,
+            vocab_size=padded_vocab_size,
+            max_sequence_length=self._model_config.seq_length,
+            fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy,
+            parallel_output=self._model_config.parallel_output,
+            share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights,
+            position_embedding_type=self._model_config.position_embedding_type,
+            rotary_percent=self._model_config.rotary_percent,
+            rotary_base=self._model_config.rotary_base,
+            rope_scaling=self._model_config.rope_scaling,
+            rope_scaling_factor=self._model_config.rope_scaling_factor,
+            seq_len_interpolation_factor=self._model_config.seq_len_interpolation_factor,
+            scatter_embedding_sequence_parallel=self._model_config.scatter_embedding_sequence_parallel,
+            pre_process=pre_process,
+            post_process=post_process,
+            pg_collection=pg_collection,
+            vp_stage=vp_stage,
+        )
+
+        return model
+
+    def build_distributed_models(
+        self,
+        pg_collection: ProcessGroupCollection,
+        ddp_config: DistributedDataParallelConfig | None = None,
+        overlap_param_gather_with_optimizer_step: bool = False,
+        use_megatron_fsdp: bool = False,
+        use_torch_fsdp2: bool = False,
+        wrap_with_ddp: bool = True,
+        data_parallel_random_init: bool = True,
+        mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule] | None = Float16Module,
+        model_type: ModelType = ModelType.encoder_or_decoder,
+    ) -> list[GPTModel]:
+        """Build model stages and wrap for distributed training.
+
+        Args:
+            pg_collection: Model communication process groups.
+            ddp_config: DistributedDataParallel configuration
+            overlap_param_gather_with_optimizer_step: Whether to overlap parameter
+                gather with optimizer step.
+            use_megatron_fsdp: Whether to use Megatron FSDP
+            use_torch_fsdp2: Whether to use Torch FSDP 2.0
+            wrap_with_ddp: Set to False to skip the DDP/FSDP wrapper.
+            data_parallel_random_init: Whether to use data parallel random initialization
+            mixed_precision_wrapper: Mixed precision wrapper, e.g. ``Float16Module``
+            model_type: Deprecated flag, only used for backwards compatibility.
+
+        Returns:
+            List of model stages.
+        """
+        transformer_config = self._model_config.transformer
+        composed_pre_wrap_hook = compose_hooks(self._model_config.pre_wrap_hooks)
+        model_list = unimodal_build_distributed_models(
+            self.build_model,
+            transformer_config,
+            pg_collection,
+            ddp_config,
+            overlap_param_gather_with_optimizer_step,
+            use_megatron_fsdp,
+            use_torch_fsdp2,
+            wrap_with_ddp,
+            data_parallel_random_init,
+            mixed_precision_wrapper,
+            composed_pre_wrap_hook,
+            model_type,
+        )
+
+        composed_post_wrap_hook = compose_hooks(self._model_config.post_wrap_hooks)
+        _model = composed_post_wrap_hook(model_list)
+        if _model is not None:
+            model_list = _model
+        else:
+            logger.warning("Final post wrap hook returned None, skipping post wrap hooks.")
+
+        return model_list
+
+
+def mtp_block_spec(
+    config: "GPTModelConfig", transformer_layer_spec: ModuleSpec, vp_stage: int | None = None
+) -> ModuleSpec | None:
+    """Create MTP block spec if model has MTP layers.
+
+    Args:
+        config: full model config
+
+    Returns:
+        ModuleSpec: The MTP module specification
+    """
+    if config.transformer.mtp_num_layers is not None:
+        from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
+
+        if hasattr(transformer_layer_spec, "layer_specs") and len(transformer_layer_spec.layer_specs) == 0:
+            # Get the decoder layer spec explicitly if no decoder layer in the last stage,
+            # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
+            spec = default_layer_spec(config)
+        else:
+            spec = transformer_layer_spec
+
+        return get_gpt_mtp_block_spec(config.transformer, spec, use_transformer_engine=True, vp_stage=vp_stage)
+    else:
+        return None

From bafec695ffb98f7011252da9ec1b3cf608e7e5cc Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Mon, 11 May 2026 13:31:12 -0700
Subject: [PATCH 02/18] update copyright year

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index 131f1b08114..77678b1cba3 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
 
 import inspect
 import logging

From bbf51356ed3748a4f490b37bad5b51430535e026 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Mon, 11 May 2026 15:06:03 -0700
Subject: [PATCH 03/18] integrate into pretrain_gpt via container

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/argument_utils.py   | 12 +++++++++++-
 megatron/training/config/container.py |  4 ++--
 pretrain_gpt.py                       |  5 +++--
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py
index e6b111f91bf..874113fa5af 100644
--- a/megatron/training/argument_utils.py
+++ b/megatron/training/argument_utils.py
@@ -30,7 +30,7 @@
     StragglerDetectionConfig,
     RerunStateMachineConfig, CheckpointConfig, ProfilingConfig
 )
-from megatron.training.models.hybrid import HybridModelConfig
+from megatron.training.models import HybridModelConfig, GPTModelConfig
 # TODO: support arg renames
 
 class TypeInferenceError(Exception):
@@ -388,6 +388,16 @@ def _default_config_from_args(cls: type, args: Namespace, return_instance: bool
         return kwargs
 
 
+def gpt_config_from_args(args: Namespace, config: TransformerConfig | None=None) -> Any:
+    """Create a GPTModelConfig from the appropriate values in the `args` Namespace."""
+    kwargs = {}
+
+    # TODO (@maanug): impl
+    pass
+
+    return GPTModelConfig(**kwargs)
+    
+
 def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None=None) -> Any:
     """Create a HybridModelConfig from the appropriate values in the `args` Namespace."""
 
diff --git a/megatron/training/config/container.py b/megatron/training/config/container.py
index 1495167ae31..c13f73f52e9 100644
--- a/megatron/training/config/container.py
+++ b/megatron/training/config/container.py
@@ -27,7 +27,7 @@
 )
 from megatron.training.config.utils import sanitize_dataclass_config
 from megatron.training.config.yaml_utils import safe_yaml_representers
-from megatron.training.models import Serializable, HybridModelConfig
+from megatron.training.models import GPTModelConfig, Serializable, HybridModelConfig
 
 T = TypeVar("T", bound="ConfigContainerBase")
 
@@ -233,7 +233,7 @@ class PretrainConfigContainer(ConfigContainerBase):
 
     train: TrainingConfig
     validation: ValidationConfig = field(default_factory=ValidationConfig)
-    model: HybridModelConfig  # TODO (@maanug): add support for GPTModelConfig 
+    model: HybridModelConfig | GPTModelConfig
     optimizer: OptimizerConfig
     scheduler: SchedulerConfig
     # dataset: GPTDatasetConfig # TODO (@maanug): add support
diff --git a/pretrain_gpt.py b/pretrain_gpt.py
index 1c2ede9fa66..944c8cca7df 100644
--- a/pretrain_gpt.py
+++ b/pretrain_gpt.py
@@ -54,7 +54,7 @@
     print_rank_0,
     set_startup_timestamps,
 )
-from megatron.training.argument_utils import pretrain_cfg_container_from_args
+from megatron.training.argument_utils import pretrain_cfg_container_from_args, gpt_config_from_args
 from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
 from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig
 from megatron.training.datasets.sft_dataset import SFTDataset
@@ -390,7 +390,8 @@ def get_embedding_ranks(pp_ranks: List[int]):
         extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
         args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
     )
-    full_config = pretrain_cfg_container_from_args(args)
+    model_cfg = gpt_config_from_args(args)
+    full_config = pretrain_cfg_container_from_args(args, model_cfg)
     pretrain(full_config,
         train_valid_test_datasets_provider,
         partial(model_provider, gpt_builder),

From 7c5d6e49f74ce26c30e4ad06391bca9173a87db1 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sat, 16 May 2026 16:21:48 -0400
Subject: [PATCH 04/18] fix import

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index 77678b1cba3..d895e087669 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -23,7 +23,7 @@
 
 from megatron.training.models.base import ModelConfig, ModelBuilder, compose_hooks
 from megatron.training.vocab_utils import calculate_padded_vocab_size
-from megatron.training.models.base import unimodal_build_distributed_models
+from megatron.training.models.dist_utils import unimodal_build_distributed_models
 
 from megatron.core.transformer.transformer_config import  TransformerConfig
 
@@ -242,9 +242,6 @@ def build_model(
 
         Returns:
             The constructed model
-
-        Note:
-            Virtual pipeline model parallelism is not supported for Mamba models.
         """
         transformer_layer_spec = self._model_config.transformer_layer_spec
         if not isinstance(transformer_layer_spec, ModuleSpec):

From 0bdeed1283d69eb241b4eb6dec86890b4fa38237 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sat, 16 May 2026 16:30:27 -0400
Subject: [PATCH 05/18] change default layer spec to none

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index d895e087669..8d52df2a135 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -120,7 +120,7 @@ class GPTModelConfig(ModelConfig):
 
     builder: ClassVar[str] = "megatron.training.models.gpt.GPTModelBuilder"
     transformer: TransformerConfig
-    transformer_layer_spec: ModuleSpec | Callable[["GPTModelConfig"], ModuleSpec] = default_layer_spec
+    transformer_layer_spec: ModuleSpec | Callable[["GPTModelConfig"], ModuleSpec] | None = None
 
     ### vocab padding related ###
     vocab_size: int | None = None
@@ -244,7 +244,9 @@ def build_model(
             The constructed model
         """
         transformer_layer_spec = self._model_config.transformer_layer_spec
-        if not isinstance(transformer_layer_spec, ModuleSpec):
+        if transformer_layer_spec is None:
+            transformer_layer_spec = default_layer_spec(self._model_config)
+        elif callable(transformer_layer_spec):
             # Check if the transformer_layer_spec function accepts vp_stage parameter
             if "vp_stage" in inspect.signature(transformer_layer_spec).parameters:
                 transformer_layer_spec = transformer_layer_spec(self._model_config, vp_stage=vp_stage)

From 40c32459aff24575e94fa35d924d02a493924985 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sat, 16 May 2026 17:13:56 -0400
Subject: [PATCH 06/18] helper to create gpt cfg from args

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/argument_utils.py | 38 +++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py
index 874113fa5af..b2051e81360 100644
--- a/megatron/training/argument_utils.py
+++ b/megatron/training/argument_utils.py
@@ -390,10 +390,43 @@ def _default_config_from_args(cls: type, args: Namespace, return_instance: bool
 
 def gpt_config_from_args(args: Namespace, config: TransformerConfig | None=None) -> Any:
     """Create a GPTModelConfig from the appropriate values in the `args` Namespace."""
+
     kwargs = {}
+    if config is None:
+        if args.yaml_cfg is not None:
+            from megatron.training.yaml_arguments import core_transformer_config_from_yaml
 
-    # TODO (@maanug): impl
-    pass
+            transformer_cfg = core_transformer_config_from_yaml(args, "language_model")
+        else:
+            transformer_cfg = core_transformer_config_from_args(args)
+    else:
+        transformer_cfg = config
+    kwargs["transformer"] = transformer_cfg
+
+    if args.spec is not None:
+        kwargs["transformer_layer_spec"] = import_module(args.spec)
+
+
+    kwargs["fp16_lm_cross_entropy"] = args.fp16_lm_cross_entropy
+    kwargs["position_embedding_type"] = args.position_embedding_type
+    kwargs["rotary_percent"] = args.rotary_percent
+    kwargs["rotary_base"] = args.rotary_base
+    kwargs["make_vocab_size_divisible_by"] = args.make_vocab_size_divisible_by
+    kwargs["rope_scaling"] = args.use_rope_scaling
+
+    kwargs["seq_len_interpolation_factor"] = args.rotary_seq_len_interpolation_factor
+    kwargs["seq_length"] = args.max_position_embeddings
+    kwargs["share_embeddings_and_output_weights"] = not args.untie_embeddings_and_output_weights
+
+    if args.padded_vocab_size is not None:
+        kwargs["vocab_size"] = args.padded_vocab_size
+        kwargs["should_pad_vocab"] = False
+    else:
+        # Megatron-Bridge uses an explicit setting "should_pad_vocab" so that 
+        # when converting model configs from HF, we can set a vocab size and disable padding. 
+        assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified."
+        kwargs["vocab_size"] = args.vocab_size
+        kwargs["should_pad_vocab"] = True
 
     return GPTModelConfig(**kwargs)
     
@@ -429,6 +462,7 @@ def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None=No
 
     if args.padded_vocab_size is not None:
         kwargs["vocab_size"] = args.padded_vocab_size
+        kwargs["should_pad_vocab"] = False
     else:
         # Megatron-Bridge uses an explicit setting "should_pad_vocab" so that 
         # when converting model configs from HF, we can set a vocab size and disable padding. 

From bc2ea8fb1196f4b6555d9241a46903d73e038588 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sat, 16 May 2026 19:43:46 -0400
Subject: [PATCH 07/18] move assertion

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 .../core/transformer/heterogeneous/heterogeneous_config.py  | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/megatron/core/transformer/heterogeneous/heterogeneous_config.py b/megatron/core/transformer/heterogeneous/heterogeneous_config.py
index 550e7b67560..109eb67524c 100644
--- a/megatron/core/transformer/heterogeneous/heterogeneous_config.py
+++ b/megatron/core/transformer/heterogeneous/heterogeneous_config.py
@@ -201,6 +201,12 @@ class HeterogeneousTransformerConfig(TransformerConfig):
     def __post_init__(self):
         super().__post_init__()
 
+        err_msg = (
+            "Cannot use inference optimized transformer implementation "
+            "with heterogeneous transformer layers."
+        )
+        assert not (self.transformer_impl == "inference_optimized"), err_msg
+
         self.heterogeneous_block_specs = True
 
         if self.heterogeneous_layers_config_encoded_json in ("", None):

From 447ae0ebf455b84e34c33c8f7377a0137681f117 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 00:43:38 -0400
Subject: [PATCH 08/18] synchronize with gpt builders

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 73 +++++++++++++++++++++++++++------
 1 file changed, 61 insertions(+), 12 deletions(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index 8d52df2a135..9461d467ecd 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -4,6 +4,8 @@
 import logging
 from typing import Any, Callable, ClassVar, Literal, override
 
+from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import get_gpt_heterogeneous_layer_spec
+from megatron.core.transformer.heterogeneous.heterogeneous_config import HeterogeneousTransformerConfig
 import torch
 from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
 from megatron.core.enums import ModelType
@@ -20,6 +22,9 @@
 from megatron.core.transformer.module import Float16Module, MegatronModule
 from megatron.core.transformer.dot_product_attention import DotProductAttention as MCoreDotProductAttention
 from megatron.core.transformer.enums import AttnBackend
+from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+    get_transformer_block_with_experimental_attention_variant_spec,
+)
 
 from megatron.training.models.base import ModelConfig, ModelBuilder, compose_hooks
 from megatron.training.vocab_utils import calculate_padded_vocab_size
@@ -34,6 +39,9 @@
 from dataclasses import dataclass
 
 from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_decoder_block_spec,
+    get_gpt_decoder_layer_specs,
+    get_gpt_layer_with_inference_spec,
     get_gpt_layer_local_spec,
     get_gpt_layer_with_transformer_engine_spec,
 )
@@ -47,10 +55,17 @@ def transformer_engine_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
     else:
         kwargs = {}
     return get_gpt_layer_with_transformer_engine_spec(
-        num_experts=config.transformer.num_moe_experts,
-        moe_grouped_gemm=config.transformer.moe_grouped_gemm,
-        qk_layernorm=config.transformer.qk_layernorm,
-        fp8=bool(config.transformer.num_moe_experts and (config.transformer.fp8 is not None)),
+        config.transformer.num_moe_experts,
+        config.transformer.moe_grouped_gemm,
+        config.transformer.qk_layernorm,
+        config.transformer.multi_latent_attention,
+        config.transformer.experimental_attention_variant,
+        qk_l2_norm=config.transformer.qk_l2_norm,
+        use_kitchen=config.transformer.use_kitchen,
+        use_te_activation_func=config.transformer.use_te_activation_func,
+        use_kitchen_attention=config.transformer.use_kitchen_attention,
+        kitchen_attention_backend=config.transformer.kitchen_attention_backend,
+        mla_down_proj_fusion=getattr(config.transformer, "mla_down_proj_fusion", False),
         **kwargs,
     )
 
@@ -92,15 +107,45 @@ def modelopt_transformer_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
     )
 
 
-def default_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
+def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec:
     """Determine the most appropriate layer specification based on availability."""
+    transformer_cfg = config.transformer
+    use_te = transformer_cfg.transformer_impl == "transformer_engine"
     if config.restore_modelopt_state:
         return modelopt_transformer_layer_spec(config)
-    # TODO (@maanug): migrate this spec from Megatron-Bridge
-    # elif config.use_transformer_engine_full_layer_spec:
-    #     return transformer_engine_full_layer_spec(config.transformer)
-    else:
+    elif transformer_cfg.experimental_attention_variant is not None:
+        return get_transformer_block_with_experimental_attention_variant_spec(config=transformer_cfg, vp_stage=vp_stage)
+    elif transformer_cfg.num_moe_experts is not None:
+        return get_gpt_decoder_block_spec(
+            transformer_cfg,
+            use_transformer_engine=use_te,
+            normalization=transformer_cfg.normalization,
+            qk_l2_norm=transformer_cfg.qk_l2_norm,
+            vp_stage=vp_stage,
+        )
+    elif isinstance(transformer_cfg, HeterogeneousTransformerConfig):
+        return get_gpt_heterogeneous_layer_spec(transformer_cfg, use_te)
+    elif use_te:
         return transformer_engine_layer_spec(config)
+    elif transformer_cfg.transformer_impl == "inference_optimized":
+        return get_gpt_layer_with_inference_spec(
+            transformer_cfg.qk_layernorm,
+            transformer_cfg.multi_latent_attention,
+            qk_l2_norm=transformer_cfg.qk_l2_norm,
+        )
+    else:
+        return get_gpt_layer_local_spec(
+            transformer_cfg.num_moe_experts,
+            transformer_cfg.moe_grouped_gemm,
+            transformer_cfg.qk_layernorm,
+            transformer_cfg.multi_latent_attention,
+            transformer_cfg.experimental_attention_variant,
+            normalization=transformer_cfg.normalization,
+            use_kitchen=transformer_cfg.use_kitchen,
+            use_kitchen_attention=transformer_cfg.use_kitchen_attention,
+            kitchen_attention_backend=transformer_cfg.kitchen_attention_backend,
+        )
+
 
 
 @dataclass(kw_only=True)
@@ -368,16 +413,20 @@ def mtp_block_spec(
     Returns:
         ModuleSpec: The MTP module specification
     """
+    transformer_cfg = config.transformer
+    use_te = config.transformer.transformer_impl == "transformer_engine"
+
     if config.transformer.mtp_num_layers is not None:
         from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
 
         if hasattr(transformer_layer_spec, "layer_specs") and len(transformer_layer_spec.layer_specs) == 0:
             # Get the decoder layer spec explicitly if no decoder layer in the last stage,
             # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
-            spec = default_layer_spec(config)
+            spec = default_layer_spec(config, vp_stage)
         else:
-            spec = transformer_layer_spec
+            decoder_specs = get_gpt_decoder_layer_specs(transformer_cfg, use_transformer_engine=use_te, normalization=transformer_cfg.normalization, qk_l2_norm=transformer_cfg.qk_l2_norm, vp_stage=vp_stage)
+            spec = decoder_specs[-1]
 
-        return get_gpt_mtp_block_spec(config.transformer, spec, use_transformer_engine=True, vp_stage=vp_stage)
+        return get_gpt_mtp_block_spec(transformer_cfg, spec, use_transformer_engine=use_te, vp_stage=vp_stage)
     else:
         return None

From ae3748266794f4bb4f17a117289157db2f9c5ef3 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 00:54:34 -0400
Subject: [PATCH 09/18] merge spec funcs

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 98 ++++++++++++---------------------
 1 file changed, 36 insertions(+), 62 deletions(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index 9461d467ecd..d8d3cd26511 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -47,72 +47,29 @@
 )
 
 
-
-def transformer_engine_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
-    """Create a Transformer Engine layer specification based on the provided config."""
-    if "use_te_op_fuser" in inspect.signature(get_gpt_layer_with_transformer_engine_spec).parameters:
-        kwargs = {"use_te_op_fuser": config.use_transformer_engine_op_fuser}
-    else:
-        kwargs = {}
-    return get_gpt_layer_with_transformer_engine_spec(
-        config.transformer.num_moe_experts,
-        config.transformer.moe_grouped_gemm,
-        config.transformer.qk_layernorm,
-        config.transformer.multi_latent_attention,
-        config.transformer.experimental_attention_variant,
-        qk_l2_norm=config.transformer.qk_l2_norm,
-        use_kitchen=config.transformer.use_kitchen,
-        use_te_activation_func=config.transformer.use_te_activation_func,
-        use_kitchen_attention=config.transformer.use_kitchen_attention,
-        kitchen_attention_backend=config.transformer.kitchen_attention_backend,
-        mla_down_proj_fusion=getattr(config.transformer, "mla_down_proj_fusion", False),
-        **kwargs,
-    )
-
-
-def local_layer_spec(config: TransformerConfig) -> ModuleSpec:
-    """Create a local layer specification without Transformer Engine.
-
-    Args:
-        config: GPT configuration object
-
-    Returns:
-        ModuleSpec: Module specification for local implementation layers
-    """
-    return get_gpt_layer_local_spec(
-        num_experts=config.num_moe_experts,
-        moe_grouped_gemm=config.moe_grouped_gemm,
-        qk_layernorm=config.qk_layernorm,
-        normalization=config.normalization,
-    )
-
-
-def modelopt_transformer_layer_spec(config: "GPTModelConfig") -> ModuleSpec:
-    """Layer specification for quantization with ModelOpt."""
-    # arbitrary attention mask is used for speculative decoding training
-    # When context parallel > 1, only causal mask type is supported
-    from megatron.core import parallel_state
-
-    use_arbitrary_attention_mask = (
-        config.use_arbitrary_attention_mask
-        if config.use_arbitrary_attention_mask is not None
-        else parallel_state.get_context_parallel_world_size() == 1
-    )
-    return get_gpt_modelopt_spec(
-        config=config.transformer,
-        local_core_attention=False,
-        remap_te_layernorm=True,
-        real_quant_cfg="None",
-        use_arbitrary_attention_mask=use_arbitrary_attention_mask,
-    )
-
-
 def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec:
     """Determine the most appropriate layer specification based on availability."""
     transformer_cfg = config.transformer
     use_te = transformer_cfg.transformer_impl == "transformer_engine"
     if config.restore_modelopt_state:
-        return modelopt_transformer_layer_spec(config)
+        ## Layer specification for quantization with ModelOpt. ##
+
+        # arbitrary attention mask is used for speculative decoding training
+        # When context parallel > 1, only causal mask type is supported
+        from megatron.core import parallel_state
+
+        use_arbitrary_attention_mask = (
+            config.use_arbitrary_attention_mask
+            if config.use_arbitrary_attention_mask is not None
+            else parallel_state.get_context_parallel_world_size() == 1
+        )
+        return get_gpt_modelopt_spec(
+            config=config.transformer,
+            local_core_attention=False,
+            remap_te_layernorm=True,
+            real_quant_cfg="None",
+            use_arbitrary_attention_mask=use_arbitrary_attention_mask,
+        )
     elif transformer_cfg.experimental_attention_variant is not None:
         return get_transformer_block_with_experimental_attention_variant_spec(config=transformer_cfg, vp_stage=vp_stage)
     elif transformer_cfg.num_moe_experts is not None:
@@ -126,7 +83,24 @@ def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec:
     elif isinstance(transformer_cfg, HeterogeneousTransformerConfig):
         return get_gpt_heterogeneous_layer_spec(transformer_cfg, use_te)
     elif use_te:
-        return transformer_engine_layer_spec(config)
+        if "use_te_op_fuser" in inspect.signature(get_gpt_layer_with_transformer_engine_spec).parameters:
+            kwargs = {"use_te_op_fuser": config.use_transformer_engine_op_fuser}
+        else:
+            kwargs = {}
+        return get_gpt_layer_with_transformer_engine_spec(
+            config.transformer.num_moe_experts,
+            config.transformer.moe_grouped_gemm,
+            config.transformer.qk_layernorm,
+            config.transformer.multi_latent_attention,
+            config.transformer.experimental_attention_variant,
+            qk_l2_norm=config.transformer.qk_l2_norm,
+            use_kitchen=config.transformer.use_kitchen,
+            use_te_activation_func=config.transformer.use_te_activation_func,
+            use_kitchen_attention=config.transformer.use_kitchen_attention,
+            kitchen_attention_backend=config.transformer.kitchen_attention_backend,
+            mla_down_proj_fusion=getattr(config.transformer, "mla_down_proj_fusion", False),
+            **kwargs,
+        )
     elif transformer_cfg.transformer_impl == "inference_optimized":
         return get_gpt_layer_with_inference_spec(
             transformer_cfg.qk_layernorm,

From 3377de212dac808e91ced4b37137587e84d7bcd8 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 01:05:33 -0400
Subject: [PATCH 10/18] fix call

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index d8d3cd26511..97c16f43dee 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -264,7 +264,7 @@ def build_model(
         """
         transformer_layer_spec = self._model_config.transformer_layer_spec
         if transformer_layer_spec is None:
-            transformer_layer_spec = default_layer_spec(self._model_config)
+            transformer_layer_spec = default_layer_spec(self._model_config, vp_stage)
         elif callable(transformer_layer_spec):
             # Check if the transformer_layer_spec function accepts vp_stage parameter
             if "vp_stage" in inspect.signature(transformer_layer_spec).parameters:

From dd56e7d78bab5a4c5133ad28df56a5e808e37891 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 01:11:31 -0400
Subject: [PATCH 11/18] add unit tests

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 .../training/models/test_gpt_builder.py       | 888 ++++++++++++++++++
 1 file changed, 888 insertions(+)
 create mode 100644 tests/unit_tests/training/models/test_gpt_builder.py

diff --git a/tests/unit_tests/training/models/test_gpt_builder.py b/tests/unit_tests/training/models/test_gpt_builder.py
new file mode 100644
index 00000000000..5a6e85e7200
--- /dev/null
+++ b/tests/unit_tests/training/models/test_gpt_builder.py
@@ -0,0 +1,888 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import inspect
+from unittest.mock import Mock, call, patch
+
+import pytest
+
+from megatron.core.transformer import ModuleSpec
+from megatron.core.transformer.enums import AttnBackend
+from megatron.core.transformer.heterogeneous.heterogeneous_config import (
+    HeterogeneousTransformerConfig,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training.models.gpt import (
+    GPTModelBuilder,
+    GPTModelConfig,
+    default_layer_spec,
+    mtp_block_spec,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_transformer(**kwargs):
+    defaults = dict(num_layers=2, hidden_size=128, num_attention_heads=1)
+    defaults.update(kwargs)
+    return TransformerConfig(**defaults)
+
+
+def _make_gpt_config(**kwargs):
+    defaults = dict(transformer=_make_transformer(), vocab_size=32000)
+    defaults.update(kwargs)
+    return GPTModelConfig(**defaults)
+
+
+def _make_dispatch_config(**transformer_kwargs):
+    """Build a Mock config for ``default_layer_spec`` dispatch tests.
+
+    Defaults route execution to the ``local`` branch; tests flip individual
+    flags to exercise the other branches.
+    """
+    config = Mock()
+    config.restore_modelopt_state = False
+    config.use_transformer_engine_op_fuser = False
+    config.use_arbitrary_attention_mask = None
+
+    transformer = Mock(spec=TransformerConfig)
+    transformer.transformer_impl = "local"
+    transformer.experimental_attention_variant = None
+    transformer.num_moe_experts = None
+    transformer.moe_grouped_gemm = False
+    transformer.qk_layernorm = False
+    transformer.qk_l2_norm = False
+    transformer.multi_latent_attention = False
+    transformer.normalization = "LayerNorm"
+    transformer.use_kitchen = False
+    transformer.use_kitchen_attention = False
+    transformer.kitchen_attention_backend = None
+    transformer.use_te_activation_func = False
+    transformer.mla_down_proj_fusion = False
+    for k, v in transformer_kwargs.items():
+        setattr(transformer, k, v)
+    config.transformer = transformer
+    return config
+
+
+# =============================================================================
+# Section 1 — default_layer_spec
+# =============================================================================
+
+
+class TestDefaultLayerSpecDispatch:
+    """Tests for ``default_layer_spec`` — branch dispatch on config flags."""
+
+    @patch("megatron.training.models.gpt.get_gpt_modelopt_spec")
+    @patch("megatron.core.parallel_state")
+    def test_returns_modelopt_spec_when_restore_modelopt_state_true(self, mock_ps, mock_get_spec):
+        mock_ps.get_context_parallel_world_size.return_value = 1
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        config = _make_dispatch_config()
+        config.restore_modelopt_state = True
+
+        result = default_layer_spec(config, vp_stage=None)
+
+        mock_get_spec.assert_called_once_with(
+            config=config.transformer,
+            local_core_attention=False,
+            remap_te_layernorm=True,
+            real_quant_cfg="None",
+            use_arbitrary_attention_mask=True,
+        )
+        assert result is spec
+
+    @patch("megatron.training.models.gpt.get_gpt_modelopt_spec")
+    @patch("megatron.core.parallel_state")
+    def test_modelopt_arbitrary_attention_mask_explicit_value(self, mock_ps, mock_get_spec):
+        mock_ps.get_context_parallel_world_size.return_value = 4
+        mock_get_spec.return_value = Mock(spec=ModuleSpec)
+        config = _make_dispatch_config()
+        config.restore_modelopt_state = True
+        config.use_arbitrary_attention_mask = True
+
+        default_layer_spec(config, vp_stage=None)
+
+        assert mock_get_spec.call_args.kwargs["use_arbitrary_attention_mask"] is True
+
+    @patch("megatron.training.models.gpt.get_gpt_modelopt_spec")
+    @patch("megatron.core.parallel_state")
+    def test_modelopt_arbitrary_attention_mask_derived_from_cp_world_size(
+        self, mock_ps, mock_get_spec
+    ):
+        mock_ps.get_context_parallel_world_size.return_value = 4
+        mock_get_spec.return_value = Mock(spec=ModuleSpec)
+        config = _make_dispatch_config()
+        config.restore_modelopt_state = True
+        config.use_arbitrary_attention_mask = None
+
+        default_layer_spec(config, vp_stage=None)
+
+        # cp_world_size=4 → use_arbitrary_attention_mask=False
+        assert mock_get_spec.call_args.kwargs["use_arbitrary_attention_mask"] is False
+
+    @patch(
+        "megatron.training.models.gpt.get_transformer_block_with_experimental_attention_variant_spec"
+    )
+    def test_returns_experimental_spec_when_experimental_attention_variant_set(self, mock_get_spec):
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        config = _make_dispatch_config(experimental_attention_variant="variant_x")
+
+        result = default_layer_spec(config, vp_stage=2)
+
+        mock_get_spec.assert_called_once_with(config=config.transformer, vp_stage=2)
+        assert result is spec
+
+    @patch("megatron.training.models.gpt.get_gpt_decoder_block_spec")
+    def test_returns_decoder_block_spec_when_num_moe_experts_set(self, mock_get_spec):
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        config = _make_dispatch_config(num_moe_experts=4)
+
+        # use_te flag depends on transformer_impl
+        config.transformer.transformer_impl = "transformer_engine"
+
+        result = default_layer_spec(config, vp_stage=1)
+
+        mock_get_spec.assert_called_once_with(
+            config.transformer,
+            use_transformer_engine=True,
+            normalization=config.transformer.normalization,
+            qk_l2_norm=config.transformer.qk_l2_norm,
+            vp_stage=1,
+        )
+        assert result is spec
+
+    @patch("megatron.training.models.gpt.get_gpt_heterogeneous_layer_spec")
+    def test_returns_heterogeneous_spec_when_heterogeneous_config(self, mock_get_spec):
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        # Use spec= so isinstance(transformer, HeterogeneousTransformerConfig) is True
+        transformer = Mock(spec=HeterogeneousTransformerConfig)
+        transformer.transformer_impl = "transformer_engine"
+        transformer.experimental_attention_variant = None
+        transformer.num_moe_experts = None
+        config = Mock()
+        config.restore_modelopt_state = False
+        config.transformer = transformer
+
+        result = default_layer_spec(config, vp_stage=None)
+
+        mock_get_spec.assert_called_once_with(transformer, True)
+        assert result is spec
+
+    @patch("megatron.training.models.gpt.get_gpt_layer_with_transformer_engine_spec")
+    def test_returns_te_spec_when_use_te_true(self, mock_get_spec):
+        mock_get_spec.__signature__ = inspect.Signature()
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        config = _make_dispatch_config()
+        config.transformer.transformer_impl = "transformer_engine"
+
+        result = default_layer_spec(config, vp_stage=None)
+
+        assert mock_get_spec.call_count == 1
+        # use_te_op_fuser is not in (empty) signature, so should NOT appear in kwargs
+        assert "use_te_op_fuser" not in mock_get_spec.call_args.kwargs
+        assert result is spec
+
+    @patch("megatron.training.models.gpt.get_gpt_layer_with_transformer_engine_spec")
+    def test_te_spec_passes_use_te_op_fuser_when_supported(self, mock_get_spec):
+        mock_get_spec.__signature__ = inspect.Signature(
+            [inspect.Parameter("use_te_op_fuser", inspect.Parameter.KEYWORD_ONLY)]
+        )
+        mock_get_spec.return_value = Mock(spec=ModuleSpec)
+        config = _make_dispatch_config()
+        config.transformer.transformer_impl = "transformer_engine"
+        config.use_transformer_engine_op_fuser = True
+
+        default_layer_spec(config, vp_stage=None)
+
+        assert mock_get_spec.call_args.kwargs.get("use_te_op_fuser") is True
+
+    @patch("megatron.training.models.gpt.get_gpt_layer_with_inference_spec")
+    def test_returns_inference_spec_when_inference_optimized(self, mock_get_spec):
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        config = _make_dispatch_config()
+        config.transformer.transformer_impl = "inference_optimized"
+
+        result = default_layer_spec(config, vp_stage=None)
+
+        mock_get_spec.assert_called_once_with(
+            config.transformer.qk_layernorm,
+            config.transformer.multi_latent_attention,
+            qk_l2_norm=config.transformer.qk_l2_norm,
+        )
+        assert result is spec
+
+    @patch("megatron.training.models.gpt.get_gpt_layer_local_spec")
+    def test_returns_local_spec_as_default(self, mock_get_spec):
+        spec = Mock(spec=ModuleSpec)
+        mock_get_spec.return_value = spec
+        config = _make_dispatch_config()
+
+        result = default_layer_spec(config, vp_stage=None)
+
+        mock_get_spec.assert_called_once_with(
+            config.transformer.num_moe_experts,
+            config.transformer.moe_grouped_gemm,
+            config.transformer.qk_layernorm,
+            config.transformer.multi_latent_attention,
+            config.transformer.experimental_attention_variant,
+            normalization=config.transformer.normalization,
+            use_kitchen=config.transformer.use_kitchen,
+            use_kitchen_attention=config.transformer.use_kitchen_attention,
+            kitchen_attention_backend=config.transformer.kitchen_attention_backend,
+        )
+        assert result is spec
+
+
+# =============================================================================
+# Section 2 — GPTModelConfig
+# =============================================================================
+
+
+class TestGPTModelConfigInitialization:
+    """Tests for GPTModelConfig field defaults and custom initialization."""
+
+    def test_builder_classvar(self):
+        assert GPTModelConfig.builder == "megatron.training.models.gpt.GPTModelBuilder"
+
+    def test_default_values(self):
+        config = GPTModelConfig(transformer=_make_transformer())
+        assert config.transformer_layer_spec is None
+        assert config.vocab_size is None
+        assert config.make_vocab_size_divisible_by == 128
+        assert config.should_pad_vocab is False
+        assert config.seq_length == 1024
+        assert config.fp16_lm_cross_entropy is False
+        assert config.parallel_output is True
+        assert config.share_embeddings_and_output_weights is False
+        assert config.position_embedding_type == "learned_absolute"
+        assert config.rotary_percent == 1.0
+        assert config.rotary_base == 10000
+        assert config.rope_scaling is False
+        assert config.rope_scaling_factor == 8.0
+        assert config.scatter_embedding_sequence_parallel is True
+        assert config.seq_len_interpolation_factor is None
+        assert config.tp_comm_overlap_cfg is None
+        assert config.use_transformer_engine_op_fuser is False
+        assert config.use_arbitrary_attention_mask is None
+
+    def test_custom_initialization(self):
+        config = GPTModelConfig(
+            transformer=_make_transformer(),
+            seq_length=4096,
+            fp16_lm_cross_entropy=True,
+            parallel_output=False,
+            share_embeddings_and_output_weights=True,
+            position_embedding_type="rope",
+            rotary_base=500000,
+            vocab_size=50000,
+        )
+        assert config.seq_length == 4096
+        assert config.fp16_lm_cross_entropy is True
+        assert config.parallel_output is False
+        assert config.share_embeddings_and_output_weights is True
+        assert config.position_embedding_type == "rope"
+        assert config.rotary_base == 500000
+        assert config.vocab_size == 50000
+
+
+class TestGPTModelConfigGetAttr:
+    """Tests for GPTModelConfig.__getattr__ — direct access vs. TransformerConfig proxy."""
+
+    def setup_method(self):
+        self.transformer = _make_transformer(hidden_size=256, num_layers=4)
+        self.config = GPTModelConfig(transformer=self.transformer, vocab_size=32000)
+
+    def test_own_attribute_not_proxied(self):
+        assert self.config.vocab_size == 32000
+
+    def test_proxies_transformer_attribute(self):
+        assert self.config.hidden_size == 256
+
+    def test_raises_attribute_error_for_unknown(self):
+        with pytest.raises(AttributeError):
+            _ = self.config.completely_unknown_attr_xyz
+
+    def test_raises_before_transformer_init(self):
+        del self.config.__dict__["transformer"]
+        with pytest.raises(AttributeError):
+            _ = self.config.hidden_size
+
+    def test_error_message_contains_attr_name(self):
+        attr_name = "completely_unknown_attr_xyz"
+        with pytest.raises(AttributeError, match=attr_name):
+            getattr(self.config, attr_name)
+
+
+class TestGPTModelConfigSetAttr:
+    """Tests for GPTModelConfig.__setattr__ — own-field writes vs. TransformerConfig proxy writes."""
+
+    def setup_method(self):
+        self.transformer = _make_transformer(hidden_size=256)
+        self.config = GPTModelConfig(transformer=self.transformer, vocab_size=32000)
+
+    def test_sets_own_attribute_on_self(self):
+        self.config.vocab_size = 50000
+        assert self.config.vocab_size == 50000
+        assert self.config.__dict__.get("vocab_size") == 50000
+
+    def test_proxies_set_to_transformer_attribute(self):
+        self.config.hidden_size = 512
+        assert self.transformer.hidden_size == 512
+
+    def test_set_proxied_attr_reflects_on_transformer(self):
+        self.config.hidden_size = 1024
+        assert self.config.hidden_size == 1024
+        assert self.transformer.hidden_size == 1024
+
+    def test_set_before_transformer_init(self):
+        del self.config.__dict__["transformer"]
+        self.config.vocab_size = 42
+        assert self.config.__dict__["vocab_size"] == 42
+
+    def test_set_transformer_itself_stores_on_self(self):
+        new_transformer = _make_transformer(hidden_size=512)
+        self.config.transformer = new_transformer
+        assert self.config.transformer is new_transformer
+        assert self.config.__dict__["transformer"] is new_transformer
+
+    def test_set_own_attr_does_not_go_to_transformer(self):
+        self.config.vocab_size = 99999
+        assert self.config.__dict__.get("vocab_size") == 99999
+        assert not hasattr(self.config.transformer, "vocab_size")
+
+    def test_proxied_write_does_not_shadow_on_self(self):
+        self.config.hidden_size = 2048
+        assert "hidden_size" not in self.config.__dict__
+
+
+class TestGPTModelConfigFinalize:
+    """Tests for GPTModelConfig.finalize() — validation logic."""
+
+    def test_calls_transformer_finalize_when_present(self):
+        config = _make_gpt_config()
+        # TransformerConfig doesn't have finalize by default; attach a stub.
+        config.transformer.finalize = Mock()
+        config.finalize()
+        config.transformer.finalize.assert_called_once()
+
+    def test_no_finalize_attribute_is_ok(self):
+        config = _make_gpt_config()
+        # Should not raise even when transformer has no finalize()
+        assert not hasattr(config.transformer, "finalize")
+        config.finalize()
+
+    def test_raises_when_cudagraph_without_te_rng_tracker(self):
+        config = _make_gpt_config()
+        config.transformer.cuda_graph_impl = "local"
+        config.transformer.use_te_rng_tracker = False
+        with pytest.raises(AssertionError, match="RNG tracker"):
+            config.finalize()
+
+    def test_no_error_when_cudagraph_with_te_rng_tracker(self):
+        config = _make_gpt_config()
+        config.transformer.cuda_graph_impl = "local"
+        config.transformer.use_te_rng_tracker = True
+        # Should not raise
+        config.finalize()
+
+    def test_vp_size_assertion_fails_on_indivisible_layers(self):
+        # 6 layers, pp=2 → 3 per stage, vp=2 → 3 % 2 != 0 → AssertionError
+        transformer = _make_transformer(
+            num_layers=6, pipeline_model_parallel_size=2, virtual_pipeline_model_parallel_size=2
+        )
+        config = GPTModelConfig(transformer=transformer, vocab_size=32000)
+        with pytest.raises(AssertionError, match="number of model chunks"):
+            config.finalize()
+
+    def test_vp_size_assertion_passes_on_divisible_layers(self):
+        # 8 layers, pp=2 → 4 per stage, vp=2 → 4 % 2 == 0 → OK
+        transformer = _make_transformer(
+            num_layers=8, pipeline_model_parallel_size=2, virtual_pipeline_model_parallel_size=2
+        )
+        config = GPTModelConfig(transformer=transformer, vocab_size=32000)
+        # Should not raise
+        config.finalize()
+
+    def test_vp_size_check_skipped_for_flexible_pp_layout(self):
+        # Indivisible layers, but account_for_embedding_in_pipeline_split=True makes
+        # the layout flexible, so the assertion is skipped.
+        transformer = _make_transformer(
+            num_layers=6,
+            pipeline_model_parallel_size=2,
+            virtual_pipeline_model_parallel_size=2,
+            account_for_embedding_in_pipeline_split=True,
+        )
+        config = GPTModelConfig(transformer=transformer, vocab_size=32000)
+        # Should not raise
+        config.finalize()
+
+
+# =============================================================================
+# Section 3 — GPTModelBuilder
+# =============================================================================
+
+
+class TestGPTModelBuilderInit:
+    """Tests for GPTModelBuilder.__init__ — config storage."""
+
+    def setup_method(self):
+        self.config = _make_gpt_config()
+        self.builder = GPTModelBuilder(self.config)
+
+    def test_stores_model_config(self):
+        assert self.builder._model_config is self.config
+
+
+class TestGPTModelBuilderBuildModel:
+    """Tests for GPTModelBuilder.build_model() — spec resolution, vocab padding,
+    vp/pp-stage inference, attention-backend override, and MCoreGPTModel kwargs.
+    """
+
+    def setup_method(self):
+        self.config = _make_gpt_config(vocab_size=32000)
+        # Provide an explicit ModuleSpec so build_model never hits default_layer_spec.
+        self._default_spec = ModuleSpec(module=object)
+        self.config.__dict__["transformer_layer_spec"] = self._default_spec
+        self.builder = GPTModelBuilder(self.config)
+        self.pg = Mock()
+        self.pg.pp = Mock()
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_raises_when_vocab_size_none(self, mock_model, *_):
+        self.config.__dict__["vocab_size"] = None
+        with pytest.raises(AssertionError, match="vocab_size"):
+            self.builder.build_model(self.pg)
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_spec_already_module_spec_used_directly(self, mock_model, *_):
+        module_spec = ModuleSpec(module=object)
+        self.config.__dict__["transformer_layer_spec"] = module_spec
+        self.builder.build_model(self.pg, pre_process=True, post_process=True)
+        assert mock_model.call_args.kwargs["transformer_layer_spec"] is module_spec
+
+    @patch("megatron.training.models.gpt.default_layer_spec")
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_spec_none_uses_default_layer_spec(self, mock_model, *patches):
+        mock_default = patches[-1]
+        default_spec = ModuleSpec(module=object)
+        mock_default.return_value = default_spec
+        self.config.__dict__["transformer_layer_spec"] = None
+
+        self.builder.build_model(self.pg, pre_process=True, post_process=True, vp_stage=3)
+
+        mock_default.assert_called_once_with(self.config, 3)
+        assert mock_model.call_args.kwargs["transformer_layer_spec"] is default_spec
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_spec_callable_without_vp_stage_param_called_without_it(self, mock_model, *_):
+        returned_spec = ModuleSpec(module=object)
+        calls = []
+
+        def no_vp_fn(config):
+            calls.append(config)
+            return returned_spec
+
+        self.config.__dict__["transformer_layer_spec"] = no_vp_fn
+        self.builder.build_model(self.pg, pre_process=True, post_process=True)
+        assert calls == [self.config]
+        assert mock_model.call_args.kwargs["transformer_layer_spec"] is returned_spec
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_spec_callable_with_vp_stage_param_called_with_it(self, mock_model, *_):
+        returned_spec = ModuleSpec(module=object)
+        received = []
+
+        def vp_fn(config, vp_stage=None):
+            received.append((config, vp_stage))
+            return returned_spec
+
+        self.config.__dict__["transformer_layer_spec"] = vp_fn
+        self.builder.build_model(self.pg, pre_process=True, post_process=True, vp_stage=2)
+        assert received == [(self.config, 2)]
+        assert mock_model.call_args.kwargs["transformer_layer_spec"] is returned_spec
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_no_vocab_padding_uses_vocab_size_directly(self, mock_model, *patches):
+        # The padding mock is the second patch from the bottom (after mtp_block_spec).
+        # Fetch it by traversing in MRO-aware order isn't straightforward; instead,
+        # patch it explicitly within this test for clarity.
+        with patch("megatron.training.models.gpt.calculate_padded_vocab_size") as mock_pad:
+            self.config.__dict__["should_pad_vocab"] = False
+            self.config.__dict__["vocab_size"] = 32000
+            self.builder.build_model(self.pg, pre_process=True, post_process=True)
+            mock_pad.assert_not_called()
+        assert mock_model.call_args.kwargs["vocab_size"] == 32000
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size", return_value=32128)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_vocab_padding_calls_calculate_padded_vocab_size(self, mock_model, mock_pad, *_):
+        self.config.__dict__["should_pad_vocab"] = True
+        self.config.__dict__["vocab_size"] = 32000
+        self.config.__dict__["make_vocab_size_divisible_by"] = 128
+        self.config.transformer.tensor_model_parallel_size = 2
+        self.builder.build_model(self.pg, pre_process=True, post_process=True)
+        mock_pad.assert_called_once_with(32000, 128, 2)
+        assert mock_model.call_args.kwargs["vocab_size"] == 32128
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_explicit_pre_post_process_passed_through(self, mock_model, *_):
+        self.builder.build_model(self.pg, pre_process=False, post_process=True)
+        kw = mock_model.call_args.kwargs
+        assert kw["pre_process"] is False
+        assert kw["post_process"] is True
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=False)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_infers_pre_process_from_vp_and_pp(
+        self, mock_model, mock_pp_first, mock_pp_last, mock_vp_first, mock_vp_last, *_
+    ):
+        self.builder.build_model(self.pg)
+        # is_vp_first_stage returns False → pre_process should be False even if pp_first is True
+        assert mock_model.call_args.kwargs["pre_process"] is False
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_infers_post_process_from_vp_and_pp(self, mock_model, mock_pp_first, mock_pp_last, *_):
+        self.builder.build_model(self.pg)
+        mock_pp_last.assert_called_once_with(self.pg.pp)
+        assert mock_model.call_args.kwargs["post_process"] is True
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_local_attn_backend_overrides_core_attention(self, mock_model, *_):
+        from megatron.core.transformer.dot_product_attention import (
+            DotProductAttention as MCoreDotProductAttention,
+        )
+
+        mock_spec = Mock(spec=ModuleSpec)
+        mock_spec.submodules = Mock()
+        mock_spec.submodules.self_attention = Mock()
+        mock_spec.submodules.self_attention.submodules = Mock()
+        self.config.__dict__["transformer_layer_spec"] = mock_spec
+        self.config.transformer.attention_backend = AttnBackend.local
+
+        self.builder.build_model(self.pg, pre_process=True, post_process=True)
+
+        assert (
+            mock_spec.submodules.self_attention.submodules.core_attention
+            is MCoreDotProductAttention
+        )
+
+    @patch("megatron.training.models.gpt.mtp_block_spec")
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_mtp_block_spec_passed_to_mcore(self, mock_model, *patches):
+        mock_mtp = patches[-1]
+        mtp_spec = ModuleSpec(module=object)
+        mock_mtp.return_value = mtp_spec
+
+        self.builder.build_model(self.pg, pre_process=True, post_process=True, vp_stage=1)
+
+        # mtp_block_spec is called with (config, transformer_layer_spec, vp_stage=vp_stage)
+        mock_mtp.assert_called_once_with(self.config, self._default_spec, vp_stage=1)
+        assert mock_model.call_args.kwargs["mtp_block_spec"] is mtp_spec
+
+    @patch("megatron.training.models.gpt.mtp_block_spec", return_value=None)
+    @patch("megatron.training.models.gpt.calculate_padded_vocab_size")
+    @patch("megatron.training.models.gpt.is_vp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_vp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_last_stage", return_value=True)
+    @patch("megatron.training.models.gpt.is_pp_first_stage", return_value=True)
+    @patch("megatron.training.models.gpt.GPTModel")
+    def test_config_params_passed_to_mcore(self, mock_model, *_):
+        config = _make_gpt_config(
+            vocab_size=32000,
+            seq_length=4096,
+            fp16_lm_cross_entropy=True,
+            parallel_output=False,
+            share_embeddings_and_output_weights=True,
+            position_embedding_type="rope",
+            rotary_percent=0.5,
+            rotary_base=500000,
+            rope_scaling=True,
+            rope_scaling_factor=4.0,
+            scatter_embedding_sequence_parallel=False,
+        )
+        config.__dict__["transformer_layer_spec"] = ModuleSpec(module=object)
+        builder = GPTModelBuilder(config)
+        pg = Mock()
+        pg.pp = Mock()
+
+        builder.build_model(pg, pre_process=True, post_process=True)
+
+        kw = mock_model.call_args.kwargs
+        assert kw["config"] is config.transformer
+        assert kw["vocab_size"] == 32000
+        assert kw["max_sequence_length"] == 4096
+        assert kw["fp16_lm_cross_entropy"] is True
+        assert kw["parallel_output"] is False
+        assert kw["share_embeddings_and_output_weights"] is True
+        assert kw["position_embedding_type"] == "rope"
+        assert kw["rotary_percent"] == 0.5
+        assert kw["rotary_base"] == 500000
+        assert kw["rope_scaling"] is True
+        assert kw["rope_scaling_factor"] == 4.0
+        assert kw["seq_len_interpolation_factor"] is None
+        assert kw["scatter_embedding_sequence_parallel"] is False
+        assert kw["pre_process"] is True
+        assert kw["post_process"] is True
+        assert kw["pg_collection"] is pg
+        assert kw["vp_stage"] is None
+
+
+class TestGPTModelBuilderBuildDistributedModels:
+    """Tests for GPTModelBuilder.build_distributed_models() — delegation to
+    unimodal helper, hook composition, and default kwargs.
+    """
+
+    def setup_method(self):
+        self.config = _make_gpt_config(vocab_size=32000)
+        self.builder = GPTModelBuilder(self.config)
+        self.pg = Mock()
+
+    @patch("megatron.training.models.gpt.compose_hooks")
+    @patch("megatron.training.models.gpt.unimodal_build_distributed_models")
+    def test_delegates_to_unimodal_build_distributed_models(self, mock_unimodal, mock_compose):
+        mock_unimodal.return_value = [Mock()]
+        mock_compose.return_value = Mock(return_value=None)
+
+        self.builder.build_distributed_models(self.pg)
+
+        assert mock_unimodal.called
+
+    @patch("megatron.training.models.gpt.compose_hooks")
+    @patch("megatron.training.models.gpt.unimodal_build_distributed_models")
+    def test_returns_model_list_from_unimodal(self, mock_unimodal, mock_compose):
+        model_list = [Mock(), Mock()]
+        mock_unimodal.return_value = model_list
+        # post_wrap hook returns None → original list kept
+        mock_compose.return_value = Mock(return_value=None)
+
+        result = self.builder.build_distributed_models(self.pg)
+
+        assert result is model_list
+
+    @patch("megatron.training.models.gpt.compose_hooks")
+    @patch("megatron.training.models.gpt.unimodal_build_distributed_models")
+    def test_pre_wrap_hooks_composed_and_passed(self, mock_unimodal, mock_compose):
+        mock_unimodal.return_value = [Mock()]
+        composed_pre = Mock()
+        composed_post = Mock(return_value=None)
+        mock_compose.side_effect = [composed_pre, composed_post]
+
+        hook1 = Mock()
+        self.config.pre_wrap_hooks = [hook1]
+        self.builder.build_distributed_models(self.pg)
+
+        # First compose_hooks call must be with the pre_wrap_hooks list
+        assert mock_compose.call_args_list[0] == call([hook1])
+        # The composed pre-wrap hook is the 11th positional arg (index 10)
+        unimodal_args = mock_unimodal.call_args.args
+        assert unimodal_args[10] is composed_pre
+
+    @patch("megatron.training.models.gpt.compose_hooks")
+    @patch("megatron.training.models.gpt.unimodal_build_distributed_models")
+    def test_post_wrap_hook_applied_to_results(self, mock_unimodal, mock_compose):
+        model_list = [Mock()]
+        wrapped_list = [Mock(), Mock()]
+        mock_unimodal.return_value = model_list
+        composed_pre = Mock()
+        composed_post = Mock(return_value=wrapped_list)
+        mock_compose.side_effect = [composed_pre, composed_post]
+
+        result = self.builder.build_distributed_models(self.pg)
+
+        composed_post.assert_called_once_with(model_list)
+        assert result is wrapped_list
+
+    @patch("megatron.training.models.gpt.compose_hooks")
+    @patch("megatron.training.models.gpt.unimodal_build_distributed_models")
+    def test_post_wrap_hook_returning_none_keeps_original_list(self, mock_unimodal, mock_compose):
+        model_list = [Mock()]
+        mock_unimodal.return_value = model_list
+        mock_compose.return_value = Mock(return_value=None)
+
+        result = self.builder.build_distributed_models(self.pg)
+
+        assert result is model_list
+
+    @patch("megatron.training.models.gpt.compose_hooks")
+    @patch("megatron.training.models.gpt.unimodal_build_distributed_models")
+    def test_default_parameters_forwarded(self, mock_unimodal, mock_compose):
+        from megatron.core.enums import ModelType
+        from megatron.core.transformer.module import Float16Module
+
+        mock_unimodal.return_value = [Mock()]
+        mock_compose.return_value = Mock(return_value=None)
+
+        self.builder.build_distributed_models(self.pg)
+
+        args = mock_unimodal.call_args.args
+        assert args[0] == self.builder.build_model
+        assert args[1] is self.config.transformer
+        assert args[2] is self.pg
+        assert args[3] is None  # ddp_config
+        assert args[7] is True  # wrap_with_ddp
+        assert args[8] is True  # data_parallel_random_init
+        assert args[9] is Float16Module  # mixed_precision_wrapper
+        assert args[11] is ModelType.encoder_or_decoder  # model_type
+
+
+# =============================================================================
+# Section 4 — mtp_block_spec
+# =============================================================================
+
+
+class TestMtpBlockSpec:
+    """Tests for mtp_block_spec() helper function."""
+
+    def _make_config(self, mtp_num_layers=None, transformer_impl="transformer_engine"):
+        config = Mock()
+        config.transformer.mtp_num_layers = mtp_num_layers
+        config.transformer.transformer_impl = transformer_impl
+        config.transformer.normalization = "LayerNorm"
+        config.transformer.qk_l2_norm = False
+        return config
+
+    def test_returns_none_when_mtp_num_layers_is_none(self):
+        config = self._make_config(mtp_num_layers=None)
+        spec = ModuleSpec(module=object)
+        result = mtp_block_spec(config, spec)
+        assert result is None
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_uses_explicit_spec_when_layer_specs_nonempty(self, mock_get_mtp):
+        config = self._make_config(mtp_num_layers=1)
+        spec = Mock(spec=ModuleSpec)
+        spec.layer_specs = [Mock()]  # Non-empty
+        mock_get_mtp.return_value = Mock(spec=ModuleSpec)
+
+        with patch(
+            "megatron.training.models.gpt.get_gpt_decoder_layer_specs"
+        ) as mock_decoder_specs:
+            mock_decoder_specs.return_value = [Mock(), Mock()]
+            mtp_block_spec(config, spec)
+
+        # When layer_specs is non-empty, use the last decoder spec (not the explicit spec arg)
+        passed_spec = mock_get_mtp.call_args.args[1]
+        assert passed_spec is mock_decoder_specs.return_value[-1]
+
+    @patch("megatron.training.models.gpt.default_layer_spec")
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_uses_default_layer_spec_for_empty_layer_specs(self, mock_get_mtp, mock_default):
+        config = self._make_config(mtp_num_layers=1)
+        spec = Mock(spec=ModuleSpec)
+        spec.layer_specs = []  # Empty → falls back to default_layer_spec
+        fallback_spec = Mock(spec=ModuleSpec)
+        mock_default.return_value = fallback_spec
+        mock_get_mtp.return_value = Mock(spec=ModuleSpec)
+
+        mtp_block_spec(config, spec, vp_stage=4)
+
+        mock_default.assert_called_once_with(config, 4)
+        passed_spec = mock_get_mtp.call_args.args[1]
+        assert passed_spec is fallback_spec
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_passes_vp_stage_and_use_te_to_get_gpt_mtp_block_spec(self, mock_get_mtp):
+        config = self._make_config(mtp_num_layers=2, transformer_impl="transformer_engine")
+        spec = ModuleSpec(module=object)
+        mock_get_mtp.return_value = Mock(spec=ModuleSpec)
+
+        with patch(
+            "megatron.training.models.gpt.get_gpt_decoder_layer_specs"
+        ) as mock_decoder_specs:
+            mock_decoder_specs.return_value = [Mock(), Mock()]
+            mtp_block_spec(config, spec, vp_stage=3)
+
+        call_kwargs = mock_get_mtp.call_args.kwargs
+        assert call_kwargs["use_transformer_engine"] is True
+        assert call_kwargs["vp_stage"] == 3
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_use_transformer_engine_false_when_impl_not_te(self, mock_get_mtp):
+        config = self._make_config(mtp_num_layers=2, transformer_impl="local")
+        spec = ModuleSpec(module=object)
+        mock_get_mtp.return_value = Mock(spec=ModuleSpec)
+
+        with patch(
+            "megatron.training.models.gpt.get_gpt_decoder_layer_specs"
+        ) as mock_decoder_specs:
+            mock_decoder_specs.return_value = [Mock(), Mock()]
+            mtp_block_spec(config, spec)
+
+        assert mock_get_mtp.call_args.kwargs["use_transformer_engine"] is False

From be89d55b8d670362dc1d788bd10d50467374550f Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 01:24:11 -0400
Subject: [PATCH 12/18] fix bug

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index 97c16f43dee..af7c0638d28 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -265,7 +265,7 @@ def build_model(
         transformer_layer_spec = self._model_config.transformer_layer_spec
         if transformer_layer_spec is None:
             transformer_layer_spec = default_layer_spec(self._model_config, vp_stage)
-        elif callable(transformer_layer_spec):
+        elif not isinstance(transformer_layer_spec, ModuleSpec) and callable(transformer_layer_spec):
             # Check if the transformer_layer_spec function accepts vp_stage parameter
             if "vp_stage" in inspect.signature(transformer_layer_spec).parameters:
                 transformer_layer_spec = transformer_layer_spec(self._model_config, vp_stage=vp_stage)

From 64b45ecba8a077b3ca7c00d037dc4e3e76faf455 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 01:24:51 -0400
Subject: [PATCH 13/18] fix tests

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 tests/unit_tests/training/models/test_gpt_builder.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/training/models/test_gpt_builder.py b/tests/unit_tests/training/models/test_gpt_builder.py
index 5a6e85e7200..48b36bc379d 100644
--- a/tests/unit_tests/training/models/test_gpt_builder.py
+++ b/tests/unit_tests/training/models/test_gpt_builder.py
@@ -4,6 +4,7 @@
 from unittest.mock import Mock, call, patch
 
 import pytest
+import torch
 
 from megatron.core.transformer import ModuleSpec
 from megatron.core.transformer.enums import AttnBackend
@@ -396,7 +397,10 @@ def test_no_error_when_cudagraph_with_te_rng_tracker(self):
     def test_vp_size_assertion_fails_on_indivisible_layers(self):
         # 6 layers, pp=2 → 3 per stage, vp=2 → 3 % 2 != 0 → AssertionError
         transformer = _make_transformer(
-            num_layers=6, pipeline_model_parallel_size=2, virtual_pipeline_model_parallel_size=2
+            num_layers=6,
+            pipeline_model_parallel_size=2,
+            virtual_pipeline_model_parallel_size=2,
+            pipeline_dtype=torch.bfloat16,
         )
         config = GPTModelConfig(transformer=transformer, vocab_size=32000)
         with pytest.raises(AssertionError, match="number of model chunks"):
@@ -405,7 +409,10 @@ def test_vp_size_assertion_fails_on_indivisible_layers(self):
     def test_vp_size_assertion_passes_on_divisible_layers(self):
         # 8 layers, pp=2 → 4 per stage, vp=2 → 4 % 2 == 0 → OK
         transformer = _make_transformer(
-            num_layers=8, pipeline_model_parallel_size=2, virtual_pipeline_model_parallel_size=2
+            num_layers=8,
+            pipeline_model_parallel_size=2,
+            virtual_pipeline_model_parallel_size=2,
+            pipeline_dtype=torch.bfloat16,
         )
         config = GPTModelConfig(transformer=transformer, vocab_size=32000)
         # Should not raise
@@ -419,6 +426,7 @@ def test_vp_size_check_skipped_for_flexible_pp_layout(self):
             pipeline_model_parallel_size=2,
             virtual_pipeline_model_parallel_size=2,
             account_for_embedding_in_pipeline_split=True,
+            pipeline_dtype=torch.bfloat16,
         )
         config = GPTModelConfig(transformer=transformer, vocab_size=32000)
         # Should not raise

From 66304bb6fc5828df55c31da3af65f62980de6570 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Sun, 17 May 2026 01:27:43 -0400
Subject: [PATCH 14/18] fix tests

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 tests/unit_tests/training/models/test_gpt_builder.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/unit_tests/training/models/test_gpt_builder.py b/tests/unit_tests/training/models/test_gpt_builder.py
index 48b36bc379d..20603e780b7 100644
--- a/tests/unit_tests/training/models/test_gpt_builder.py
+++ b/tests/unit_tests/training/models/test_gpt_builder.py
@@ -419,10 +419,12 @@ def test_vp_size_assertion_passes_on_divisible_layers(self):
         config.finalize()
 
     def test_vp_size_check_skipped_for_flexible_pp_layout(self):
-        # Indivisible layers, but account_for_embedding_in_pipeline_split=True makes
-        # the layout flexible, so the assertion is skipped.
+        # num_layers=3, pp=2, vp=2: (3 // 2) % 2 = 1 ≠ 0, so the gpt.py vp-size
+        # assertion would fail. But account_for_embedding_in_pipeline_split=True
+        # marks the layout as flexible, so the check is skipped.
+        # (num_layers + 1 = 4 is divisible by pp_size=2, so TransformerConfig accepts.)
         transformer = _make_transformer(
-            num_layers=6,
+            num_layers=3,
             pipeline_model_parallel_size=2,
             virtual_pipeline_model_parallel_size=2,
             account_for_embedding_in_pipeline_split=True,

From bb5af57fd603cc3efe61b151937c4ddc0c65fcd7 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Thu, 21 May 2026 14:49:13 -0700
Subject: [PATCH 15/18] reapply lost fixes for agent comments

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/__init__.py | 2 +-
 megatron/training/models/gpt.py      | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/megatron/training/models/__init__.py b/megatron/training/models/__init__.py
index 39e52b0a371..151fab1addc 100644
--- a/megatron/training/models/__init__.py
+++ b/megatron/training/models/__init__.py
@@ -23,5 +23,5 @@
     "MambaModelConfig",
     "MambaModelBuilder",
     "GPTModelConfig",
-    "GPTModelBuilder"
+    "GPTModelBuilder",
 ]
diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index af7c0638d28..92c973b0aef 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -179,10 +179,10 @@ def __getattr__(self, name: str, /) -> Any:
         try:
             transformer = object.__getattribute__(self, "transformer")
         except AttributeError:
-            raise AttributeError(f"MambaModelConfig has no attribute '{name}'")
+            raise AttributeError(f"GPTModelConfig has no attribute '{name}'")
         if hasattr(transformer, name):
             return getattr(transformer, name)
-        raise AttributeError(f"Neither MambaModelConfig nor TransformerConfig has any attribute '{name}'.")
+        raise AttributeError(f"Neither GPTModelConfig nor TransformerConfig has any attribute '{name}'.")
 
     @override
     def __setattr__(self, name: str, value: Any, /) -> None:

From aaa76dceecf6f3b136ae57669f81228c6d5d38b5 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Mon, 1 Jun 2026 12:30:29 -0700
Subject: [PATCH 16/18] revise comments

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/argument_utils.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py
index b2051e81360..2cfb3f0f17b 100644
--- a/megatron/training/argument_utils.py
+++ b/megatron/training/argument_utils.py
@@ -418,12 +418,13 @@ def gpt_config_from_args(args: Namespace, config: TransformerConfig | None=None)
     kwargs["seq_length"] = args.max_position_embeddings
     kwargs["share_embeddings_and_output_weights"] = not args.untie_embeddings_and_output_weights
 
+    # GPTModelConfig supports either automatically padding vocab size or using exact provided
+    # vocab size via "should_pad_vocab" to support loading third-party checkpoints. Here,
+    # that is just mapped to settings in args appropriately.
     if args.padded_vocab_size is not None:
         kwargs["vocab_size"] = args.padded_vocab_size
         kwargs["should_pad_vocab"] = False
     else:
-        # Megatron-Bridge uses an explicit setting "should_pad_vocab" so that 
-        # when converting model configs from HF, we can set a vocab size and disable padding. 
         assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified."
         kwargs["vocab_size"] = args.vocab_size
         kwargs["should_pad_vocab"] = True
@@ -460,12 +461,13 @@ def hybrid_config_from_args(args: Namespace, config: TransformerConfig | None=No
     kwargs["seq_length"] = args.max_position_embeddings
     kwargs["share_embeddings_and_output_weights"] = not args.untie_embeddings_and_output_weights
 
+    # HybridModelConfig supports either automatically padding vocab size or using exact provided
+    # vocab size via "should_pad_vocab" to support loading third-party checkpoints. Here,
+    # that is just mapped to settings in args appropriately.
     if args.padded_vocab_size is not None:
         kwargs["vocab_size"] = args.padded_vocab_size
         kwargs["should_pad_vocab"] = False
     else:
-        # Megatron-Bridge uses an explicit setting "should_pad_vocab" so that 
-        # when converting model configs from HF, we can set a vocab size and disable padding. 
         assert args.vocab_size is not None, "Either --padded-vocab-size or --vocab-size must be specified."
         kwargs["vocab_size"] = args.vocab_size
         kwargs["should_pad_vocab"] = True

From ca0fdb9b3d3ac08037568161e34ce31fe1a5671f Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Mon, 1 Jun 2026 13:23:03 -0700
Subject: [PATCH 17/18] Revert "move assertion"

This reverts commit bc2ea8fb1196f4b6555d9241a46903d73e038588.
---
 .../core/transformer/heterogeneous/heterogeneous_config.py  | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/megatron/core/transformer/heterogeneous/heterogeneous_config.py b/megatron/core/transformer/heterogeneous/heterogeneous_config.py
index 109eb67524c..550e7b67560 100644
--- a/megatron/core/transformer/heterogeneous/heterogeneous_config.py
+++ b/megatron/core/transformer/heterogeneous/heterogeneous_config.py
@@ -201,12 +201,6 @@ class HeterogeneousTransformerConfig(TransformerConfig):
     def __post_init__(self):
         super().__post_init__()
 
-        err_msg = (
-            "Cannot use inference optimized transformer implementation "
-            "with heterogeneous transformer layers."
-        )
-        assert not (self.transformer_impl == "inference_optimized"), err_msg
-
         self.heterogeneous_block_specs = True
 
         if self.heterogeneous_layers_config_encoded_json in ("", None):

From 511ba6e2a4df05df64a04d3073fa991818c83dc5 Mon Sep 17 00:00:00 2001
From: Maanu Grover <maanug@nvidia.com>
Date: Mon, 1 Jun 2026 13:24:50 -0700
Subject: [PATCH 18/18] revise priority in default layer spec helper

Signed-off-by: Maanu Grover <maanug@nvidia.com>
---
 megatron/training/models/gpt.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py
index 92c973b0aef..633b2ad0b27 100644
--- a/megatron/training/models/gpt.py
+++ b/megatron/training/models/gpt.py
@@ -51,7 +51,13 @@ def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec:
     """Determine the most appropriate layer specification based on availability."""
     transformer_cfg = config.transformer
     use_te = transformer_cfg.transformer_impl == "transformer_engine"
-    if config.restore_modelopt_state:
+    if transformer_cfg.transformer_impl == "inference_optimized":
+        return get_gpt_layer_with_inference_spec(
+            transformer_cfg.qk_layernorm,
+            transformer_cfg.multi_latent_attention,
+            qk_l2_norm=transformer_cfg.qk_l2_norm,
+        )
+    elif config.restore_modelopt_state:
         ## Layer specification for quantization with ModelOpt. ##
 
         # arbitrary attention mask is used for speculative decoding training
@@ -101,12 +107,6 @@ def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec:
             mla_down_proj_fusion=getattr(config.transformer, "mla_down_proj_fusion", False),
             **kwargs,
         )
-    elif transformer_cfg.transformer_impl == "inference_optimized":
-        return get_gpt_layer_with_inference_spec(
-            transformer_cfg.qk_layernorm,
-            transformer_cfg.multi_latent_attention,
-            qk_l2_norm=transformer_cfg.qk_l2_norm,
-        )
     else:
         return get_gpt_layer_local_spec(
             transformer_cfg.num_moe_experts,