Repository: pytorch/torchtitan
Pull request #887: Enable FlexAttention for llama model
Author: login: fegin
Labels: name: CLA Signed
description: This label is managed by the Meta Open Source bot.
name: fb-exported
Created: 2025-02-25T18:00:52Z | Merged: 2025-03-20T22:09:12Z

Summary: This PR enables using FlexAttention for llama model.

Differential Revision: D70186496

SDPA
```
[rank2]:[titan] 2025-03-20 00:15:44,653 - root - INFO - step: 20  loss:  8.4828  memory: 36.06GiB(45.56%)  tps: 2,546  tflops: 147.44  mfu: 47.26%
```
FlexAttention
```
[rank0]:[titan] 2025-03-20 00:38:14,024 - root - INFO - step: 20  loss:  8.4238  memory: 36.06GiB(45.56%)  tps: 2,168  tflops: 125.53  mfu: 40.23%
```
FlexAttention + max auto tune without cuda graph
```
[rank0]:[titan] 2025-03-20 00:33:37,008 - root - INFO - step: 20  loss:  8.5293  memory: 36.75GiB(46.43%)  tps: 2,421  tflops: 140.20  mfu: 44.94%
```


Diff:
diff --git a/tests/integration_tests.py b/tests/integration_tests.py
index 8453579df1..59a73efbf9 100755
--- a/tests/integration_tests.py
+++ b/tests/integration_tests.py
@@ -296,6 +296,18 @@ def build_test_list():
             "hsdp",
             ngpu=4,
         ),
