Migrate SelfAttention and CrossAttention off of ModuleSpec by nschank · Pull Request #2454 · NVIDIA/Megatron-LM

nschank · 2025-12-02T15:34:08Z

What does this PR do ?

Demonstrates what a migration away from ModuleSpec would look like. Draft is mostly for demonstration - can break it up or target the first updates differently depending on request.

Associated design doc: https://docs.google.com/document/d/1shyv0iKEzRdevLOlouF_NktbdJazvWifqxUwPXFigQE/edit?tab=t.0#heading=h.uwes2zo47yg6

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2025-12-02T15:34:12Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Skylion007 · 2025-12-11T18:07:38Z

FYI, this typing utility might be useful. I need to clean it up a bit though: https://github.com/pytorch/pytorch/pull/163418/files let's you copy signatures form super classes and validate it.

Skylion007 · 2025-12-11T18:08:36Z

    if padding:
        attn_mask_type = AttnMaskType.padding_causal

+    assert TELayerNormColumnParallelLinear is not None


Creating a not_none typing utility (or just using the one in PyTorch) might make your lives a bit easier.

Yeah that'd be fine! I don't actually want to do it this way (it's clearly sorta ugly), it was just a cheap way to get the type-checker to pass locally while I was drafting. Sorry for the confusion :)

Skylion007 · 2025-12-11T18:09:17Z

-    def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type:
+    def layer_norm(
+        self, rms_norm: bool = False, for_qk: bool = False
+    ) -> Union[type[FusedLayerNorm], type[TENorm]]:


Yikes, have a TypeAlias here that references all known LayerNorm classes here if you aren't going to have a parent class.

Ah so this is something that will take some thinking when we get to it - I'm taking advantage of covariance here and simply having the subclass specify what it would actually return, specifically because it allows for the type-checker to actually examine the underlying types for the correct signatures in practice. This is not a desired end-state, but was needed (at least while I was playing around) to actually ensure type-checking occurred at all the call-sites where implementations and Protocols needed to be matched up.

In reality, we wouldn't want the BackendProvider interface to specify a base class or a union of classes (which would prevent users from providing custom implementations without subclassing), we would want it to actually have methods which return the specific Protocol that needs to be satisfied; but this would require a bit more of a refactor since each individual method right now probably needs to be split up in order to describe the different ways in which the layers are used. For the moment, since there are only a few BackendProviders and they are mostly used non-polymorphically, it was cheaper to simply push the subclasses to declare the exact types they were returning, since that made ~almost all callers then type-check that the return types actually satisfy the Protocols.

In coming PRs, expect this to be done more cleanly

Skylion007 · 2025-12-11T18:10:55Z

-    def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type:
+    def layer_norm(
+        self, rms_norm: bool = False, for_qk: bool = False
+    ) -> Union[type['FusedLayerNorm'], type[WrappedTorchNorm]]:


Import this other class Is it defined? Why is it stringified is it defined here too? Type checker is not gooing to like this and fall back to Any, wont it?

Type checker doesn't care about strings vs. not - I stringified because if the apex import fails, then this type is unbound.

However, this is all more of a draft PR just kind of showing the broader ideas here - I'd like to clean up the imports instead when I actually update this, to avoid needing to be weird like this.

Skylion007 · 2025-12-11T18:11:56Z

+            linear_qkv = TELayerNormColumnParallelLinear
+            core_attention = TEDotProductAttention
+            linear_proj = TERowParallelLinear


Suggested change

linear_qkv = TELayerNormColumnParallelLinear

core_attention = TEDotProductAttention

linear_proj = TERowParallelLinear

linear_qkv = not_none(TELayerNormColumnParallelLinear)

core_attention = not_none(TEDotProductAttention)

linear_proj = not_none(TERowParallelLinear)

Skylion007 · 2025-12-11T18:13:51Z

+        else:
+            linear_qkv = ColumnParallelLinear
+            core_attention = DotProductAttention
+            linear_proj = RowParallelLinear


Use not_none from torch typing utilities or copy it here

Suggested change

linear_proj = RowParallelLinear

linear_qkv = not_none(TELayerNormColumnParallelLinear)

core_attention = not_none(TEDotProductAttention)

linear_proj = not_none(TERowParallelLinear)

Skylion007 · 2025-12-11T18:14:25Z

+        bias: bool,
+        skip_bias_add: bool,


boolean trap, bias coould be an enum that had a skip_bias_add_option

I am explicitly not going to try to fix the existing interfaces in these PRs - these Protocols will be exact replicas of the current way in which the underlying modules will be used, to ensure backwards compatibility. (I'm happy to also help fix things, but I don't want to get caught in the trap of doing both at once and confusing everyone)

Skylion007 · 2025-12-11T18:15:55Z

+            *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)


