Ensure type-checker understands use of Submodules in unit tests

tests/functional_tests/test_cases/common/ckpt_converter/__main__.py

@@ -8,7 +8,6 @@
 import types
 import typing as T
 from collections import namedtuple
-from copy import deepcopy
 from functools import partial
 
 import numpy as np
@@ -19,14 +18,10 @@
 from megatron.core import parallel_state
 from megatron.core.datasets.gpt_dataset import _get_ltor_masks_and_position_ids
 from megatron.core.enums import ModelType
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
-from megatron.core.models.multimodal.llava_model import DEFAULT_IMAGE_TOKEN_INDEX, LLaVAModel
-from megatron.core.models.vision.vit_layer_specs import get_vit_layer_with_transformer_engine_spec
+from megatron.core.models.multimodal.llava_model import DEFAULT_IMAGE_TOKEN_INDEX
 from megatron.core.pipeline_parallel import get_forward_backward_func
 from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.utils import get_attr_wrapped_model
 from megatron.training import get_args, get_tokenizer
 from megatron.training.arguments import parse_args, validate_args
 from megatron.training.checkpointing import load_checkpoint as _load_checkpoint

tests/functional_tests/test_cases/common/moe_perf/__main__.py

@@ -8,18 +8,16 @@
 import os
 import statistics
 from contextlib import nullcontext
-from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, cast
+from typing import Any, Dict, Mapping, Sequence, cast
 
 import pytest  # type: ignore[import]
 import torch
 
 from megatron.core.config import set_experimental_flag
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.models.gpt.gpt_layer_specs import (
-    get_gpt_layer_local_spec,
-    get_gpt_layer_with_transformer_engine_spec,
+    get_gpt_layer_with_transformer_engine_submodules,
 )
 from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP, HAVE_HYBRIDEP
 from megatron.core.transformer.moe.moe_layer import MoELayer
@@ -89,10 +87,9 @@ def _build_transformer_config(case: MoEPerformanceCase) -> TransformerConfig:
 
 # NOTE: Only TE backend is covered in this test.
 def _resolve_moe_submodules(case: MoEPerformanceCase):
-    layer_spec = get_gpt_layer_with_transformer_engine_spec(
+    return get_gpt_layer_with_transformer_engine_submodules(
         num_experts=case.model.num_experts, moe_grouped_gemm=True
-    )
-    return layer_spec.submodules.mlp.submodules
+    ).mlp.submodules
 
 
 def _load_baselines() -> Dict[str, Dict[str, float]]:

tests/unit_tests/dist_checkpointing/models/test_mlp_glu.py

@@ -13,7 +13,9 @@
     get_param_id_to_sharded_param_map,
     optim_state_to_sharding_state,
 )
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_with_transformer_engine_submodules,
+)
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.mlp import MLP, apply_swiglu_sharded_factory
 from megatron.core.transformer.transformer_config import TransformerConfig
@@ -32,7 +34,7 @@ def initialize_mlp(glu=True):
         gated_linear_unit=glu,
     )
     return MLP(
-        transformer_config, get_gpt_layer_with_transformer_engine_spec().submodules.mlp.submodules
+        transformer_config, get_gpt_layer_with_transformer_engine_submodules().mlp.submodules
     )
 
 

tests/unit_tests/dist_checkpointing/models/test_moe_experts.py

@@ -1,6 +1,4 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-import os
-
 import pytest
 import torch
 from transformer_engine.pytorch.fp8 import check_fp8_support, fp8_autocast
