From abdf998c8e9e1b4291af030550d1a9dc8fa787a2 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Mon, 1 Dec 2025 20:33:23 +0000 Subject: [PATCH 1/4] Make ModuleSpec callable --- megatron/core/transformer/spec_utils.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/megatron/core/transformer/spec_utils.py b/megatron/core/transformer/spec_utils.py index b3de8541734..ae5f24d4ace 100644 --- a/megatron/core/transformer/spec_utils.py +++ b/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 @@ -26,6 +26,15 @@ class ModuleSpec: params: dict = field(default_factory=lambda: {}) submodules: type = 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]): """Import a named object from a module in the context of this function. From 0005e96ce5c20d0f8875b0a908cc1e84f04e6dc0 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Mon, 1 Dec 2025 20:42:00 +0000 Subject: [PATCH 2/4] Create a typed_torch type-hinting library --- megatron/core/typed_torch.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 megatron/core/typed_torch.py diff --git a/megatron/core/typed_torch.py b/megatron/core/typed_torch.py new file mode 100644 index 00000000000..8ca163315a8 --- /dev/null +++ b/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 From ba632c2deabd4bd06b7d3b140925121e5eb9359d Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Mon, 1 Dec 2025 21:04:30 +0000 Subject: [PATCH 3/4] Fix the type of submodules in ModuleSpec --- megatron/core/transformer/spec_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/transformer/spec_utils.py b/megatron/core/transformer/spec_utils.py index ae5f24d4ace..73f46a8df54 100644 --- a/megatron/core/transformer/spec_utils.py +++ b/megatron/core/transformer/spec_utils.py @@ -24,7 +24,7 @@ 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. From 908728f3449f9a23da1845e966d3712c5f757ca9 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Wed, 10 Dec 2025 04:07:27 +0000 Subject: [PATCH 4/4] Demonstrate ModuleSpec replacement for SelfAttention --- megatron/core/models/gpt/gpt_layer_specs.py | 12 +++-- megatron/core/transformer/attention.py | 11 +++-- .../core/transformer/transformer_layer.py | 49 ++++++++++++++++--- 3 files changed, 57 insertions(+), 15 deletions(-) diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index cffeb9234f1..968c41b4cc8 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/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}, diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 1051799db94..fd0f690316a 100644 --- a/megatron/core/transformer/attention.py +++ b/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, diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index edddc82d7ae..d37596da982 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/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,