Baichuan-Audio-Base / sequence_parallel_utils.py

Upload folder using huggingface_hub

d6c94dc verified over 1 year ago

7.83 kB

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()]

	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()]