Update on "[cp] dispatch flex_attention_backward to CP impl in TorchD…

…ispatchMode" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k [ghstack-poisoned]
pytorch · XilunWu · Apr 28, 2025 · Apr 28, 2025 · Apr 28, 2025 · Apr 28, 2025
commit 38d5e2dfd57da61bcc5b6ac9ceabc3e46cfb61bc
diff --git a/test/distributed/tensor/test_attention.py b/test/distributed/tensor/test_attention.py
@@ -449,7 +449,6 @@ def world_size(self) -> int:
     def test_ring_flex_attention(self) -> None:
         def causal_mask(b, h, q_idx, kv_idx):
             return q_idx >= kv_idx
-            # return q_idx >= 0
 
         from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 

diff --git a/torch/distributed/tensor/experimental/_attention.py b/torch/distributed/tensor/experimental/_attention.py
@@ -1509,9 +1509,14 @@ def context_parallel(
     ):
         yield
 
-    return
     for buffer, original_buffer in zip(buffers, original_buffers):
         if original_buffer is not None:
+            # tensor cannot resize if requires_grad is True
+            # key and value's requires_grad has been set to False in manual comm calls
+            # unless via DTensor.
+            if buffer.requires_grad:
+                buffer.requires_grad = False
+
             buffer.resize_(original_buffer.shape)
             buffer.copy_(original_buffer)