[cp][flex_attention] integration test trial

ghstack-source-id: 7e12a16
Pull-Request-resolved: #1160

Author: XilunWu

torchtitan/distributed/utils.py

@@ -9,13 +9,16 @@
 import os
 from collections.abc import Generator, Iterable
 from datetime import timedelta
+from typing import Optional
 
 import torch
 import torch.distributed._functional_collectives as funcol
 import torch.distributed.distributed_c10d as c10d
 from torch import distributed as dist
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.experimental._attention import FlexAttentionSharder
+from torch.nn.attention.flex_attention import BlockMask
 
 from torchtitan.tools.logging import logger
 from torchtitan.tools.utils import device_module, device_type
@@ -154,22 +157,31 @@ def create_context_parallel_ctx(
     cp_seq_dims: list[int],
     cp_no_restore_buffers: set[torch.Tensor],
     cp_rotate_method: str,
+    block_mask: Optional[BlockMask] = None,
+    sharder: Optional[FlexAttentionSharder] = None,
 ):
     try:
         from torch.distributed.tensor.experimental import context_parallel
-        from torch.distributed.tensor.experimental._attention import set_rotate_method
+        from torch.distributed.tensor.experimental._attention import (
+            _dispatch_mode,
+            _DispatchMode,
+            set_rotate_method,
+        )
     except ImportError:
         print(
             f"PyTorch version {torch.__version__} does not include the experimental "
             "Context Parallel API. Please update to a newer version."
         )
 
     set_rotate_method(cp_rotate_method)
+    _dispatch_mode = _DispatchMode.TORCH_DISPATCH
     return context_parallel(
         cp_mesh,
         buffers=cp_buffers,
         buffer_seq_dims=cp_seq_dims,
         no_restore_buffers=cp_no_restore_buffers,
+        block_mask=block_mask,
+        sharder=sharder,
     )
 
 

torchtitan/experiments/llama4/model/args.py

@@ -74,12 +74,13 @@ def update_from_config(self, job_config: JobConfig, tokenizer: Tokenizer) -> Non
                 "FlexAttention is not compatible with selective AC yet. "
                 "See https://github.com/pytorch/pytorch/issues/147879"
             )
-
+        """
         if job_config.parallelism.context_parallel_degree > 1 and self.use_flex_attn:
             raise ValueError(
                 "FlexAttention is not compatible with CP yet. "
                 "We are still working on this."
             )
+        """
 
     def get_nparams_and_flops(
         self, model: nn.Module, seq_len: int

torchtitan/train.py

@@ -11,10 +11,13 @@
 from typing import Any, Generator, Iterable, Optional
 
 import torch
-from torch.distributed.elastic.multiprocessing.errors import record
 
 import torchtitan.components.ft as ft
 import torchtitan.protocols.train_spec as train_spec_module
+from torch.distributed.elastic.multiprocessing.errors import record
+from torch.distributed.tensor.experimental._attention import (
+    FlexAttentionContiguousSharder,
+)
 
 from torchtitan.components.checkpoint import CheckpointManager
 from torchtitan.components.metrics import (
@@ -133,7 +136,9 @@ def __init__(self, job_config: JobConfig):
 
         # build model (using meta init)
         model_cls = self.train_spec.cls
+        # NOTE (xilunwu): need to store model_args.use_flex_attn for train_step
         model_args = self.train_spec.config[job_config.model.flavor]
+        self.model_args = model_args
         # set the model args from training job configs
         model_args.update_from_config(job_config, tokenizer)
 
@@ -319,13 +324,29 @@ def train_step(self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor):
         # apply context parallelism if cp is enabled
         # ensure CP handles the separate freqs_cis buffer for each pp stage
         inputs = input_dict["input"]
+
+        # TODO: move this into `create_context_parallel_ctx`
+        # init block_mask for flex_attention
+        block_mask = None
+        if self.model_args.use_flex_attn:
+            from torchtitan.models.attention import FlexAttention
+
+            mask_mod = FlexAttention._get_causal_mask_mod()
+            batch_dimension = 1
+            seq_len = inputs.shape[1]
+            block_mask = FlexAttention.compiled_create_block_mask(
+                mask_mod, batch_dimension, None, seq_len, seq_len
+            )
+
         optional_context_parallel_ctx = (
             dist_utils.create_context_parallel_ctx(
                 cp_mesh=world_mesh["cp"],
                 cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts],
                 cp_seq_dims=[1, 1] + [0 for _ in model_parts],
                 cp_no_restore_buffers={inputs, labels},
                 cp_rotate_method=self.job_config.parallelism.context_parallel_rotate_method,
+                block_mask=block_mask,
+                sharder=FlexAttentionContiguousSharder(),
             )
             if parallel_dims.cp_enabled
             else None