wehub-resource-sync
543 lines
18 KiB
Python
543 lines
18 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License
|
|
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
import paddle.nn.functional as F
|
|
from paddle import _C_ops
|
|
|
|
|
|
def shard_seq_load_balance(tensor, seq_dim):
|
|
# dtensor Replicate() -> reorder -> Shard(seq_dim)
|
|
placements = tensor.placements
|
|
process_mesh = tensor.process_mesh
|
|
cp_index = process_mesh.dim_names.index('sep')
|
|
cp_degree = process_mesh.shape[cp_index]
|
|
if cp_degree > 1:
|
|
# split
|
|
sliced_datas = paddle.split(
|
|
tensor, num_or_sections=cp_degree * 2, axis=seq_dim
|
|
)
|
|
# resort [q0,q1,q2,q3] -> [q0,q3,q1,q2]
|
|
indices = []
|
|
for i in range(cp_degree):
|
|
indices.append(i)
|
|
indices.append(cp_degree * 2 - 1 - i)
|
|
reorder_indices = indices
|
|
reordered = [sliced_datas[i] for i in reorder_indices]
|
|
reordered_tensor = paddle.concat(reordered, axis=seq_dim)
|
|
# reshard q/k/v -> Shard(seq_dim)
|
|
placements[cp_index] = paddle.distributed.Shard(seq_dim) # seq_dim:1
|
|
tensor = paddle.distributed.reshard(
|
|
reordered_tensor, process_mesh, placements
|
|
)
|
|
return tensor
|
|
|
|
|
|
def unshard_seq_load_balance(tensor, seq_dim):
|
|
# dtensor Shard(seq_dim) -> Replicate() -> reorder
|
|
placements = tensor.placements
|
|
process_mesh = tensor.process_mesh
|
|
cp_index = process_mesh.dim_names.index('sep')
|
|
cp_degree = process_mesh.shape[cp_index]
|
|
all_tensor = dist.reshard(tensor, process_mesh, [dist.Replicate()])
|
|
sliced_datas = paddle.split(
|
|
all_tensor, num_or_sections=cp_degree * 2, axis=seq_dim
|
|
)
|
|
reorder_indices = []
|
|
for i in range(cp_degree):
|
|
reorder_indices.append(i)
|
|
reorder_indices.append(cp_degree * 2 - 1 - i)
|
|
inverse_indices = [0] * len(reorder_indices)
|
|
for idx, v in enumerate(reorder_indices):
|
|
inverse_indices[v] = idx
|
|
restored = [sliced_datas[i] for i in inverse_indices]
|
|
return paddle.concat(restored, axis=seq_dim)
|
|
|
|
|
|
class RingCommunicator:
|
|
def __init__(self, group, local_key, local_value):
|
|
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
|
|
|
|
self.group = group
|
|
mesh = dist.auto_parallel.get_mesh()
|
|
process_id = dist.get_rank()
|
|
self.group_rank = mesh.get_rank_by_dim_and_process_id("sep", process_id)
|
|
self.cp_size = mesh.get_dim_size("sep")
|
|
cp_index = mesh.dim_names.index("sep")
|
|
|
|
self.send_rank = self.group.ranks[
|
|
(self.group_rank + 1) % self.cp_size
|
|
] # 1%2=1
|
|
self.recv_rank = self.group.ranks[(self.group_rank - 1) % self.cp_size]
|
|
|
|
self._reqs = []
|
|
|
|
def wait(self):
|
|
paddle.device.synchronize()
|
|
|
|
def add_to_buffers(self, key, value):
|
|
if key.shape != self._k_buffer[self._next_buffer_idx].shape:
|
|
self._k_buffer[self._next_buffer_idx][:, : key.shape[1], :, :].add_(
|
|
key
|
|
)
|
|
self._v_buffer[self._next_buffer_idx][:, : key.shape[1], :, :].add_(
|
|
value
|
|
)
|
|
else:
|
|
self._k_buffer[self._next_buffer_idx].add_(key)
|
|
self._v_buffer[self._next_buffer_idx].add_(value)
|
|
|
|
def get_buffers(self):
|
|
return (
|
|
self._k_buffer[self._next_buffer_idx],
|
|
self._v_buffer[self._next_buffer_idx],
|
|
)
|
|
|
|
def send_recv(self):
|
|
send_k_op = dist.P2POp(
|
|
dist.isend,
|
|
self._k_buffer[self._next_buffer_idx].contiguous(),
|
|
self.send_rank,
|
|
self.group,
|
|
)
|
|
send_v_op = dist.P2POp(
|
|
dist.isend,
|
|
self._v_buffer[self._next_buffer_idx].contiguous(),
|
|
self.send_rank,
|
|
self.group,
|
|
)
|
|
recv_k_op = dist.P2POp(
|
|
dist.irecv,
|
|
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],
|
|
self.recv_rank,
|
|
self.group,
|
|
)
|
|
|
|
self._next_buffer_idx = (self._next_buffer_idx + 1) % 2
|
|
|
|
ops = [send_k_op, send_v_op, recv_k_op, recv_v_op]
|
|
|
|
self._reqs = dist.batch_isend_irecv(ops)
|
|
|
|
|
|
def update_out_and_lse(
|
|
old_out, old_lse, block_out, block_lse, second_chunk_only=False
|
|
):
|
|
if second_chunk_only:
|
|
second_chunk_out = old_out[:, old_out.shape[1] // 2 :, :, :]
|
|
second_chunk_lse = old_lse[:, old_lse.shape[1] // 2 :, :, :]
|
|
second_chunk_out, second_chunk_lse = update_out_and_lse(
|
|
second_chunk_out, second_chunk_lse, block_out, block_lse
|
|
)
|
|
old_out[:, old_out.shape[1] // 2 :, :, :] = second_chunk_out
|
|
old_lse[:, old_lse.shape[1] // 2 :, :, :] = second_chunk_lse
|
|
return old_out, old_lse
|
|
else:
|
|
block_out, block_lse = (
|
|
paddle.cast(block_out, "float32"),
|
|
paddle.cast(block_lse, "float32"),
|
|
)
|
|
with paddle.amp.auto_cast(enable=False):
|
|
return old_out - (old_out - block_out) * F.sigmoid(
|
|
block_lse - old_lse
|
|
), old_lse - F.log_sigmoid(old_lse - block_lse)
|
|
|
|
|
|
def get_chunk_id(rank, cp_size):
|
|
return rank, (2 * cp_size - 1 - rank)
|
|
|
|
|
|
def concat_masks(attn_masks_list, rank, cp_size):
|
|
assert len(attn_masks_list) == 2 * cp_size
|
|
first_chunk_id, second_chunk_id = get_chunk_id(rank, cp_size)
|
|
return paddle.concat(
|
|
[attn_masks_list[first_chunk_id], attn_masks_list[second_chunk_id]],
|
|
axis=3,
|
|
)
|
|
|
|
|
|
def ring_flash_attention_forward_func(
|
|
group,
|
|
local_query,
|
|
local_key,
|
|
local_value,
|
|
attn_mask=None,
|
|
dropout=0.0,
|
|
is_causal=False,
|
|
fixed_seed_offset=None,
|
|
training=True,
|
|
):
|
|
cp_size = group.world_size
|
|
group_rank = group.rank
|
|
|
|
comm_buffer = RingCommunicator(group, local_key, local_value)
|
|
local_q_seq_len = local_query.shape[1]
|
|
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 :, :, :
|
|
].contiguous()
|
|
for step in range(cp_size):
|
|
block_k, block_v = comm_buffer.get_buffers()
|
|
if step != cp_size - 1:
|
|
comm_buffer.send_recv()
|
|
if not is_causal:
|
|
# out [bs, seq, nhead, headdim]
|
|
# lse [bs, nhead, seq]
|
|
block_out, _, block_lse, _ = _C_ops.flash_attn(
|
|
local_query,
|
|
block_k,
|
|
block_v,
|
|
fixed_seed_offset,
|
|
(
|
|
None
|
|
if attn_mask is None
|
|
else concat_masks(
|
|
attn_masks_list, (group_rank - step) % cp_size, cp_size
|
|
)
|
|
),
|
|
dropout,
|
|
False,
|
|
False,
|
|
not training,
|
|
"",
|
|
)
|
|
paddle.unsqueeze_(paddle.transpose_(block_lse, [0, 2, 1]), axis=-1)
|
|
|
|
if step == 0:
|
|
out, lse = block_out, block_lse
|
|
else:
|
|
out, lse = update_out_and_lse(out, lse, block_out, block_lse)
|
|
else:
|
|
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,
|
|
"",
|
|
)
|
|
paddle.unsqueeze_(
|
|
paddle.transpose_(block_lse, [0, 2, 1]), axis=-1
|
|
)
|
|
out, lse = block_out, block_lse
|
|
elif step > group_rank:
|
|
block_out, _, block_lse, _ = _C_ops.flash_attn(
|
|
local_query_second_chunk,
|
|
block_k,
|
|
block_v,
|
|
fixed_seed_offset,
|
|
None,
|
|
dropout,
|
|
False,
|
|
False,
|
|
not training,
|
|
"",
|
|
)
|
|
block_lse = block_lse[:, :, 0 : (local_q_seq_len // 2)]
|
|
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_out, _, block_lse, _ = _C_ops.flash_attn(
|
|
local_query,
|
|
block_k[:, : local_q_seq_len // 2, :, :],
|
|
block_v[:, : local_q_seq_len // 2, :, :],
|
|
fixed_seed_offset,
|
|
None,
|
|
dropout,
|
|
False,
|
|
False,
|
|
not training,
|
|
"",
|
|
)
|
|
paddle.unsqueeze_(
|
|
paddle.transpose_(block_lse, [0, 2, 1]), axis=-1
|
|
)
|
|
out, lse = update_out_and_lse(out, lse, block_out, block_lse)
|
|
|
|
paddle.device.synchronize()
|
|
out = paddle.cast(out, local_query.dtype)
|
|
lse = paddle.transpose_(paddle.squeeze(lse, axis=-1), [0, 2, 1])
|
|
return out, lse
|
|
|
|
|
|
def ring_flash_attention_backward_func(
|
|
group,
|
|
local_out_grad,
|
|
local_query,
|
|
local_key,
|
|
local_value,
|
|
local_out,
|
|
lse,
|
|
attn_mask,
|
|
dropout=0.0,
|
|
is_causal=False,
|
|
fixed_seed_offset=None,
|
|
):
|
|
cp_size = group.world_size
|
|
group_rank = group.rank
|
|
|
|
lse = lse.contiguous()
|
|
|
|
local_q_seq_len = local_query.shape[1]
|
|
query_grad_buffer = paddle.zeros_like(local_query)
|
|
key_grad_buffer = paddle.zeros_like(local_key)
|
|
value_grad_buffer = paddle.zeros_like(local_value)
|
|
|
|
kv_comm_buffer = RingCommunicator(group, local_key, local_value)
|
|
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 :, :, :]
|
|
local_out_second_chunk = local_out[:, 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:
|
|
kv_comm_buffer.send_recv()
|
|
|
|
if not is_causal:
|
|
block_q_grad, block_k_grad, block_v_grad = _C_ops.flash_attn_grad(
|
|
local_query,
|
|
block_k,
|
|
block_v,
|
|
local_out,
|
|
lse,
|
|
fixed_seed_offset,
|
|
(
|
|
None
|
|
if attn_mask is None
|
|
else concat_masks(
|
|
attn_masks_list, (group_rank - step) % cp_size, cp_size
|
|
)
|
|
),
|
|
local_out_grad,
|
|
dropout,
|
|
False,
|
|
)
|
|
query_grad_buffer.add_(block_q_grad)
|
|
else:
|
|
if step == 0:
|
|
block_q_grad, block_k_grad, block_v_grad = (
|
|
_C_ops.flash_attn_grad(
|
|
local_query,
|
|
block_k,
|
|
block_v,
|
|
local_out,
|
|
lse,
|
|
fixed_seed_offset,
|
|
None,
|
|
local_out_grad,
|
|
dropout,
|
|
True,
|
|
)
|
|
)
|
|
query_grad_buffer.add_(block_q_grad)
|
|
elif step > group_rank:
|
|
block_q_grad, block_k_grad, block_v_grad = (
|
|
_C_ops.flash_attn_grad(
|
|
local_query_second_chunk,
|
|
block_k,
|
|
block_v,
|
|
local_out_second_chunk,
|
|
lse_second_chunk,
|
|
fixed_seed_offset,
|
|
None,
|
|
out_grad_second_chunk,
|
|
dropout,
|
|
False,
|
|
)
|
|
)
|
|
query_grad_buffer[:, local_q_seq_len // 2 :, :, :].add_(
|
|
block_q_grad
|
|
)
|
|
else:
|
|
block_q_grad, block_k_grad, block_v_grad = (
|
|
_C_ops.flash_attn_grad(
|
|
local_query,
|
|
block_k[:, : local_q_seq_len // 2, :, :],
|
|
block_v[:, : local_q_seq_len // 2, :, :],
|
|
local_out,
|
|
lse,
|
|
fixed_seed_offset,
|
|
None,
|
|
local_out_grad,
|
|
dropout,
|
|
False,
|
|
)
|
|
)
|
|
query_grad_buffer.add_(block_q_grad)
|
|
paddle.device.synchronize()
|
|
|
|
grad_comm_buffer.add_to_buffers(
|
|
block_k_grad.contiguous(), block_v_grad.contiguous()
|
|
)
|
|
grad_comm_buffer.send_recv()
|
|
|
|
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
|
|
|
|
|
|
class RingFlashAttention(paddle.autograd.PyLayer):
|
|
@staticmethod
|
|
def forward(
|
|
ctx,
|
|
query,
|
|
key,
|
|
value,
|
|
attn_mask=None,
|
|
dropout=0.0,
|
|
is_causal=False,
|
|
fixed_seed_offset=None,
|
|
training=True,
|
|
):
|
|
if dropout > 0.0:
|
|
raise NotImplementedError(
|
|
"Dropout is not supported in ring attention yet."
|
|
)
|
|
mesh = dist.auto_parallel.get_mesh()
|
|
cp_index = mesh.dim_names.index('sep')
|
|
process_id = dist.get_rank()
|
|
rank = mesh.get_rank_by_dim_and_process_id("sep", process_id)
|
|
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,
|
|
local_query,
|
|
local_key,
|
|
local_value,
|
|
attn_mask,
|
|
dropout,
|
|
is_causal,
|
|
fixed_seed_offset,
|
|
training,
|
|
)
|
|
ctx.save_for_backward(group, query, key, value, out, lse, attn_mask)
|
|
ctx.fixed_seed_offset = fixed_seed_offset
|
|
ctx.dropout = dropout
|
|
ctx.is_causal = is_causal
|
|
out_dtensor = dist.auto_parallel.api.dtensor_from_local(
|
|
out, query.process_mesh, query.placements
|
|
)
|
|
return out_dtensor.contiguous()
|
|
|
|
@staticmethod
|
|
def backward(ctx, out_grad):
|
|
mesh = dist.auto_parallel.get_mesh()
|
|
cp_index = mesh.dim_names.index('sep')
|
|
group, query, key, value, out, lse, attn_mask = ctx.saved_tensor()
|
|
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
|
|
)
|
|
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,
|
|
local_out_grad,
|
|
local_query,
|
|
local_key,
|
|
local_value,
|
|
out,
|
|
lse,
|
|
attn_mask,
|
|
dropout,
|
|
is_causal,
|
|
fixed_seed_offset,
|
|
)
|
|
query_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
|
|
query_grad, query.process_mesh, query.placements
|
|
)
|
|
key_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
|
|
key_grad, key.process_mesh, key.placements
|
|
)
|
|
value_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
|
|
value_grad, value.process_mesh, value.placements
|
|
)
|
|
|
|
if attn_mask is not None and not attn_mask.stop_gradient:
|
|
return (
|
|
query_grad_dtensor,
|
|
key_grad_dtensor,
|
|
value_grad_dtensor,
|
|
None,
|
|
)
|
|
else:
|
|
return query_grad_dtensor, key_grad_dtensor, value_grad_dtensor
|