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/spec_utils.py b/megatron/core/transformer/spec_utils.py index b3de8541734..73f46a8df54 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 @@ -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]): 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, 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