wehub-resource-sync
148 lines
6.0 KiB
Python
148 lines
6.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch
|
|
import torch.distributed as dist
|
|
from typing import Any, Tuple
|
|
|
|
|
|
# Code borrowed from deepspeed, here is why:
|
|
# 1. Reduce the dependency
|
|
# 2. The original code is complex
|
|
def _generate_layout_params(scatter_idx, seq_world_size, input):
|
|
if scatter_idx < 2:
|
|
bs, global_seq_len, num_local_head, head_dim = input.shape
|
|
pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim]
|
|
pre_all2all_permute_idx = (1, 0, 2, 3, 4)
|
|
|
|
post_all2all_permute_idx = (1, 2, 0, 3, 4)
|
|
post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim]
|
|
else:
|
|
bs, local_seq_len, num_total_head, head_dim = input.shape
|
|
assert num_total_head % seq_world_size == 0, (f'Number of heads ({num_total_head}) must be divisible '
|
|
f'by the sequence parallel size ({seq_world_size})!')
|
|
pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim]
|
|
pre_all2all_permute_idx = (2, 0, 1, 3, 4)
|
|
|
|
post_all2all_permute_idx = (1, 0, 2, 3, 4)
|
|
post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim]
|
|
|
|
return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape
|
|
|
|
|
|
def post_all2all(permute_idx, res_shape):
|
|
"""
|
|
Post-processing function for `all2all` communication.
|
|
"""
|
|
|
|
def post_func(input):
|
|
if permute_idx is not None:
|
|
input = input.permute(permute_idx).contiguous()
|
|
output = input.reshape(res_shape).contiguous()
|
|
|
|
return output
|
|
|
|
return post_func
|
|
|
|
|
|
def pre_all2all_fun(permute_idx, inp_shape, input):
|
|
"""
|
|
Pre-processing function for `all2all` communication.
|
|
"""
|
|
input_t = input.reshape(inp_shape).contiguous()
|
|
if permute_idx is not None:
|
|
input_t = input_t.permute(permute_idx).contiguous()
|
|
return input_t
|
|
|
|
|
|
def single_all_to_all(input, scatter_idx, gather_idx, group, **kwargs):
|
|
seq_world_size = dist.get_world_size(group)
|
|
num_heads = input.shape[2]
|
|
if num_heads % seq_world_size != 0 and not scatter_idx < 2:
|
|
raise NotImplementedError(f'num_heads {num_heads} cannot be split by sp world size {seq_world_size}')
|
|
pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = (
|
|
_generate_layout_params(scatter_idx, seq_world_size, input))
|
|
|
|
input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input)
|
|
|
|
post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape)
|
|
output = torch.empty_like(input_t)
|
|
dist.all_to_all_single(output, input_t, group=group)
|
|
|
|
res = post_all2all_fun(output)
|
|
return res
|
|
|
|
|
|
class _SeqAllToAll(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(
|
|
ctx: Any,
|
|
group: dist.ProcessGroup,
|
|
input: torch.Tensor,
|
|
scatter_idx: int,
|
|
gather_idx: int,
|
|
) -> torch.Tensor:
|
|
ctx.group = group
|
|
ctx.scatter_idx = scatter_idx
|
|
ctx.gather_idx = gather_idx
|
|
res = single_all_to_all(input, scatter_idx, gather_idx, group)
|
|
return res
|
|
|
|
@staticmethod
|
|
def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]:
|
|
return None, _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), None, None
|
|
|
|
|
|
class DistributedAttention(torch.nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
local_attention,
|
|
sequence_parallel,
|
|
scatter_idx: int = 2,
|
|
gather_idx: int = 1,
|
|
) -> None:
|
|
super(DistributedAttention, self).__init__()
|
|
self.local_attn = local_attention
|
|
self.sequence_parallel = sequence_parallel
|
|
self.scatter_idx = scatter_idx
|
|
self.gather_idx = gather_idx
|
|
|
|
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor, *args:
|
|
Any, **kwargs) -> torch.Tensor:
|
|
if self.sequence_parallel.world_size == 1:
|
|
return self.local_attn(query, key, value, attention_mask, *args, **kwargs)
|
|
|
|
# gather ulysses first, ring-attention next
|
|
if self.sequence_parallel.sp_world_size > 1:
|
|
query_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, query, self.scatter_idx, self.gather_idx)
|
|
key_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, key, self.scatter_idx, self.gather_idx)
|
|
value_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, value, self.scatter_idx, self.gather_idx)
|
|
else:
|
|
query_layer, key_layer, value_layer = query, key, value
|
|
|
|
if self.sequence_parallel.rp_world_size > 1:
|
|
# if need ring-attention
|
|
kwargs.pop('position_ids', None)
|
|
# Get the real position ids, this is filled by `sequence_parallel.prepare_inputs`
|
|
# real position ids is different from the position_ids when model uses mrope
|
|
position_ids = self.sequence_parallel.real_position_ids
|
|
# pad and split it by zigzag method
|
|
position_ids = self.sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
|
|
else:
|
|
# only ulysses
|
|
position_ids = kwargs.pop('position_ids')
|
|
if position_ids is not None:
|
|
# Reuse the generic gather path so 2D and 3D position_ids share the same SP behavior.
|
|
position_ids = self.sequence_parallel.gather(position_ids.contiguous(), dim=-1, position_ids=None)
|
|
|
|
context_layer = self.local_attn(
|
|
query_layer, key_layer, value_layer, attention_mask, *args, position_ids=position_ids, **kwargs)
|
|
|
|
if self.sequence_parallel.sp_world_size > 1:
|
|
output = _SeqAllToAll.apply(self.sequence_parallel.sp_group, context_layer, self.gather_idx,
|
|
self.scatter_idx)
|
|
else:
|
|
output = context_layer
|
|
|
|
return output
|