remove explicit mask def

Author: liangel-02

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,7 +13,7 @@
 from torch import nn
 from torch.nn.attention.flex_attention import and_masks, BlockMask
 
-from torch.nn.attention.varlen import varlen_attn, VarlenMetadata
+from torch.nn.attention.varlen import varlen_attn
 
 from torchtitan.components.tokenizer import BaseTokenizer
 from torchtitan.models.attention import (
@@ -251,8 +251,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 +493,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}"