187 lines
7.6 KiB
Python
187 lines
7.6 KiB
Python
from typing import Any, Tuple
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch import Tensor
|
||
from flash_attn import flash_attn_varlen_func
|
||
try:
|
||
import deepspeed.comm as dist
|
||
except:
|
||
dist = None
|
||
|
||
|
||
try:
|
||
from utils import (
|
||
get_sequence_parallel_group,
|
||
get_sequence_parallel_size,
|
||
get_sequence_parallel_rank
|
||
)
|
||
except (ModuleNotFoundError, ImportError):
|
||
# 从 utils 获取seq parallel设置，import不成功默认为不开启
|
||
get_sequence_parallel_group = lambda : None
|
||
get_sequence_parallel_size = lambda : 1
|
||
get_sequence_parallel_rank = lambda : 0
|
||
|
||
|
||
def single_all_to_all(input, scatter_idx, gather_idx, group):
|
||
seq_world_size = dist.get_world_size(group)
|
||
inp_shape = list(input.shape)
|
||
inp_shape[scatter_idx] = inp_shape[scatter_idx] // seq_world_size
|
||
if scatter_idx < 2:
|
||
input_t = input.reshape(
|
||
[seq_world_size, inp_shape[scatter_idx]] + \
|
||
inp_shape[scatter_idx + 1:]
|
||
).contiguous()
|
||
else:
|
||
# transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
|
||
input_t = input.reshape(
|
||
[-1, seq_world_size, inp_shape[scatter_idx]] + \
|
||
inp_shape[scatter_idx + 1:]
|
||
).transpose(0, 1).contiguous()
|
||
|
||
output = torch.empty_like(input_t)
|
||
dist.all_to_all_single(output, input_t, group=group)
|
||
|
||
# if scattering the seq-dim, transpose the heads back to the original dimension
|
||
# [sp_size, seq_len//sp_size, batch_size, head_num // sp_size, head_dim] -->
|
||
# [seq_len//sp_size,batch_size, sp_size, head_num // sp_size, head_dim]
|
||
if scatter_idx < 2:
|
||
output = output.transpose(0, 1).transpose(1, 2).contiguous()
|
||
|
||
return output.reshape(
|
||
inp_shape[: gather_idx] + \
|
||
[inp_shape[gather_idx] * seq_world_size,] + \
|
||
inp_shape[gather_idx + 1:]).contiguous()
|
||
|
||
|
||
class _SeqAllToAll(torch.autograd.Function):
|
||
|
||
@staticmethod
|
||
def forward(ctx: Any, group: 'dist.ProcessGroup', input: Tensor, scatter_idx: int, gather_idx: int) -> Tensor:
|
||
ctx.group = group
|
||
ctx.scatter_idx = scatter_idx
|
||
ctx.gather_idx = gather_idx
|
||
|
||
return single_all_to_all(input, scatter_idx, gather_idx, group)
|
||
|
||
@staticmethod
|
||
def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
|
||
return (None, _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), None, None)
|
||
|
||
|
||
# import from https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/sequence/layer.py
|
||
# but fix some bugs for 符合训练的维度设置
|
||
class DistributedAttention(nn.Module):
|
||
"""Initialization.
|
||
|
||
Arguments:
|
||
local_attention (Module): local attention with q,k,v
|
||
sequence_process_group (ProcessGroup): sequence parallel process group
|
||
scatter_idx (int): scatter_idx for all2all comm
|
||
gather_idx (int): gather_idx for all2all comm
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
local_attention: nn.Module,
|
||
sequence_process_group: 'dist.ProcessGroup',
|
||
scatter_idx: int = 2,
|
||
gather_idx: int = 0,
|
||
) -> None:
|
||
|
||
super(DistributedAttention, self).__init__()
|
||
self.local_attn = local_attention
|
||
self.spg = sequence_process_group
|
||
self.scatter_idx = scatter_idx
|
||
self.gather_idx = gather_idx
|
||
|
||
def pad_attention_head(self, query: Tensor, key: Tensor, value: Tensor):
|
||
# 将输入的head 维度pad到sp_size的倍数
|
||
sp_size = torch.distributed.get_world_size(self.spg)
|
||
pad_size = (sp_size - query.size(1) % sp_size) % sp_size
|
||
if pad_size > 0:
|
||
# [bs, num_head, seq_len, head_dim] -> [bs, num_head+pad_size, seq_len, head_dim]
|
||
query = torch.nn.functional.pad(query, (0,0,0,0,0,pad_size), value = 0.01)
|
||
key = torch.nn.functional.pad(key, (0,0,0,0,0,pad_size), value = 0.01)
|
||
value = torch.nn.functional.pad(value, (0,0,0,0,0,pad_size),value=0.0)
|
||
return query, key, value
|
||
|
||
def forward(self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs) -> Tensor:
|
||
""" forward
|
||
|
||
Arguments:
|
||
query (Tensor): query input to the layer [batch_size, num_head, seq_len, head_dim]
|
||
key (Tensor): key input to the layer
|
||
value (Tensor): value input to the layer
|
||
args: other args
|
||
|
||
Returns:
|
||
* output (Tensor): context output
|
||
"""
|
||
# TODO Merge three alltoall calls into one
|
||
# TODO (Reza): change the api on the megatron-deepspeed side so that we only receive all data (q,k, and v) together!
|
||
# [batch_size,num_head,seq_len, head_dim ]trans to [seq_len,batch_size,num_head,head_dim]
|
||
origin_num_head = query.size(1)
|
||
query, key, value = self.pad_attention_head(query,key,value)
|
||
|
||
query = query.transpose(1,2).transpose(0,1)
|
||
key = key.transpose(1,2).transpose(0,1)
|
||
value = value.transpose(1,2).transpose(0,1)
|
||
#in shape : e.g., [s/p,bs,h,head_dim]
|
||
query_layer = _SeqAllToAll.apply(self.spg, query, self.scatter_idx, self.gather_idx).transpose(0,1).transpose(1,2).contiguous()
|
||
key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx).transpose(0,1).transpose(1,2).contiguous()
|
||
value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx).transpose(0,1).transpose(1,2).contiguous()
|
||
|
||
context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs)
|
||
context_layer = context_layer.transpose(0,1).contiguous()
|
||
# [seq_len, batch_size, num_head, head_dim]
|
||
output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx)
|
||
return output.transpose(0,1)[:,:,:origin_num_head,:]
|
||
|
||
|
||
class LocalAttention(nn.Module):
|
||
def __init__(self, hidden_size, num_heads, head_dim):
|
||
super().__init__()
|
||
self.hidden_size = hidden_size
|
||
self.num_heads = num_heads
|
||
self.head_dim = head_dim
|
||
|
||
def forward(self, q, k, v, *args, use_flash=True, **kwargs):
|
||
# input q,k,v [batch_size, num_head, seq_len, head_dim]
|
||
# output [batch_size, seq_len, num_head, head_dim]
|
||
if use_flash:
|
||
q_len, num_heads = q.shape[2], q.shape[1]
|
||
q = q.transpose(1,2).reshape(-1, num_heads, self.head_dim)
|
||
k = k.transpose(1,2).reshape(-1, num_heads, self.head_dim)
|
||
v = v.transpose(1,2).reshape(-1, num_heads, self.head_dim)
|
||
return flash_attn_varlen_func(q,k,v,*args, **kwargs).reshape(-1,q_len, num_heads, self.head_dim)
|
||
else:
|
||
with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
|
||
attn_output = F.scaled_dot_product_attention(
|
||
q,k,v, *args, **kwargs)
|
||
attn_output = attn_output.transpose(1, 2)
|
||
return attn_output
|
||
|
||
|
||
def create_attention_layer(hidden_size, num_heads, head_dim):
|
||
if get_sequence_parallel_group() is None:
|
||
return LocalAttention(hidden_size, num_heads, head_dim)
|
||
else:
|
||
return DistributedAttention(
|
||
local_attention=LocalAttention(hidden_size, num_heads, head_dim),
|
||
sequence_process_group=get_sequence_parallel_group()
|
||
)
|
||
|
||
|
||
def get_sequence_parallel_chunk(tensor, dim=1, shift=0):
|
||
assert tensor.size(dim) % get_sequence_parallel_size() == 0
|
||
original_size = tensor.size(dim)
|
||
if shift:
|
||
tensor = tensor.split([shift, tensor.size(dim) - shift], dim=dim)[1]
|
||
if get_sequence_parallel_group() is None:
|
||
return tensor
|
||
else:
|
||
chunk_size = original_size // get_sequence_parallel_size()
|
||
return tensor.split(chunk_size, dim=dim)[get_sequence_parallel_rank()]
|