@@ -17,11 +15,18 @@
     FullyParallelSaveStrategyWrapper,
 )
 from megatron.core.models.gpt.gpt_layer_specs import (
-    get_gpt_layer_local_spec,
-    get_gpt_layer_with_transformer_engine_spec,
+    get_gpt_layer_local_submodules,
+    get_gpt_layer_with_transformer_engine_submodules,
 )
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP, TEGroupedMLP
+from megatron.core.transformer.mlp import MLPSubmodules
+from megatron.core.transformer.moe.experts import (
+    GroupedMLP,
+    SequentialMLP,
+    TEGroupedMLP,
+    TEGroupedMLPSubmodules,
+)
+from megatron.core.transformer.moe.moe_layer import MoESubmodules
 from megatron.core.transformer.moe.moe_utils import get_default_pg_collection
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
@@ -54,33 +59,41 @@ def initialize_expert_layer(seed, glu=True, expert_type='sequential', fp8=False,
     if expert_type == 'grouped':
         model = GroupedMLP(num_local_experts, transformer_config, pg_collection)
     elif expert_type == 'te_grouped':
-        transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+        layer_submodules = get_gpt_layer_with_transformer_engine_submodules(
             num_experts=num_moe_experts, moe_grouped_gemm=True
         )
+        assert isinstance(layer_submodules.mlp.submodules, MoESubmodules)
+        assert isinstance(
+            layer_submodules.mlp.submodules.experts.submodules, TEGroupedMLPSubmodules
+        )
         model = TEGroupedMLP(
             num_local_experts,
             transformer_config,
-            transformer_layer_spec.submodules.mlp.submodules.experts.submodules,
+            layer_submodules.mlp.submodules.experts.submodules,
             pg_collection,
         )
     elif expert_type == 'sequential':
-        transformer_layer_spec = get_gpt_layer_local_spec(
+        layer_submodules = get_gpt_layer_local_submodules(
             num_experts=num_moe_experts, moe_grouped_gemm=False
         )
+        assert isinstance(layer_submodules.mlp.submodules, MoESubmodules)
+        assert isinstance(layer_submodules.mlp.submodules.experts.submodules, MLPSubmodules)
         model = SequentialMLP(
             num_local_experts,
             transformer_config,
-            transformer_layer_spec.submodules.mlp.submodules.experts.submodules,
+            layer_submodules.mlp.submodules.experts.submodules,
             pg_collection,
         )
     elif expert_type == 'te_sequential':
-        transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+        layer_submodules = get_gpt_layer_with_transformer_engine_submodules(
             num_experts=num_moe_experts, moe_grouped_gemm=False
         )
+        assert isinstance(layer_submodules.mlp.submodules, MoESubmodules)
+        assert isinstance(layer_submodules.mlp.submodules.experts.submodules, MLPSubmodules)
         model = SequentialMLP(
             num_local_experts,
             transformer_config,
-            transformer_layer_spec.submodules.mlp.submodules.experts.submodules,
+            layer_submodules.mlp.submodules.experts.submodules,
             pg_collection,
         )
     else:

tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py

@@ -8,11 +8,12 @@
 
 from megatron.core import parallel_state
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
-from megatron.core.distributed.param_and_grad_buffer import partition_buckets
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_with_transformer_engine_submodules,
+)
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.moe.moe_layer import MoELayer
-from tests.unit_tests.test_utilities import TestModel, Utils
+from tests.unit_tests.test_utilities import Utils
 
 
 class TestMoEModel(torch.nn.Module):
@@ -41,15 +42,13 @@ def __init__(
             params_dtype=torch.bfloat16,
             add_bias_linear=False,
         )
-        transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+        submodules = get_gpt_layer_with_transformer_engine_submodules(
             num_experts=num_moe_experts, moe_grouped_gemm=moe_grouped_gemm
         )
         super().__init__()
         self.layers = torch.nn.ModuleList(
             [
-                MoELayer(
-                    transformer_config, transformer_layer_spec.submodules.mlp.submodules
-                ).cuda()
+                MoELayer(transformer_config, submodules.mlp.submodules).cuda()
                 for _ in range(num_layers)
             ]
         )

tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py

@@ -1,7 +1,5 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from argparse import Namespace
-
 import pytest
 import torch
 
@@ -10,10 +8,7 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.models.gpt.gpt_layer_specs import (
-    get_gpt_layer_local_spec,
-    get_gpt_layer_with_transformer_engine_spec,
-)
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.transformer_config import TransformerConfig

tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py

@@ -1,13 +1,11 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
-import os
 import random
 import string
 import time
-from argparse import Namespace
 from collections import OrderedDict
-from typing import Dict
+from typing import Dict, List
 from unittest import mock
 
 import pytest
@@ -22,12 +20,13 @@
 from megatron.core.inference.text_generation_controllers.vlm_text_generation_controller import (
     VLMTextGenerationController,
 )
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules
 from megatron.core.models.multimodal.llava_model import LLaVAModel
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer.enums import AttnBackend
 from megatron.core.transformer.module import Float16Module
+from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import TransformerLayer
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -68,15 +67,19 @@ def setup_method(self, method):
             bf16=True,
         )
 
