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

XilunWu · XilunWu · commit 38d5e2dfd57d · 2025-04-29T17:20:58.000-07:00
cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k

[ghstack-poisoned]
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)
 
