Refactor attention v2 #10623

lucylq · 2025-05-01T20:33:57Z

Stack from ghstack (oldest at bottom):

-> Refactor attention v2 #10623

Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer.

The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well.

This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py.

I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer.

It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221

Differential Revision: D73538697

Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/) [ghstack-poisoned]

Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/) ghstack-source-id: 279884551 Pull Request resolved: #10623

pytorch-bot · 2025-05-01T20:34:01Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10623

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

CI workflows being skipped on PR

❌ 2 New Failures

As of commit a3fc78a with merge base 4129ebe ():

NEW FAILURES - The following jobs have failed:

pull / android / run-emulator (gh)
The process '/usr/bin/sh' failed with exit code 255
pull / unittest / macos / macos-job (gh)
backends/xnnpack/test/ops/test_add.py::TestAdd::test_qs8_add_relu_seq

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2025-05-01T20:34:09Z

This pull request was exported from Phabricator. Differential Revision: D73538697

iseeyuan

It makes sense to me to pass attention to the transformer, and unblock LoRA.

iseeyuan · 2025-05-02T20:03:43Z

examples/models/llama/llama_transformer.py

@@ -83,25 +84,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:


 class TransformerBlock(nn.Module):
-    def __init__(self, layer_id: int, args: ModelArgs, rope: Rope):
+    def __init__(self, args: ModelArgs, attention: Attention):


Can you add doc string on each argument, especially the attention? I think it makes sense to me that Attention type is required, so that the API of user-defined attention is compatible with our transformer.

Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/) [ghstack-poisoned]

Pull Request resolved: #10623 Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Previously here: D73474110 ghstack-source-id: 281805091 @exported-using-ghexport Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/)

facebook-github-bot · 2025-05-03T01:12:59Z

This pull request was exported from Phabricator. Differential Revision: D73538697

Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/) [ghstack-poisoned]

Pull Request resolved: #10623 Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Previously here: D73474110 ghstack-source-id: 282064227 @exported-using-ghexport Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/)

facebook-github-bot · 2025-05-05T17:34:36Z

This pull request was exported from Phabricator. Differential Revision: D73538697

jackzhxng · 2025-05-05T18:00:31Z

examples/models/llama/llama_transformer.py

@@ -117,7 +138,15 @@ def forward(self, x, freqs_cos, freqs_sin, attn_options: ForwardOptions):  # x:


 class Transformer(nn.Module):
-    def __init__(self, params: ModelArgs):
+    def __init__(self, params: ModelArgs, layers: nn.ModuleList, rope: Rope):


I think if you are going to do this, might as well lift all of the major model components out as well, such as the embedding layer and rms norm, even though they are not customizable by model args at the moment

I think we can, but would prefer to have it in a separate PR if it's something we want to do. Is there a use-case, or more to make Transformer more modular?

Up to you, no use case atm, just for modularity. Just feels a bit weird to me seeing layers and rope be the only lifted inputs for Transformer

jackzhxng · 2025-05-05T18:01:40Z

examples/models/llama/llama_transformer.py

@@ -212,3 +239,23 @@ def forward(
            return logits, attn_options_update

        return logits
+
+
+def construct_transformer(model_args: ModelArgs) -> Transformer:


Why not @classmethod?

discussed offline; construct_transformer is likely going to be more high-level; not quite at model-creation, but will contain eg. lora instantiation so may not make sense for it to be part of the transformer class itself.

Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/) [ghstack-poisoned]

Pull Request resolved: #10623 Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Previously here: D73474110 ghstack-source-id: 282118266 @exported-using-ghexport Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/)

facebook-github-bot · 2025-05-05T21:02:22Z

This pull request was exported from Phabricator. Differential Revision: D73538697

Pull Request resolved: #10623 Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Previously here: D73474110 ghstack-source-id: 282118266 @exported-using-ghexport Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/)

Pull Request resolved: pytorch#10623 Pull attention creation out of Transformer/TransformerBlock. Instead, pass the layers into Transformer. The motivation is to customize linear layers in attention for LoRA (eg. make wq into a LoraLinear instead of a regular linear). In the next diff (D73517350), we pull wq,wk,wv,wo out of the attention and pass those in as well. This allows us to customize attention parameters without passing in ModelArgs and doing the customization deep inside attention.py. I think this modularizes our attention/transformer components, though also means that users have to do some more work to construct the attention layers and pass it to transformer. It follows the torchtune structure more closely, eg. https://github.com/pytorch/torchtune/blob/main/torchtune/models/llama3_2/_component_builders.py#L221 Previously here: D73474110 ghstack-source-id: 282118266 @exported-using-ghexport Differential Revision: [D73538697](https://our.internmc.facebook.com/intern/diff/D73538697/)

lucylq requested a review from jackzhxng as a code owner May 1, 2025 20:33

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 1, 2025

facebook-github-bot added the fb-exported label May 1, 2025

iseeyuan approved these changes May 2, 2025

View reviewed changes

iseeyuan reviewed May 2, 2025

View reviewed changes

lucylq added the topic: not user facing label May 3, 2025

jackzhxng reviewed May 5, 2025

View reviewed changes

facebook-github-bot merged commit a4d5fb9 into gh/lucylq/74/base May 6, 2025
161 of 167 checks passed

facebook-github-bot deleted the gh/lucylq/74/head branch May 6, 2025 03:59

facebook-github-bot temporarily deployed to cherry-pick-bot May 6, 2025 03:59 — with GitHub Actions Inactive

pytorchbot mentioned this pull request May 6, 2025

Refactor attention v2 #10707

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Refactor attention v2 #10623

Refactor attention v2 #10623

Uh oh!

lucylq commented May 1, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented May 1, 2025 •

edited

Loading

Uh oh!

facebook-github-bot commented May 1, 2025

Uh oh!

iseeyuan left a comment

Uh oh!

iseeyuan May 2, 2025

Uh oh!

facebook-github-bot commented May 3, 2025

Uh oh!

facebook-github-bot commented May 5, 2025

Uh oh!

jackzhxng May 5, 2025

Uh oh!

lucylq May 5, 2025 •

edited

Loading

Uh oh!

jackzhxng May 5, 2025

Uh oh!

jackzhxng May 5, 2025

Uh oh!

lucylq May 5, 2025

Uh oh!

facebook-github-bot commented May 5, 2025

Uh oh!

Uh oh!

Uh oh!

Refactor attention v2 #10623

Refactor attention v2 #10623

Uh oh!

Conversation

lucylq commented May 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10623

❗ 1 Active SEVs

❌ 2 New Failures

Uh oh!

facebook-github-bot commented May 1, 2025

Uh oh!

iseeyuan left a comment

Choose a reason for hiding this comment

Uh oh!

iseeyuan May 2, 2025

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented May 3, 2025

Uh oh!

facebook-github-bot commented May 5, 2025

Uh oh!

jackzhxng May 5, 2025

Choose a reason for hiding this comment

Uh oh!

lucylq May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jackzhxng May 5, 2025

Choose a reason for hiding this comment

Uh oh!

jackzhxng May 5, 2025

Choose a reason for hiding this comment

Uh oh!

lucylq May 5, 2025

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented May 5, 2025

Uh oh!

Uh oh!

Uh oh!

lucylq commented May 1, 2025 •

edited

Loading

pytorch-bot bot commented May 1, 2025 •

edited

Loading

lucylq May 5, 2025 •

edited

Loading