-        language_layer_spec = get_gpt_layer_local_spec()
-        vision_layer_spec = copy.deepcopy(language_layer_spec)
-        vision_projection_spec = copy.deepcopy(language_layer_spec.submodules.mlp.submodules)
+        language_layer_submodules = get_gpt_layer_local_submodules()
+        vision_layer_spec = ModuleSpec(
+            module=TransformerLayer, submodules=copy.deepcopy(language_layer_submodules)
+        )
+        vision_projection_spec = copy.deepcopy(language_layer_submodules.mlp.submodules)
 
         language_config.language_model_type = "dummy"
         vision_config.vision_model_type = "clip"
         self.model = LLaVAModel(
             language_transformer_config=language_config,
-            language_transformer_layer_spec=language_layer_spec,
+            language_transformer_layer_spec=ModuleSpec(
+                module=TransformerLayer, submodules=language_layer_submodules
+            ),
             language_vocab_size=self.language_vocab_size,
             language_max_sequence_length=self.language_max_sequence_length,
             vision_transformer_config=vision_config,

tests/unit_tests/models/test_bert_model.py

@@ -1,21 +1,22 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
 import os
-from importlib.metadata import version
 
 import pytest
 import torch
 from packaging.version import Version as PkgVersion
-from pytest_mock import mocker
 
 from megatron.core.models.bert.bert_layer_specs import (
     bert_layer_local_spec,
-    bert_layer_with_transformer_engine_spec,
+    get_bert_layer_with_transformer_engine_spec,
+    get_bert_layer_with_transformer_engine_submodules,
 )
 from megatron.core.models.bert.bert_model import BertModel
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.enums import AttnBackend, AttnMaskType
+from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import TransformerLayer
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -40,7 +41,7 @@ def setup_method(self, method):
         self.bert_model = BertModel(
             config=transformer_config,
             num_tokentypes=0,
-            transformer_layer_spec=bert_layer_with_transformer_engine_spec,
+            transformer_layer_spec=get_bert_layer_with_transformer_engine_spec(),
             vocab_size=100,
             max_sequence_length=4,
         )
@@ -112,7 +113,7 @@ def setup_method(self, method):
         self.bert_model = BertModel(
             config=self.transformer_config,
             num_tokentypes=0,
-            transformer_layer_spec=bert_layer_with_transformer_engine_spec,
+            transformer_layer_spec=get_bert_layer_with_transformer_engine_spec(),
             vocab_size=100,
             max_sequence_length=4,
         )
@@ -139,16 +140,15 @@ def test_local_spec_exception(self, mocker):
 
     @pytest.mark.internal
     def test_transformer_engine_version_1_10(self, mocker):
-        bert_layer_with_transformer_engine_spec.submodules.self_attention.params[
-            'attn_mask_type'
-        ] == AttnMaskType.arbitrary
+        submodules = get_bert_layer_with_transformer_engine_submodules()
+        submodules.self_attention.params['attn_mask_type'] = AttnMaskType.arbitrary
 
         mocker.patch("megatron.core.utils.get_te_version", return_value=PkgVersion("1.10"))
-        self.bert_model.transformer_layer_spec = bert_layer_with_transformer_engine_spec
+        self.bert_model.transformer_layer_spec = ModuleSpec(
+            module=TransformerLayer, submodules=submodules
+        )
         attn_mask_dimensions = self.bert_model._sanity_check_attention_and_get_attn_mask_dimension()
-        attn_mask_type = self.bert_model.transformer_layer_spec.submodules.self_attention.params[
-            'attn_mask_type'
-        ]
+        attn_mask_type = submodules.self_attention.params['attn_mask_type']
         assert (
             attn_mask_type == AttnMaskType.padding
         ), f"Exepcted attn mask type to be padding, but got {attn_mask_type}"
@@ -160,7 +160,7 @@ def test_transformer_engine_version_1_10(self, mocker):
     def test_transformer_engine_version_1_7_to_1_10_flash_attn(self, mocker):
         self.bert_model.config.attention_backend = AttnBackend.flash
         mocker.patch("megatron.core.utils.get_te_version", return_value=PkgVersion("1.8"))
-        self.bert_model.transformer_layer_spec = bert_layer_with_transformer_engine_spec
+        self.bert_model.transformer_layer_spec = get_bert_layer_with_transformer_engine_spec()
         attn_mask_dimensions = self.bert_model._sanity_check_attention_and_get_attn_mask_dimension()
         assert (
             attn_mask_dimensions == "b11s"
@@ -170,15 +170,14 @@ def test_transformer_engine_version_1_7_to_1_10_flash_attn(self, mocker):
     @pytest.mark.flaky
     @pytest.mark.flaky_in_dev
     def test_transformer_engine_version_1_7_to_1_10_rng_error(self, mocker):
-        bert_layer_with_transformer_engine_spec.submodules.self_attention.params[
-            'attn_mask_type'
-        ] == AttnMaskType.padding
+        submodules = get_bert_layer_with_transformer_engine_submodules()
+        submodules.self_attention.params['attn_mask_type'] = AttnMaskType.padding
         mocker.patch("megatron.core.utils.get_te_version", return_value=PkgVersion("1.8"))
         with pytest.raises(Exception) as exc_info:
             self.bert_model = BertModel(
                 config=self.transformer_config,
                 num_tokentypes=0,
-                transformer_layer_spec=bert_layer_with_transformer_engine_spec,
+                transformer_layer_spec=ModuleSpec(module=TransformerLayer, submodules=submodules),
                 vocab_size=100,
                 max_sequence_length=4,
             )
@@ -191,15 +190,14 @@ def test_transformer_engine_version_1_7_to_1_10_rng_error(self, mocker):
     @pytest.mark.internal
     def test_transformer_engine_version_1_7_to_1_10_unfused_attention(self, mocker):
         self.bert_model.config.attention_backend = AttnBackend.unfused
-        bert_layer_with_transformer_engine_spec.submodules.self_attention.params[
-            'attn_mask_type'
-        ] == AttnMaskType.padding
+        submodules = get_bert_layer_with_transformer_engine_submodules()
+        submodules.self_attention.params['attn_mask_type'] = AttnMaskType.padding
         mocker.patch("megatron.core.utils.get_te_version", return_value=PkgVersion("1.8"))
-        self.bert_model.transformer_layer_spec = bert_layer_with_transformer_engine_spec
+        self.bert_model.transformer_layer_spec = ModuleSpec(
+            module=TransformerLayer, submodules=submodules
+        )
         attn_mask_dimensions = self.bert_model._sanity_check_attention_and_get_attn_mask_dimension()
-        attn_mask_type = self.bert_model.transformer_layer_spec.submodules.self_attention.params[
-            'attn_mask_type'
-        ]
+        attn_mask_type = submodules.self_attention.params['attn_mask_type']
         assert (
             attn_mask_type == AttnMaskType.arbitrary
         ), f"Exepcted attn mask type to be arbitrary, but got {attn_mask_type}"
@@ -218,7 +216,7 @@ def test_transformer_engine_version_less_than_1_7(self, mocker):
             self.bert_model = BertModel(
                 config=self.transformer_config,
                 num_tokentypes=0,
-                transformer_layer_spec=bert_layer_with_transformer_engine_spec,
+                transformer_layer_spec=get_bert_layer_with_transformer_engine_spec(),
                 vocab_size=100,
                 max_sequence_length=4,
             )