+        OverrideDefinitions(
+            [
+                [
+                    "--parallelism.data_parallel_shard_degree=4",
+                    "--activation_checkpoint.mode='full'",
+                    "--model.use_flex_attn",
+                ]
+            ],
+            "FSDP+FLEX_ATTN",
+            "fsdp+flex_attn",
+            ngpu=4,
+        ),
         OverrideDefinitions(
             [
                 [
diff --git a/torchtitan/config_manager.py b/torchtitan/config_manager.py
index 1515985ff9..b2134ecc6b 100644
--- a/torchtitan/config_manager.py
+++ b/torchtitan/config_manager.py
@@ -193,6 +193,11 @@ def __init__(self):
             choices=["layernorm", "np_layernorm", "rmsnorm"],
             help="Type of layer normalization to use [layernorm, np_layernorm, rmsnorm]",
         )
+        self.parser.add_argument(
+            "--model.use_flex_attn",
+            action="store_true",
+            help="Whether to use Flex Attention.",
+        )
         self.parser.add_argument(
             "--model.tokenizer_path",
             type=str,
diff --git a/torchtitan/models/llama/model.py b/torchtitan/models/llama/model.py
index a80e5663c6..591d1f33ea 100644
--- a/torchtitan/models/llama/model.py
+++ b/torchtitan/models/llama/model.py
@@ -8,11 +8,16 @@
 
 
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Callable, ClassVar, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.nn.attention.flex_attention import (
+    BlockMask,
+    create_block_mask,
+    flex_attention,
+)
 
 from torchtitan.models.norms import build_norm
 from torchtitan.protocols.train_spec import BaseModelArgs, ModelProtocol
@@ -36,6 +41,8 @@ class TransformerModelArgs(BaseModelArgs):
     depth_init: bool = True
     norm_type: str = "rmsnorm"
 
+    use_flex_attn: bool = False
+
 
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
     """
@@ -146,6 +153,13 @@ class Attention(nn.Module):
 
     """
 
+    # We registered flex_attention related attributes as class variables as we
+    # need to amortize the cost of compilation. Enabling per-instance flex_attention
+    # is not supported.
+    block_mask: ClassVar[Optional[BlockMask]] = None
+    use_flex_attn: ClassVar[bool] = False
+    flex_attn: ClassVar[Optional[Callable]] = None
+
     def __init__(self, model_args: TransformerModelArgs):
         super().__init__()
         self.n_heads = model_args.n_heads
@@ -165,6 +179,7 @@ def __init__(self, model_args: TransformerModelArgs):
         self.wo = nn.Linear(
             model_args.n_heads * self.head_dim, model_args.dim, bias=False
         )
+        self.use_flex_attn = model_args.use_flex_attn
 
     def init_weights(self, init_std: float):
         for linear in (self.wq, self.wk, self.wv):
@@ -187,6 +202,7 @@ def forward(
             torch.Tensor: Output tensor after attention.
 
         """
+
         bs, seqlen, _ = x.shape
         xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
 
@@ -208,13 +224,34 @@ def forward(
         xv = values.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
 
         # we use casual mask for training
-        output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
+        if self.use_flex_attn:
+            # assert False, (type(xq), type(xk), type(xv))
+            self._init_flex_attn(seqlen=seqlen)
+            output = self.flex_attn(xq, xk, xv, block_mask=self.block_mask)
+        else:
+            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
         output = output.transpose(
             1, 2
         ).contiguous()  # (bs, seqlen, n_local_heads, head_dim)
         output = output.view(bs, seqlen, -1)
         return self.wo(output)
 
+    @torch.no_grad()
+    def _init_flex_attn(self, seqlen: int) -> None:
+        if self.block_mask is not None:
+            return
+
+        def causal_mask(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx
+
+        compiled_create_block_mask = torch.compile(create_block_mask)
+        self.block_mask = compiled_create_block_mask(
+            causal_mask, None, None, seqlen, seqlen
+        )
+        self.flex_attn = torch.compile(
+            flex_attention, mode="max-autotune-no-cudagraphs"
+        )
+
 
 class FeedForward(nn.Module):
     """
diff --git a/torchtitan/models/llama/parallelize_llama.py b/torchtitan/models/llama/parallelize_llama.py
index f247138216..182a76d7c8 100644
--- a/torchtitan/models/llama/parallelize_llama.py
+++ b/torchtitan/models/llama/parallelize_llama.py
@@ -73,6 +73,14 @@ def parallelize_llama(
         )
 
     if job_config.activation_checkpoint.mode != "none":
+        if (
+            job_config.activation_checkpoint.mode == "selective"
+            and job_config.model.use_flex_attn
+        ):
+            raise ValueError(
+                "FlexAttention is not compatible with selective AC yet. "
+                "See https://github.com/pytorch/pytorch/issues/147879"
+            )
         apply_ac(model, job_config.activation_checkpoint)
 
     # turn on per-TransformerBlock compile after AC wrapping and before FSDP
diff --git a/torchtitan/train.py b/torchtitan/train.py
index b1da364d3a..682a3ba4f0 100644
--- a/torchtitan/train.py
+++ b/torchtitan/train.py
@@ -147,6 +147,7 @@ def __init__(self, job_config: JobConfig):
         model_config.norm_type = job_config.model.norm_type
         model_config.vocab_size = tokenizer.n_words
         model_config.max_seq_len = job_config.training.seq_len
+        model_config.use_flex_attn = job_config.model.use_flex_attn
 
         logger.info(
             f"Building {self.train_spec.name} {job_config.model.flavor} with {model_config}"


Reviews:
user: login: drisspg
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T23:30:20Z
user: login: drisspg
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T23:33:59Z
user: login: drisspg
body: does torchtitan have tests? lol
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-14T23:34:22Z
user: login: tianyu-l
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T20:18:39Z
user: login: fegin
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T21:58:56Z
user: login: wanchaol
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T04:49:55Z

Review comments:
diff_hunk: @@ -208,13 +224,32 @@ def forward(
         xv = values.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
 
         # we use casual mask for training
-        output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
+        if self.use_flex_attn:
+            # assert False, (type(xq), type(xk), type(xv))
+            self._init_flex_attn(seqlen=seqlen)
+            output = self.flex_attn(xq, xk, xv, block_mask=self.block_mask)
+        else:
+            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
         output = output.transpose(
             1, 2
         ).contiguous()  # (bs, seqlen, n_local_heads, head_dim)
         output = output.view(bs, seqlen, -1)
         return self.wo(output)
 
+    @torch.no_grad()
+    def _init_flex_attn(self, seqlen: int) -> None:
+        if self.block_mask is not None:
+            return
+
+        def causal_mask(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx
+
+        compiled_create_block_mask = torch.compile(create_block_mask)
path: torchtitan/models/llama/model.py
user: login: drisspg
body: compiling w/ max-autotune can have a non zero impact to e2e perf just fyi
created_at: 2025-03-14T23:30:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -146,6 +153,13 @@ class Attention(nn.Module):
 
     """
 
+    # We registered flex_attention related attributes as class variables as we
+    # need to amortize the cost of compilation. Enabling per-instance flex_attention
+    # is not supported.
path: torchtitan/models/llama/model.py
user: login: drisspg
body: Looks good, this important
created_at: 2025-03-14T23:33:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -208,13 +224,34 @@ def forward(
         xv = values.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
 
         # we use casual mask for training
-        output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
+        if self.use_flex_attn:
+            # assert False, (type(xq), type(xk), type(xv))
+            self._init_flex_attn(seqlen=seqlen)
+            output = self.flex_attn(xq, xk, xv, block_mask=self.block_mask)
+        else:
+            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
         output = output.transpose(
             1, 2
         ).contiguous()  # (bs, seqlen, n_local_heads, head_dim)
         output = output.view(bs, seqlen, -1)
         return self.wo(output)
 
+    @torch.no_grad()
+    def _init_flex_attn(self, seqlen: int) -> None:
path: torchtitan/models/llama/model.py
user: login: tianyu-l
body: Given how independent this piece of code could be, do you think it makes sense to put it in a separate file, e.g. `flex_attention.py`? It is like `norms.py` (which other models might use too).
created_at: 2025-03-20T20:18:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -73,6 +73,14 @@ def parallelize_llama(
         )
 
     if job_config.activation_checkpoint.mode != "none":
+        if (
+            job_config.activation_checkpoint.mode == "selective"
+            and job_config.model.use_flex_attn
+        ):
+            raise ValueError(
path: torchtitan/models/llama/parallelize_llama.py
user: login: tianyu-l
body: can we also error out when FlexAttention is used with CP, for now?
created_at: 2025-03-20T20:18:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -73,6 +73,14 @@ def parallelize_llama(
         )
 
     if job_config.activation_checkpoint.mode != "none":
+        if (
+            job_config.activation_checkpoint.mode == "selective"
+            and job_config.model.use_flex_attn
+        ):
+            raise ValueError(
path: torchtitan/models/llama/parallelize_llama.py
user: login: fegin
body: Will have a separate PR to add the warning. Already initiated the diff landing from the internal code base.
created_at: 2025-03-20T21:58:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -208,13 +224,34 @@ def forward(
         xv = values.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
 
         # we use casual mask for training
-        output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
+        if self.use_flex_attn:
+            # assert False, (type(xq), type(xk), type(xv))
+            self._init_flex_attn(seqlen=seqlen)
+            output = self.flex_attn(xq, xk, xv, block_mask=self.block_mask)
+        else:
+            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
         output = output.transpose(
             1, 2
         ).contiguous()  # (bs, seqlen, n_local_heads, head_dim)
         output = output.view(bs, seqlen, -1)
         return self.wo(output)
 
+    @torch.no_grad()
+    def _init_flex_attn(self, seqlen: int) -> None:
path: torchtitan/models/llama/model.py
user: login: wanchaol
body: IMO maybe making a SDPA module that switches between fsdp and flex attention might be better instead of embedding them to each model files
created_at: 2025-03-21T04:49:55Z
author_association: CONTRIBUTOR

PR comments:
user: login: facebook-github-bot
created_at: 2025-02-25T18:01:19Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-13T21:23:18Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-14T06:19:18Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-14T06:28:37Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-14T06:40:50Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-14T22:49:19Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: fegin
created_at: 2025-03-17T06:44:40Z
body: Yes, TorchTitan has integration tests. I'll enable FlexAttention with E2E test before landing. :)
author_association: CONTRIBUTOR
user: login: facebook-github-bot
created_at: 2025-03-20T07:33:04Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-20T07:40:47Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-20T15:47:11Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-20T15:49:59Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-03-20T19:59:03Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D70186496](https://www.internalfb.com/diff/D70186496)
author_association: NONE