chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,542 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user