Just copy the signature from build_module or do:

Suggested change

return build_module(self, *args, **kwargs)

__call__ = build_module

What do you mean 'copy the signature'? Happy to do it either way since they're equivalent, but the existing methods has no type hints or docstring!

Skylion007 · 2025-12-11T18:22:08Z

+    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]:


You can use a typing overload here to have check_subclass=True be a static assertion here. Now if check_subclass is True it will require it to follow the _TrueModule Protocol and complain if it doesn't. Mypy, pyright, etc prefer the more specific overload so this should work.

Suggested change

def apply_module(m: _Module[P, R_co], *, check_subclass: bool = True) -> Callable[P, R_co]:

class _TrueModule(torch.nn.Module, _Module, Protocol):

pass

@typing.overload

def apply_module(m: _TrueModule[P, R_co], *, check_subclass: Literal[True] = True) -> Callable[P, R_co]:

pass

@typing.overload

def apply_module(m: _Module[P, R_co], *, check_subclass: Literal[True] = True) -> typing.Never:

pass

This is a good idea, I'll add this!

Oh wait no I tried this when I originally wrote it - Protocols cannot extend non-Protocol classes. Unfortunately python typing is fairly limited and doesn't allow for higher-kinded or constrained types :(

nschank · 2025-12-15T22:52:51Z

@Skylion007 I just saw this, sorry - I was mostly sending this PR as an example while the doc was being discussed. I broke out the first few commits into a prep PR: #2668 happy to discuss further there!

nschank · 2025-12-16T14:58:46Z

FYI, this typing utility might be useful. I need to clean it up a bit though: https://github.com/pytorch/pytorch/pull/163418/files let's you copy signatures form super classes and validate it.

You're doing god's work there, but unfortunately that wouldn't help me in the short-term - until all existing modules (both inside and outside of Megatron, plus all custom modules provided by users) are consistently annotated correctly, I will still need a way to figure out the correct signature of a module nonintrusively. I view the typed_torch library as a necessary evil which can be discarded easily when the time comes!

nschank added 8 commits December 2, 2025 15:26

Make ModuleSpec callable

abdf998

Create a typed_torch type-hinting library

0005e96

Fix the type of submodules in ModuleSpec

ba632c2

Add Protocol-based typing for SelfAttention.linear_qkv

f4b32d5

Add Protocol-based typing for Self/CrossAttention.core_attention

2c1c709

Add Protocol-based typing for CrossAttention.linear_q/kv

ae27bc5

Add Protocol-based typing for Cross/SelfAttention.linear_proj

4ea6cd5

Add Protocol-based typing for SelfAttention.layernorm

9256eb5

github-actions Bot added the community-request label Dec 2, 2025

ericharper requested a review from asolergi-nv December 2, 2025 18:53

nschank mentioned this pull request Dec 10, 2025

Demonstrate how to do typed ModuleSpec but for SelfAttention nschank/Megatron-LM#1

Closed

Skylion007 reviewed Dec 11, 2025

View reviewed changes

nschank closed this Jan 9, 2026

nschank deleted the typespec1 branch January 9, 2026 17:27

nschank mentioned this pull request Feb 11, 2026

Split layer_specs to return Submodules instead of ModuleSpecs #3255

Merged

6 tasks

-            linear_proj = RowParallelLinear
+        linear_qkv = not_none(TELayerNormColumnParallelLinear)
+        core_attention = not_none(TEDotProductAttention)
+        linear_proj = not_none(TERowParallelLinear)

	return build_module(self, args, *kwargs)
	__call__ = build_module

		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]:

-def apply_module(m: _Module[P, R_co], *, check_subclass: bool = True) -> Callable[P, R_co]:
+class _TrueModule(torch.nn.Module, _Module, Protocol):
+    pass
+@typing.overload
+def apply_module(m: _TrueModule[P, R_co], *, check_subclass: Literal[True] = True) -> Callable[P, R_co]:
+    pass
+@typing.overload
+def apply_module(m: _Module[P, R_co], *, check_subclass: Literal[True] = True) -> typing.Never:
+    pass

Uh oh!

Conversation

nschank commented Dec 2, 2025

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Dec 2, 2025

Uh oh!

Skylion007 commented Dec 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nschank commented Dec 15, 2025

Uh oh!

nschank commented Dec 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

(Step 1): Add PR label `Expert Review`