From 43e7e63b4630f4c65ab0953ecd3674d0f7830c71 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Mon, 1 Dec 2025 20:33:23 +0000 Subject: [PATCH 1/6] 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 2d560afd36c337ec24cb8fa53cef353396975150 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Mon, 1 Dec 2025 20:42:00 +0000 Subject: [PATCH 2/6] 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 c2ec6bc099db582e05cdcc2c76304366d0b59eb7 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Mon, 1 Dec 2025 21:04:30 +0000 Subject: [PATCH 3/6] 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 ab36a9804878ba4d9e952916ff238fc4b09075c2 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Thu, 18 Dec 2025 20:57:15 +0000 Subject: [PATCH 4/6] Fix lint errors --- megatron/core/transformer/spec_utils.py | 13 ++++++++++++- megatron/core/typed_torch.py | 4 +++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/spec_utils.py b/megatron/core/transformer/spec_utils.py index 73f46a8df54..160cad10ed7 100644 --- a/megatron/core/transformer/spec_utils.py +++ b/megatron/core/transformer/spec_utils.py @@ -1,9 +1,12 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +import logging import types from dataclasses import dataclass, field from typing import Any, Tuple, Union +logger = logging.getLogger(__name__) + @dataclass class ModuleSpec: @@ -46,12 +49,13 @@ def import_module(module_path: Tuple[str]): try: module = __import__(base_path, globals(), locals(), [name]) except ImportError as e: - print(f"couldn't import module due to {e}") + logger.error(f"couldn't import module due to {e}") return None return vars(module)[name] def get_module(spec_or_module: Union[ModuleSpec, type], **additional_kwargs): + """Returns or imports the provided module.""" # If a module clas is already provided return it as is if isinstance(spec_or_module, (type, types.FunctionType)): return spec_or_module @@ -65,6 +69,13 @@ def get_module(spec_or_module: Union[ModuleSpec, type], **additional_kwargs): def build_module(spec_or_module: Union[ModuleSpec, type], *args, **kwargs): + """Builds an instance of the module from the spec. + + Args: + spec_or_module: The module spec or module class to build. + *args: Positional arguments to be passed to the module init. + **kwargs: Keyword arguments to be passed to the module init. + """ # If the passed `spec_or_module` is # a `Function`, then return it as it is # NOTE: to support an already initialized module add the following condition diff --git a/megatron/core/typed_torch.py b/megatron/core/typed_torch.py index 8ca163315a8..d1f7524cdca 100644 --- a/megatron/core/typed_torch.py +++ b/megatron/core/typed_torch.py @@ -12,7 +12,9 @@ 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 forward(self, *args: P.args, **kwargs: P.kwargs) -> R_co: + """Forward method of the matching torch.nn.Module.""" + ... def apply_module(m: _Module[P, R_co], *, check_subclass: bool = True) -> Callable[P, R_co]: From d12cce03864f1546f080b6149df2a1323685ebe1 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Fri, 19 Dec 2025 19:41:35 +0000 Subject: [PATCH 5/6] Add copyright notice to typed_torch --- megatron/core/typed_torch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/core/typed_torch.py b/megatron/core/typed_torch.py index d1f7524cdca..dd14eabd1c1 100644 --- a/megatron/core/typed_torch.py +++ b/megatron/core/typed_torch.py @@ -1,3 +1,4 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Utilities for improved type hinting with torch interfaces.""" from collections.abc import Callable From dd191b4be53347a8affcf954d84946974bf2c677 Mon Sep 17 00:00:00 2001 From: Nick Schank Date: Wed, 24 Dec 2025 19:10:09 +0000 Subject: [PATCH 6/6] Fix unittest depending on ModuleSpec not being callable --- tests/unit_tests/dist_checkpointing/models/test_bert_model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit_tests/dist_checkpointing/models/test_bert_model.py b/tests/unit_tests/dist_checkpointing/models/test_bert_model.py index 27f01447851..81b01c8f886 100644 --- a/tests/unit_tests/dist_checkpointing/models/test_bert_model.py +++ b/tests/unit_tests/dist_checkpointing/models/test_bert_model.py @@ -24,13 +24,11 @@ def initialize_bert_model( - seed, layer_spec_fn=bert_layer_with_transformer_engine_spec, vocab_size=128, **config_kwargs + seed, layer_spec=bert_layer_with_transformer_engine_spec, vocab_size=128, **config_kwargs ): torch.manual_seed(seed) model_parallel_cuda_manual_seed(seed) - layer_spec = layer_spec_fn() if callable(layer_spec_fn) else layer_spec_fn - default_config_kwargs = dict( num_layers=8, hidden_size=16,