【auto parallel】fix cp ut (#73735)

Author: danleifeng

python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py

@@ -30,7 +30,7 @@
     from paddle.nn import Layer
 
 
-def c_split(x, process_mesh, need_transpose):
+def c_split(x, process_mesh, need_transpose, split_type="sp"):
     mp_index = process_mesh.dim_names.index('mp')  # get the axis for the split
     dp_index = process_mesh.dim_names.index('dp')
     if isinstance(x, tuple):
@@ -44,17 +44,23 @@ def c_split(x, process_mesh, need_transpose):
     placements = target_x.placements
     if placements is None:
         placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
-    if placements[dp_index] == dist.Shard(0):
-        # NOTE(zhangwl):if shard(0) , input shape should be [b,s,h]
-        split_dims = dist.Shard(1)
-    elif placements[dp_index] == dist.Shard(1):
-        # NOTE(zhangwl):if shard(1) , input shape should be [s,b,h]
-        split_dims = dist.Shard(0)
+    if split_type == "sp":
+        if placements[dp_index] == dist.Shard(0):
+            # NOTE(zhangwl):if shard(0) , input shape should be [b,s,h]
+            split_dims = dist.Shard(1)
+        elif placements[dp_index] == dist.Shard(1):
+            # NOTE(zhangwl):if shard(1) , input shape should be [s,b,h]
+            split_dims = dist.Shard(0)
+        else:
+            logging.warning(
+                f"parallel api don't know {target_x.shape} which dimension is batch, default is to cut to the 0th dimension"
+            )
+            split_dims = dist.Shard(0)
+    elif split_type == "mp":
+        split_dims = dist.Shard(2)  # split h [b,s,h]
     else:
-        logging.warning(
-            f"parallel api don't know {target_x.shape} which dimension is batch, default is to cut to the 0th dimension"
-        )
-        split_dims = dist.Shard(0)
+        raise ValueError(f"Unsupported split type {split_type}")
+
     placements[mp_index] = split_dims
     target_x = dist.reshard(target_x, process_mesh, placements)
     if isinstance(x, tuple):
@@ -251,7 +257,7 @@ def __init__(self, is_input_parallel: bool = True) -> None:
 
     def split_input_hook(self, process_mesh):
         def split_hook(layer, input):
-            return c_split(input, process_mesh, False)
+            return c_split(input, process_mesh, False, split_type="mp")
 
         return split_hook
 

python/paddle/distributed/auto_parallel/ring_attention.py

@@ -69,13 +69,14 @@ def unshard_seq_load_balance(tensor, seq_dim):
 
 class RingCommunicator:
     def __init__(self, group, local_key, local_value):
-        self._k_buffer = [paddle.zeros_like(local_key) for _ in range(2)]
-        self._v_buffer = [paddle.zeros_like(local_value) for _ in range(2)]
-
-        local_key = local_key.contiguous()
-        local_value = local_value.contiguous()
-        self._k_buffer[0] = local_key.clone()
-        self._v_buffer[0] = local_value.clone()
+        self._k_buffer = [
+            local_key.clone().contiguous(),
+            local_key.clone().contiguous(),
+        ]
+        self._v_buffer = [
+            local_value.clone().contiguous(),
+            local_value.clone().contiguous(),
+        ]
 
         self._next_buffer_idx = 0
 
@@ -97,8 +98,6 @@ def wait(self):
         paddle.device.synchronize()
 
     def add_to_buffers(self, key, value):
-        key = key.contiguous()
-        value = value.contiguous()
         if key.shape != self._k_buffer[self._next_buffer_idx].shape:
             self._k_buffer[self._next_buffer_idx][:, : key.shape[1], :, :].add_(
                 key
@@ -112,8 +111,8 @@ def add_to_buffers(self, key, value):
 
     def get_buffers(self):
         return (
-            self._k_buffer[self._next_buffer_idx].contiguous(),
-            self._v_buffer[self._next_buffer_idx].contiguous(),
+            self._k_buffer[self._next_buffer_idx],
+            self._v_buffer[self._next_buffer_idx],
         )
 
     def send_recv(self):
@@ -131,13 +130,13 @@ def send_recv(self):
         )
         recv_k_op = dist.P2POp(
             dist.irecv,
-            self._k_buffer[(self._next_buffer_idx + 1) % 2].contiguous(),
+            self._k_buffer[(self._next_buffer_idx + 1) % 2],
             self.recv_rank,
             self.group,
         )
         recv_v_op = dist.P2POp(
             dist.irecv,
-            self._v_buffer[(self._next_buffer_idx + 1) % 2].contiguous(),
+            self._v_buffer[(self._next_buffer_idx + 1) % 2],
             self.recv_rank,
             self.group,
         )
@@ -186,9 +185,9 @@ def concat_masks(attn_masks_list, rank, cp_size):
 
 def ring_flash_attention_forward_func(
     group,
-    query,
-    key,
-    value,
+    local_query,
+    local_key,
+    local_value,
     attn_mask=None,
     dropout=0.0,
     is_causal=False,
@@ -197,20 +196,6 @@ def ring_flash_attention_forward_func(
 ):
     cp_size = group.world_size
     group_rank = group.rank
-    mesh = dist.auto_parallel.get_mesh()
-
-    local_query = dist.auto_parallel.api.dtensor_to_local(
-        query, mesh, query.placements
-    )
-    local_key = dist.auto_parallel.api.dtensor_to_local(
-        key, mesh, key.placements
-    )
-    local_value = dist.auto_parallel.api.dtensor_to_local(
-        value, mesh, value.placements
-    )
-    local_query = local_query.contiguous()
-    local_key = local_key.contiguous()
-    local_value = local_value.contiguous()
 
     comm_buffer = RingCommunicator(group, local_key, local_value)
     local_q_seq_len = local_query.shape[1]
@@ -219,7 +204,9 @@ def ring_flash_attention_forward_func(
             attn_mask, num_or_sections=cp_size * 2, axis=3
         )
     if is_causal:
-        local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :]
+        local_query_second_chunk = local_query[
+            :, local_q_seq_len // 2 :, :, :
+        ].contiguous()
     for step in range(cp_size):
         block_k, block_v = comm_buffer.get_buffers()
         if step != cp_size - 1:
@@ -315,10 +302,10 @@ def ring_flash_attention_forward_func(
 
 def ring_flash_attention_backward_func(
     group,
-    out_grad,
-    query,
-    key,
-    value,
+    local_out_grad,
+    local_query,
+    local_key,
+    local_value,
     local_out,
     lse,
     attn_mask,
@@ -328,25 +315,7 @@ def ring_flash_attention_backward_func(
 ):
     cp_size = group.world_size
     group_rank = group.rank
-    mesh = dist.auto_parallel.get_mesh()
 
-    local_query = dist.auto_parallel.api.dtensor_to_local(
-        query, mesh, query.placements
-    )
-    local_key = dist.auto_parallel.api.dtensor_to_local(
-        key, mesh, key.placements
-    )
-    local_value = dist.auto_parallel.api.dtensor_to_local(
-        value, mesh, value.placements
-    )
-    local_out_grad = dist.auto_parallel.api.dtensor_to_local(
-        out_grad, mesh, out_grad.placements
-    )
-    local_query = local_query.contiguous()
-    local_key = local_key.contiguous()
-    local_value = local_value.contiguous()
-    local_out = local_out.contiguous()
-    local_out_grad = local_out_grad.contiguous()
     lse = lse.contiguous()
 
     local_q_seq_len = local_query.shape[1]
@@ -361,14 +330,13 @@ def ring_flash_attention_backward_func(
     if is_causal:
         local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :]
         local_out_second_chunk = local_out[:, local_q_seq_len // 2 :, :, :]
-        lse_second_chunk = lse[:, :, local_q_seq_len // 2 :]
+        lse_second_chunk = lse[:, :, local_q_seq_len // 2 :].contiguous()
         out_grad_second_chunk = local_out_grad[:, local_q_seq_len // 2 :, :, :]
 
     if attn_mask is not None:
         attn_masks_list = paddle.split(
             attn_mask, num_or_sections=cp_size * 2, axis=3
         )
-
     for step in range(cp_size):
         block_k, block_v = kv_comm_buffer.get_buffers()
         if step != cp_size - 1:
@@ -445,13 +413,14 @@ def ring_flash_attention_backward_func(
                     )
                 )
                 query_grad_buffer.add_(block_q_grad)
-
         paddle.device.synchronize()
 
-        grad_comm_buffer.add_to_buffers(block_k_grad, block_v_grad)
+        grad_comm_buffer.add_to_buffers(
+            block_k_grad.contiguous(), block_v_grad.contiguous()
+        )
         grad_comm_buffer.send_recv()
 
-    grad_comm_buffer.wait()
+        grad_comm_buffer.wait()
     key_grad_buffer, value_grad_buffer = grad_comm_buffer.get_buffers()
 
     return query_grad_buffer, key_grad_buffer, value_grad_buffer
@@ -481,14 +450,23 @@ def forward(
         dist.init_parallel_env()
 
         group = mesh._get_group("sep")
-
+        local_query = dist.auto_parallel.api.dtensor_to_local(
+            query, query.process_mesh, query.placements
+        )
+        local_key = dist.auto_parallel.api.dtensor_to_local(
+            key, key.process_mesh, key.placements
+        )
+        local_value = dist.auto_parallel.api.dtensor_to_local(
+            value, value.process_mesh, value.placements
+        )
         if attn_mask is not None:
             is_causal = False
+
         out, lse = ring_flash_attention_forward_func(
             group,
-            query,
-            key,
-            value,
+            local_query,
+            local_key,
+            local_value,
             attn_mask,
             dropout,
             is_causal,
@@ -500,7 +478,7 @@ def forward(
         ctx.dropout = dropout
         ctx.is_causal = is_causal
         out_dtensor = dist.auto_parallel.api.dtensor_from_local(
-            out.contiguous(), query.process_mesh, query.placements
+            out, query.process_mesh, query.placements
         )
         return out_dtensor.contiguous()
 
@@ -517,27 +495,39 @@ def backward(ctx, out_grad):
             fixed_seed_offset = paddle.to_tensor(
                 [0, 0], place=paddle.CPUPlace(), dtype=paddle.int64
             )
+        local_query = dist.auto_parallel.api.dtensor_to_local(
+            query, query.process_mesh, query.placements
+        )
+        local_key = dist.auto_parallel.api.dtensor_to_local(
+            key, key.process_mesh, key.placements
+        )
+        local_value = dist.auto_parallel.api.dtensor_to_local(
+            value, value.process_mesh, value.placements
+        )
+        local_out_grad = dist.auto_parallel.api.dtensor_to_local(
+            out_grad, out_grad.process_mesh, out_grad.placements
+        )
         query_grad, key_grad, value_grad = ring_flash_attention_backward_func(
             group,
-            out_grad,
-            query,
-            key,
-            value,
+            local_out_grad,
+            local_query,
+            local_key,
+            local_value,
             out,
-            lse.contiguous(),
+            lse,
             attn_mask,
             dropout,
             is_causal,
             fixed_seed_offset,
         )
         query_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
-            query_grad.contiguous(), query.process_mesh, query.placements
+            query_grad, query.process_mesh, query.placements
         )
         key_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
-            key_grad.contiguous(), key.process_mesh, key.placements
+            key_grad, key.process_mesh, key.placements
         )
         value_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
-            value_grad.contiguous(), value.process_mesh, value.placements
+            value_grad, value.process_mesh, value.placements
         )
 
         if attn_mask is not None and not attn_mask.stop_gradient:

test/auto_parallel/hybrid_strategy/CMakeLists.txt

@@ -137,6 +137,14 @@ if((WITH_GPU) AND (LINUX))
   set_tests_properties(test_parallel_api_with_llama_2d
                        PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
 endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_2d_sep MODULES
+    test_parallel_api_with_llama_2d_sep ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_2d_sep
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
+endif()
 if((WITH_GPU) AND (LINUX))
   py_test_modules(
     test_parallel_api_with_llama_3d MODULES test_parallel_api_with_llama_3d

test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d_sep.py

@@ -37,7 +37,7 @@ def setUp(self):
             "amp_dtype": ["bfloat16"],
             "amp_master_grad": ["true"],
             "use_lazy_init": ["true"],
-            "sequence_parallel": ["true"],
+            "sequence_parallel": ["false"],
             "context_parallel": ["true"],
             "prepare_input_output": ["false"],
             "sharding_stage": ["0"],
@@ -83,7 +83,7 @@ def setUp(self):
             "amp_dtype": ["bfloat16"],
             "amp_master_grad": ["true"],
             "use_lazy_init": ["true"],
-            "sequence_parallel": ["true"],
+            "sequence_parallel": ["false"],
             "sep_parallel": ["true"],
             "context_parallel": ["false"],
             "prepare_input_output": ["false"],

test/auto_parallel/hybrid_strategy/testslist.csv

@@ -15,6 +15,7 @@ test_semi_auto_llama_acc_align,LINUX,GPU,300,HYBRID,test_runner.py,,,http_proxy=
 test_semi_auto_llama_save_load,LINUX,GPU,180,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..;FLAGS_enable_pir_api=1,
 test_parallel_api_with_llama_1d,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
 test_parallel_api_with_llama_2d,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_parallel_api_with_llama_2d_sep,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
 test_parallel_api_with_llama_3d,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
 test_parallel_api_with_llama_lora,LINUX,GPU,360,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
 test_process_mesh,LINUX,GPU,60,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,