Prep work for migrating to types from ModuleSpec

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/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