[RFC] Refactor attention and make attention mask an argument to the model

**Status**
The PR is not landable yet but server as a RFC. If people are okay with
this design, this PR requires following changes and verifications:

1. Change all models, including the experimental ones.
2. E2E loss verification (this has been done for functional check, but loss verification is noot done yet).
3. We should add an unittest for attention. But since we don't have GPU unittest, this can be done in a seperate PR.

**Summary**
This PR aims to refactor how TorchTitan build the attention masks and pass to model. Before this PR, init_attention_masks() is called in Trainer but the masks are stored as a class variable of FlexAttentionWrapper(). We chose this shortcut to support the case where a single model requires multiple masks.

The previous design has several issues, one particular one is #1723.

Now that pytorch/pytorch#164111 proves that we can let PP split BlockMask, this PR performs the refactor to pass masks as an argument of model.forward().

The new design:
1. Model needs to provide `get_attention_masks()` that accepts `create_mask_fn`, `batch`, and `eos_id`. If the attention op is SDPA, then this API should return None as SDPA currently doesn't support varlen. But once it does, we may have to return some tuple of int that represents the mask.

Justification: attention logic is technically a part of the model, but requires some information from trainer/dataloader. So it's model author's responsibility to provide some API that let trainer calls to get the masks.

2. `get_attention_masks()` will be called from the trainer and the resulting masks are passed to the model.forward().

Justification: this will allow us to
fix #1723 with pytorch/pytorch#164111 and this PR.

3. Provide a single AttentionOp instead of two.

Justification: since the masking logic is moved outside, we don't need
to do bookkeeping of masks in FlexAttentionWrapper. The logic is so
simple that one AttentionOp makes things cleaner.

Note: we still have two very very thin op wrappers that are used for CP. I keep these two for the CP education purpose. But this certinaly can be confusion for Titan's users. I'm opn to merge them to AttentionOp.

See the discussion in #1723.
ghstack-source-id: 35aa425
Pull-Request-resolved: #1776

Author: fegin

torchtitan/experiments/llama4/model/model.py

@@ -4,14 +4,21 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from typing import Callable
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-from torchtitan.models.attention import build_attention
+from torchtitan.models.attention import (
+    AttentionOp,
+    get_causal_mask_mod,
+    get_document_mask_mod,
+    get_fixed_block_mask_mod,
+)
 from torchtitan.models.moe import MoE
-from torchtitan.protocols import ModelProtocol
+from torchtitan.protocols.model import AttentionMasksType
+from torchtitan.protocols.train_spec import ModelProtocol
 
 from .args import TransformerModelArgs
 
@@ -155,9 +162,7 @@ def __init__(
         # values of these two variables.
         self.use_rope = use_rope
 
-        self.sdpa = build_attention(
-            model_args.use_flex_attn, model_args.attn_mask_type, fixed_block_size
-        )
+        self.attention_op = AttentionOp(model_args.use_flex_attn)
 
     def init_weights(self, init_std: float):
         for linear in (self.wq, self.wk, self.wv):
@@ -168,6 +173,7 @@ def forward(
         self,
         x: torch.Tensor,
         freqs_cis: torch.Tensor,
+        attention_masks: AttentionMasksType,
     ):
         """
         Forward pass of the attention module.
@@ -202,7 +208,12 @@ def forward(
         xk = keys.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
         xv = values.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
 
-        output = self.sdpa(xq, xk, xv)
+        assert isinstance(attention_masks, dict) or attention_masks is None
+        if self.attention_op.use_flex_attn:
+            attention_mask = attention_masks["rope" if self.use_rope else "nope"]
+            output = self.attention_op(xq, xk, xv, attention_mask=attention_mask)
+        else:
+            output = self.attention_op(xq, xk, xv, attention_mask=None)
 
         output = output.transpose(
             1, 2
@@ -335,6 +346,7 @@ def forward(
         self,
         x: torch.Tensor,
         freqs_cis: torch.Tensor,
+        attention_masks: AttentionMasksType,
     ):
         """
         Perform a forward pass through the TransformerBlock.
@@ -347,7 +359,7 @@ def forward(
             torch.Tensor: Output tensor after applying attention and feedforward layers.
 
         """
-        h = x + self.attention(self.attention_norm(x), freqs_cis)
+        h = x + self.attention(self.attention_norm(x), freqs_cis, attention_masks)
         if self.moe_enabled:
             out = h + self.moe(self.ffn_norm(h))
         else:
@@ -447,9 +459,36 @@ def _precompute_freqs_cis(self) -> torch.Tensor:
             self.model_args.rope_theta,
         )
 
+    def get_attention_masks(
+        self, create_mask_fn: Callable, batch: torch.Tensor, eos_id: int
+    ) -> AttentionMasksType:
+        if not self.model_args.use_flex_attn:
+            return None
+
+        nope_mask_mod = get_causal_mask_mod()
+        match self.model_args.attn_mask_type:
+            case "causal":
+                B = 1
+            case "block_causal":
+                B = batch.shape[0]
+                rope_mask_mod = get_document_mask_mod(nope_mask_mod, batch, eos_id)
+            case _:
+                raise ValueError(f"Unknown attention mask type: {self.attn_mask_type}")
+
+        rope_mask_mod = get_fixed_block_mask_mod(
+            nope_mask_mod, self.model_args.fixed_attn_block_size
+        )
+
+        seqlen = batch.shape[1]
+        return {
+            "rope": create_mask_fn(rope_mask_mod, B, None, seqlen, seqlen),
+            "nope": create_mask_fn(nope_mask_mod, B, None, seqlen, seqlen),
+        }
+
     def forward(
         self,
         tokens: torch.Tensor,
+        attention_masks: AttentionMasksType,
         input_batch: torch.Tensor | None = None,
     ):
         """
@@ -473,7 +512,7 @@ def forward(
         h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens
 
         for layer in self.layers.values():
-            h = layer(h, self.freqs_cis)
+            h = layer(h, self.freqs_cis, attention_masks)
 
         h = self.norm(h) if self.norm else h
         output = self.output(h) if self.output else h