update

Author: zhangyuqin1998

paddlenlp/transformers/ring_flash_attention.py

@@ -50,9 +50,7 @@ def __init__(self, group, local_key, local_value):
         self._reqs = []
 
     def wait(self):
-        # for req in self._reqs:
-        #     req.wait()
-        # self._reqs = None
+        # TODO(zhangyuqin1998)：batch_isend_irecv异步流下，无法wait，需要修复。对性能有影响。
         paddle.device.synchronize()
 
     def add_to_buffers(self, key, value):
@@ -126,12 +124,14 @@ def balanced_ring_flash_attention_fwd_func(
     comm_buffer = RingCommunicator(group, local_key, local_value)
     local_q_seq_len = local_query.shape[1]
 
-    out, lse = None, None
+    out, lse, k_cache, v_cache = None, None, dict(), dict()
 
     if attn_mask is not None:
         attn_masks_list = paddle.split(attn_mask, num_or_sections=cp_size * 2, axis=3)
     if is_causal:
-        local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :].clone().contiguous()
+        local_query_second_chunk = paddle.slice(
+            local_query, axes=[1], starts=[local_q_seq_len // 2], ends=[local_q_seq_len]
+        )
     for step in range(cp_size):
         block_k, block_v = comm_buffer.get_buffers()
 
@@ -153,14 +153,15 @@ def balanced_ring_flash_attention_fwd_func(
                 not training,
                 "",
             )
-            block_lse = paddle.unsqueeze(paddle.transpose(block_lse, [0, 2, 1]), axis=-1)
+            block_lse = paddle.unsqueeze_(paddle.transpose_(block_lse, [0, 2, 1]), axis=-1)
             out, lse = update_out_and_lse(out, lse, block_out, block_lse)
         else:
+            # block_k and block_v is from rank (group.rank - step) % cp_size
             if step == 0:
                 block_out, _, block_lse, _ = _C_ops.flash_attn(
                     local_query, block_k, block_v, fixed_seed_offset, None, dropout, True, False, not training, ""
                 )
-                block_lse = paddle.unsqueeze(paddle.transpose(block_lse, [0, 2, 1]), axis=-1)
+                block_lse = paddle.unsqueeze_(paddle.transpose_(block_lse, [0, 2, 1]), axis=-1)
                 out, lse = update_out_and_lse(out, lse, block_out, block_lse)
             elif step > rank:
                 block_out, _, block_lse, _ = _C_ops.flash_attn(
@@ -175,14 +176,16 @@ def balanced_ring_flash_attention_fwd_func(
                     not training,
                     "",
                 )
-                block_lse = block_lse[:, :, 0 : (local_q_seq_len // 2)]
-                block_lse = paddle.unsqueeze(paddle.transpose(block_lse, [0, 2, 1]), axis=-1)
+                block_lse = paddle.slice(block_lse, axes=[1], starts=[0], ends=[local_q_seq_len // 2])
+                block_lse = paddle.unsqueeze_(paddle.transpose_(block_lse, [0, 2, 1]), axis=-1)
                 out, lse = update_out_and_lse(out, lse, block_out, block_lse, True)
             else:
+                block_k = paddle.slice(block_k, axes=[1], starts=[0], ends=[local_q_seq_len // 2])
+                block_v = paddle.slice(block_v, axes=[1], starts=[0], ends=[local_q_seq_len // 2])
                 block_out, _, block_lse, _ = _C_ops.flash_attn(
                     local_query,
-                    block_k[:, : local_q_seq_len // 2, :, :],
-                    block_v[:, : local_q_seq_len // 2, :, :],
+                    block_k,
+                    block_v,
                     fixed_seed_offset,
                     None,
                     dropout,
@@ -191,20 +194,23 @@ def balanced_ring_flash_attention_fwd_func(
                     not training,
                     "",
                 )
-                block_lse = paddle.unsqueeze(paddle.transpose(block_lse, [0, 2, 1]), axis=-1)
+                block_lse = paddle.unsqueeze_(paddle.transpose_(block_lse, [0, 2, 1]), axis=-1)
                 out, lse = update_out_and_lse(out, lse, block_out, block_lse)
+                k_cache[step] = block_k
+                v_cache[step] = block_v
 
-        # if step != cp_size - 1:
-        #     comm_buffer.wait()
+        # TODO(zhangyuqin1998)：batch_isend_irecv异步流下，无法wait，需要修复。对性能有影响。
         paddle.device.synchronize()
 
     out = out.to(local_query.dtype)
-    lse = paddle.transpose(paddle.squeeze(lse, axis=-1), [0, 2, 1])
-    return out, lse
+    lse = paddle.transpose_(paddle.squeeze_(lse, axis=-1), [0, 2, 1])
+    return out, lse, k_cache, v_cache
 
 
 def balanced_ring_flash_attention_bwd_func(
     group,
+    k_cache,
+    v_cache,
     out_grad,
     local_query,
     local_key,
@@ -228,10 +234,14 @@ def balanced_ring_flash_attention_bwd_func(
     grad_comm_buffer = RingCommunicator(group, key_grad_buffer, value_grad_buffer)
 
     if is_causal:
-        local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :].clone().contiguous()
-        local_out_second_chunk = local_out[:, local_q_seq_len // 2 :, :, :].clone().contiguous()
-        lse_second_chunk = lse[:, :, local_q_seq_len // 2 :].clone().contiguous()
-        out_grad_second_chunk = out_grad[:, local_q_seq_len // 2 :, :, :].clone().contiguous()
+        local_query_second_chunk = paddle.slice(
+            local_query, axes=[1], starts=[local_q_seq_len // 2], ends=[local_q_seq_len]
+        )
+        local_out_second_chunk = paddle.slice(
+            local_out, axes=[1], starts=[local_q_seq_len // 2], ends=[local_q_seq_len]
+        )
+        lse_second_chunk = paddle.slice(lse, axes=[2], starts=[local_q_seq_len // 2], ends=[local_q_seq_len])
+        out_grad_second_chunk = paddle.slice(out_grad, axes=[1], starts=[local_q_seq_len // 2], ends=[local_q_seq_len])
 
     if attn_mask is not None:
         attn_masks_list = paddle.split(attn_mask, num_or_sections=cp_size * 2, axis=3)
@@ -279,8 +289,8 @@ def balanced_ring_flash_attention_bwd_func(
             else:
                 block_q_grad, block_k_grad, block_v_grad = flash_attn_bwd(
                     local_query,
-                    block_k[:, : local_q_seq_len // 2, :, :],
-                    block_v[:, : local_q_seq_len // 2, :, :],
+                    k_cache[step],
+                    v_cache[step],
                     local_out,
                     lse,
                     fixed_seed_offset,
@@ -291,10 +301,7 @@ def balanced_ring_flash_attention_bwd_func(
                 )
                 query_grad_buffer += block_q_grad
 
-        # if step != cp_size - 1:
-        #     kv_comm_buffer.wait()
-        # if step != 0:
-        #     grad_comm_buffer.wait()
+        # TODO(zhangyuqin1998)：batch_isend_irecv异步流下，无法wait，需要修复。对性能有影响。
         paddle.device.synchronize()
 
         grad_comm_buffer.add_to_buffers(block_k_grad, block_v_grad)
@@ -328,10 +335,10 @@ def forward(
         if attn_mask is not None:
             is_causal = False
 
-        out, lse = balanced_ring_flash_attention_fwd_func(
+        out, lse, k_cache, v_cache = balanced_ring_flash_attention_fwd_func(
             group, query, key, value, fixed_seed_offset, attn_mask, dropout, is_causal, training
         )
-        ctx.save_for_backward(query, key, value, out, lse, attn_mask)
+        ctx.save_for_backward(query, key, value, out, lse, attn_mask, k_cache, v_cache)
         ctx.group = group
         ctx.fixed_seed_offset = fixed_seed_offset
         ctx.dropout = dropout
@@ -340,17 +347,29 @@ def forward(
 
     @staticmethod
     def backward(ctx, out_grad):
-        query, key, value, out, lse, attn_mask = ctx.saved_tensor()
+        query, key, value, out, lse, attn_mask, k_cache, v_cache = ctx.saved_tensor()
         group = ctx.group
         fixed_seed_offset = ctx.fixed_seed_offset
         dropout = ctx.dropout
         is_causal = ctx.is_causal
 
         if fixed_seed_offset is None:
-            fixed_seed_offset = paddle.to_tensor([0, 0], place=paddle.CPUPlace(), dtype=paddle.int64).contiguous()
+            fixed_seed_offset = paddle.to_tensor([0, 0], place=paddle.CPUPlace(), dtype=paddle.int64)
 
         query_grad, key_grad, value_grad = balanced_ring_flash_attention_bwd_func(
-            group, out_grad, query, key, value, out, lse, fixed_seed_offset, attn_mask, dropout, is_causal
+            group,
+            k_cache,
+            v_cache,
+            out_grad,
+            query,
+            key,
+            value,
+            out,
+            lse,
+            fixed_seed_offset,
+            attn_mask,
+            dropout,
+            is_causal,
         )
         if attn_mask is not None and not attn_mask.stop_gradient:
             return query_grad, key_grad, value_grad, None