Demonstrate how to do typed ModuleSpec but for SelfAttention

megatron/core/models/gpt/gpt_layer_specs.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
+import functools
 import warnings
 from typing import Optional, Union
 
@@ -138,9 +138,10 @@ def get_gpt_layer_with_inference_spec(
         return ModuleSpec(
             module=TransformerLayer,
             submodules=TransformerLayerSubmodules(
-                self_attention=ModuleSpec(
-                    module=SelfAttention,
-                    params={"attn_mask_type": AttnMaskType.causal},
+                # To fix type hinting, simply swap ModuleSpec for functools.partial.
+                self_attention=functools.partial(
+                    SelfAttention,
+                    attn_mask_type=AttnMaskType.causal,
                     submodules=SelfAttentionSubmodules(
                         linear_qkv=backend.column_parallel_layer_norm_linear(),
                         core_attention=backend.core_attention(),
@@ -240,6 +241,9 @@ def get_gpt_layer_with_transformer_engine_spec(
             module=TransformerLayer,
             submodules=TransformerLayerSubmodules(
                 input_layernorm=backend.layer_norm(),
+                # ModuleSpec is still legal to provide and does not trigger a type error - it is
+                # callable with any arguments, so inherently matches any Builder Protocol. (It
+                # doesn't get type-checked, though, so doesn't *benefit* from the Protocol.)
                 self_attention=ModuleSpec(
                     module=MLASelfAttention,
                     params={"attn_mask_type": AttnMaskType.causal},

megatron/core/transformer/attention.py

@@ -293,7 +293,9 @@ def _get_pp_layer_offset_for_inference(self):
         ), "Virtual pipeline parallelism is not supported for inference"
 
         # Import here to avoid circular imports
-        from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+        from megatron.core.transformer.transformer_layer import (
+            get_transformer_layer_offset,
+        )
 
         return get_transformer_layer_offset(
             self.config, vp_stage=None, pp_rank=get_pg_rank(self.pg_collection.pp)
@@ -677,6 +679,7 @@ def forward(
         rotary_pos_cos_sin: Optional[Tensor] = None,
         attention_bias: Optional[Tensor] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
+        # TODO(nschank): This parameter is inconsistently either an int or a Tensor in various places
         sequence_len_offset: Optional[int] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
@@ -984,9 +987,9 @@ def __init__(
         config: TransformerConfig,
         submodules: SelfAttentionSubmodules,
         layer_number: int,
-        attn_mask_type=AttnMaskType.padding,
-        cp_comm_type: str = None,
-        pg_collection: ProcessGroupCollection = None,
+        attn_mask_type: AttnMaskType = AttnMaskType.padding,
+        cp_comm_type: Optional[str] = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         super().__init__(
             config=config,

megatron/core/transformer/spec_utils.py

@@ -2,7 +2,7 @@
 
 import types
 from dataclasses import dataclass, field
-from typing import Tuple, Union
+from typing import Any, Tuple, Union
 
 
 @dataclass
@@ -24,7 +24,16 @@ class ModuleSpec:
 
     module: Union[Tuple, type]
     params: dict = field(default_factory=lambda: {})
-    submodules: type = None
+    submodules: object = None
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Builds an instance of the module from the spec.
+
+        Args:
+            *args: Positional arguments to be passed to the module init.
+            **kwargs: Keyword arguments to be passed to the module init.
+        """
+        return build_module(self, *args, **kwargs)
 
 
 def import_module(module_path: Tuple[str]):

megatron/core/transformer/transformer_layer.py

@@ -4,13 +4,13 @@
 import warnings
 from abc import ABC
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Protocol, Tuple, Union, cast
 
 import torch
 import torch.distributed
 from torch import Tensor
 
-from megatron.core import parallel_state, tensor_parallel
+from megatron.core import parallel_state, tensor_parallel, typed_torch
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
 from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
 from megatron.core.packed_seq_params import PackedSeqParams
@@ -191,6 +191,37 @@ def get_transformer_layer_offset(
     return offset
 
 
+class SelfAttention(Protocol):
+    def forward(
+        self,
+        input_layernorm_output: Tensor,
+        /,
+        *,
+        attention_mask: Tensor,
+        inference_context: Optional[Any],
+        rotary_pos_emb: Optional[Tensor],
+        rotary_pos_cos: Optional[Tensor],
+        rotary_pos_sin: Optional[Tensor],
+        rotary_pos_cos_sin: Optional[Tensor],
+        attention_bias: Optional[Tensor],
+        packed_seq_params: Optional[PackedSeqParams],
+        sequence_len_offset: Optional[int],
+    ) -> Tuple[Tensor, Tensor]: ...
+
+
+class SelfAttentionBuilder(Protocol):
+    """Protocol for building SelfAttention modules."""
+
+    def __call__(
+        self,
+        *,
+        config: TransformerConfig,
+        layer_number: int,
+        pg_collection: ProcessGroupCollection,
+        cp_comm_type: Optional[str] = None,
+    ) -> SelfAttention: ...
+
+
 @dataclass
 class TransformerLayerSubmodules:
     """
@@ -220,7 +251,7 @@ class TransformerLayerSubmodules:
     """
 
     input_layernorm: Union[ModuleSpec, type] = IdentityOp
-    self_attention: Union[ModuleSpec, type] = IdentityOp
+    self_attention: SelfAttentionBuilder = IdentityOp
     self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
 
     pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
@@ -298,10 +329,10 @@ def __init__(
         attention_optional_kwargs["pg_collection"] = pg_collection
 
         # [Module 2: SelfAttention]
-        self.self_attention = build_module(
-            submodules.self_attention,
+        self.self_attention = submodules.self_attention(
             config=self.config,
             layer_number=self.layer_number,
+            # TODO(nschank): Remove this when cross_attention also uses a builder protocol.
             **attention_optional_kwargs,
         )
 
@@ -382,6 +413,9 @@ def __init__(
                 ):
                     self.recompute_input_layernorm = True
                     if self.config.fp8 or self.config.fp4:
+                        # TODO(nschank): For optional things like this, you could maybe use a
+                        # Protocol to check for the method instead of isinstance checks to make
+                        # checking friendlier.
                         self.self_attention.set_for_recompute_input_layernorm()
                 if not isinstance(self.pre_mlp_layernorm, IdentityOp):
                     self.recompute_pre_mlp_layernorm = True
@@ -496,8 +530,9 @@ def _forward_attention(
 
         # Self attention.
         nvtx_range_push(suffix="self_attention")
-        attention_output_with_bias = self.self_attention(
-            input_layernorm_output,
+        # TODO(nschank): attention_mask and sequence_len_offset have type errors.
+        attention_output_with_bias = typed_torch.apply_module(self.self_attention)(
+            cast(Tensor, input_layernorm_output),
             attention_mask=attention_mask,
             inference_context=inference_context,
             rotary_pos_emb=rotary_pos_emb,

megatron/core/typed_torch.py

@@ -0,0 +1,31 @@
+"""Utilities for improved type hinting with torch interfaces."""
+
+from collections.abc import Callable
+from typing import Generic, ParamSpec, Protocol, TypeVar
+
+import torch
+
+P = ParamSpec('P')
+R_co = TypeVar('R_co', covariant=True)
+
+
+class _Module(Generic[P, R_co], Protocol):
+    """Protocol allowing us to unwrap `forward`."""
+
+    def forward(self, *args: P.args, **kwargs: P.kwargs) -> R_co: ...
+
+
+def apply_module(m: _Module[P, R_co], *, check_subclass: bool = True) -> Callable[P, R_co]:
+    """Returns the provided module unchanged, but with correct type hints.
+
+    Args:
+      m: An instance of a subclass of `torch.nn.Module`.
+      check_subclass: If `True`, checks that `m` is a subclass of
+            `torch.nn.Module` and raises a `TypeError` if not.
+
+    Returns:
+      That module unchanged, but with correct type hints.
+    """
+    if check_subclass and not issubclass(type(m), torch.nn.Module):
+        raise TypeError(f'{type(m)} is not a subclass of torch.nn.Module')
+    return m  # type: ignore