remove explicit mask def

Author: liangel-02

torchtitan/models/attention.py

@@ -8,7 +8,7 @@
 
 import functools
 from collections.abc import Callable
-from typing import ClassVar
+from typing import ClassVar, NamedTuple
 
 import torch
 import torch.nn.functional as F
@@ -20,13 +20,14 @@
     flex_attention,
 )
 
-from torch.nn.attention.varlen import varlen_attn, VarlenMetadata
+from torch.nn.attention.varlen import varlen_attn
 
 
 __all__ = [
     "FlexAttentionWrapper",
     "ScaledDotProductAttentionWrapper",
     "VarlenAttentionWrapper",
+    "VarlenMetadata",
     "get_causal_mask_mod",
     "get_document_mask_mod",
     "get_sliding_window_mask_mod",
@@ -35,6 +36,18 @@
 ]
 
 
+class VarlenMetadata(NamedTuple):
+    """
+    Cumulative sequence positions for queries and keys/values.
+
+    """
+
+    cu_seq_q: torch.Tensor
+    cu_seq_k: torch.Tensor
+    max_q: int
+    max_k: int
+
+
 class VarlenAttentionWrapper(torch.nn.Module):
     _compiled_varlen_attn: ClassVar[Callable] = torch.compile(
         varlen_attn, mode="max-autotune-no-cudagraphs"

torchtitan/models/llama3/__init__.py

@@ -46,7 +46,6 @@
         vocab_size=2048,
         rope_theta=500000,
         use_varlen_attn=True,
-        attn_mask_type="varlen_attn",
     ),
     "8B": TransformerModelArgs(
         dim=4096,
@@ -77,7 +76,6 @@
         multiple_of=1024,
         rope_theta=500000,
         use_varlen_attn=True,
-        attn_mask_type="varlen_attn",
     ),
     "70B": TransformerModelArgs(
         dim=8192,

torchtitan/models/llama3/model/model.py

@@ -13,8 +13,6 @@
 from torch import nn
 from torch.nn.attention.flex_attention import and_masks, BlockMask
 
-from torch.nn.attention.varlen import varlen_attn, VarlenMetadata
-
 from torchtitan.components.tokenizer import BaseTokenizer
 from torchtitan.models.attention import (
     create_attention_mask,
@@ -23,6 +21,7 @@
     get_causal_mask_mod,
     get_document_mask_mod,
     ScaledDotProductAttentionWrapper,
+    VarlenAttentionWrapper,
 )
 from torchtitan.protocols.model import AttentionMasksType
 from torchtitan.protocols.train_spec import ModelProtocol
@@ -199,7 +198,7 @@ def __init__(self, model_args: TransformerModelArgs):
         if self.use_flex_attn:
             self.inner_attention = FlexAttentionWrapper()
         elif self.use_varlen_attn:
-            self.inner_attention = varlen_attn
+            self.inner_attention = VarlenAttentionWrapper()
         else:
             self.inner_attention = ScaledDotProductAttentionWrapper()
 
@@ -251,8 +250,6 @@ def forward(
             assert isinstance(attention_masks, BlockMask), attention_masks
             output = self.inner_attention(xq, xk, xv, block_mask=attention_masks)
         elif self.use_varlen_attn:
-            assert isinstance(attention_masks, VarlenMetadata), attention_masks
-
             cu_seq_q = attention_masks.cu_seq_q
             cu_seq_k = attention_masks.cu_seq_k
             max_q = attention_masks.max_q
@@ -495,14 +492,14 @@ def get_attention_masks(
         extra_inputs: dict[str, torch.Tensor] | None = None,
     ) -> AttentionMasksType:
         mask_mods = [get_causal_mask_mod()]
+        if self.model_args.use_varlen_attn:
+            return create_varlen_cu_seqs(input_batch, tokenizer.eos_id)
         match self.model_args.attn_mask_type:
             case "causal":
                 B = 1
             case "block_causal":
                 B = input_batch.shape[0]
                 mask_mods.append(get_document_mask_mod(input_batch, tokenizer.eos_id))
-            case "varlen_attn":
-                return create_varlen_cu_seqs(input_batch, tokenizer.eos_id)
             case _:
                 raise ValueError(
                     f"Unknown attention mask type: {self.model_args.attn_mask_type}"

torchtitan/protocols/model.py

@@ -12,11 +12,11 @@
 import torch.nn as nn
 
 from torch.nn.attention.flex_attention import BlockMask
-from torch.nn.attention.varlen import VarlenMetadata
 
 from torchtitan.components.tokenizer import BaseTokenizer
 
 from torchtitan.config import JobConfig
+from torchtitan.models.attention import VarlenMetadata
 
 
 AttentionMasksType = dict[str, BlockMask] | BlockMask | VarlenMetadata