wehub-resource-sync
724 lines
35 KiB
Python
724 lines
35 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import inspect
|
|
import math
|
|
import torch
|
|
import torch.distributed as dist
|
|
from functools import lru_cache, partial
|
|
from torch.distributed import init_device_mesh
|
|
from transformers import PreTrainedTokenizer
|
|
from types import SimpleNamespace
|
|
from typing import Optional
|
|
|
|
from swift.utils import HfConfigFactory, get_cu_seqlens_from_position_ids, get_device, get_dist_setting
|
|
from .ulysses import DistributedAttention
|
|
from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def get_signature_parameters(fn):
|
|
try:
|
|
return inspect.signature(fn).parameters
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def has_signature_parameter(fn, parameter: str) -> bool:
|
|
parameters = get_signature_parameters(fn)
|
|
return parameters is not None and parameter in parameters
|
|
|
|
|
|
def call_with_supported_kwargs(fn, *args, **kwargs):
|
|
parameters = get_signature_parameters(fn)
|
|
if parameters is None:
|
|
return fn(*args, **kwargs)
|
|
if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values()):
|
|
kwargs = {key: value for key, value in kwargs.items() if key in parameters}
|
|
return fn(*args, **kwargs)
|
|
|
|
|
|
def _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args,
|
|
**kwargs):
|
|
if has_signature_parameter(fn, 'cache_position'):
|
|
return call_with_supported_kwargs(
|
|
fn,
|
|
config,
|
|
input_embeds,
|
|
attention_mask,
|
|
cache_position_or_past_key_values,
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
if cache_position_or_past_key_values is None and 'past_key_values' in kwargs:
|
|
return call_with_supported_kwargs(fn, config, input_embeds, attention_mask, *args, **kwargs)
|
|
return call_with_supported_kwargs(
|
|
fn,
|
|
config,
|
|
input_embeds,
|
|
attention_mask,
|
|
cache_position_or_past_key_values,
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
class SequenceParallel:
|
|
|
|
_global_inited: bool = False
|
|
|
|
def __init__(self):
|
|
self.sp_world_size = None
|
|
self.dp_world_size = None
|
|
self.rp_world_size = None
|
|
self.world_size = None
|
|
self.model_dtype = None
|
|
self.tokenizer = None
|
|
self.device_mesh = None
|
|
self.num_heads = None
|
|
self.causal_mask_func = None
|
|
self.extra_kwargs = {}
|
|
|
|
@property
|
|
def real_position_ids(self) -> torch.Tensor:
|
|
"""The real position ids, this is different from the position_ids in mrope"""
|
|
return self.extra_kwargs.get('text_position_ids')
|
|
|
|
def _prepare_flash_attn(self, base_model: torch.nn.Module):
|
|
try:
|
|
from transformers import masking_utils
|
|
|
|
_origin_flash_attention_mask = masking_utils.flash_attention_mask
|
|
|
|
def flash_attention_mask(*args, **kwargs):
|
|
if self.world_size == 1:
|
|
return _origin_flash_attention_mask(*args, **kwargs)
|
|
attention_mask = kwargs.get('attention_mask')
|
|
if attention_mask is not None:
|
|
if attention_mask.all():
|
|
attention_mask = None
|
|
|
|
return attention_mask
|
|
|
|
masking_utils.flash_attention_mask = flash_attention_mask
|
|
masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['flash_attention_2'] = flash_attention_mask
|
|
|
|
def sdpa_mask(*args, **kwargs):
|
|
if self.world_size == 1:
|
|
return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs)
|
|
if 'cache_position' in kwargs:
|
|
device = kwargs['cache_position'].device
|
|
else:
|
|
# transformers>=5.4.0
|
|
device = kwargs['device']
|
|
cache_position = self.real_position_ids[0]
|
|
cache_position = self.pad(cache_position, padding_value=-1, position_ids=self.real_position_ids, dim=0)
|
|
cache_position = torch.arange(0, cache_position.shape[0], device=device)
|
|
kwargs['kv_length'] = cache_position.shape[0]
|
|
if 'cache_position' in kwargs:
|
|
kwargs['cache_position'] = cache_position
|
|
else:
|
|
kwargs['q_length'] = kwargs['kv_length']
|
|
return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs)
|
|
|
|
masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping[
|
|
'sdpa_origin'] = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa']
|
|
masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask
|
|
|
|
def create_causal_mask(config,
|
|
input_embeds,
|
|
attention_mask,
|
|
cache_position_or_past_key_values=None,
|
|
*args,
|
|
**kwargs):
|
|
if self.world_size == 1:
|
|
return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds,
|
|
attention_mask, cache_position_or_past_key_values, *args, **kwargs)
|
|
input_embeds = torch.ones(
|
|
(input_embeds.shape[0], input_embeds.shape[1] * self.sp_world_size, input_embeds.shape[2]),
|
|
dtype=input_embeds.dtype,
|
|
device=input_embeds.device)
|
|
if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'):
|
|
cache_position_or_past_key_values = torch.arange(
|
|
0,
|
|
input_embeds.shape[1],
|
|
device=input_embeds.device,
|
|
)
|
|
return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds,
|
|
attention_mask, cache_position_or_past_key_values, *args, **kwargs)
|
|
|
|
masking_utils.origin_create_causal_mask = masking_utils.create_causal_mask
|
|
masking_utils.create_causal_mask = create_causal_mask
|
|
except ImportError:
|
|
pass
|
|
|
|
if hasattr(base_model, 'language_model'):
|
|
text_model = base_model.language_model
|
|
else:
|
|
text_model = base_model
|
|
|
|
from transformers.modeling_flash_attention_utils import is_flash_attn_available
|
|
if is_flash_attn_available():
|
|
# TODO this works for multi-modal models like qwen2.5-vl
|
|
# SDPA is not supported, because we need to copy the code to our project, which will bring
|
|
# more works for maintaining.
|
|
from transformers import modeling_flash_attention_utils
|
|
from transformers.modeling_flash_attention_utils import _flash_attention_forward
|
|
_distributed_flash_attention = DistributedAttention(_flash_attention_forward, self)
|
|
|
|
modeling_flash_attention_utils._flash_attention_forward_origin = _flash_attention_forward
|
|
|
|
def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor,
|
|
value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len,
|
|
*args, **kwargs):
|
|
if self.world_size == 1:
|
|
return _flash_attention_forward(query_states, key_states, value_states, attention_mask, q_len,
|
|
*args, **kwargs)
|
|
return _distributed_flash_attention(query_states, key_states, value_states, attention_mask,
|
|
q_len * self.sp_world_size, *args, **kwargs)
|
|
|
|
modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward
|
|
|
|
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
|
|
|
|
def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
|
|
dist_attn, **kwargs):
|
|
if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
|
|
return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states,
|
|
value_states, attention_mask, *args,
|
|
**kwargs)
|
|
if dist_attn.local_attn is None:
|
|
|
|
def _attention(query, key, value, *args, **kwargs):
|
|
query = query.transpose(1, 2)
|
|
key = key.transpose(1, 2)
|
|
value = value.transpose(1, 2)
|
|
if self.rp_world_size is not None and self.rp_world_size > 1:
|
|
position_ids = kwargs['position_ids']
|
|
cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
|
|
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
|
|
position_ids = self._split_packed(position_ids, cu_seqlens)
|
|
mask = position_ids != -1
|
|
query = query.transpose(1, 2)
|
|
key = key.transpose(1, 2)
|
|
value = value.transpose(1, 2)
|
|
# this is important
|
|
# the length is correct, but the value is wrong, by mask qkv, we do:
|
|
# mask the padded values of q and v to zero
|
|
# mask the padded values of k to -1e5
|
|
# to keep the attention correct
|
|
query, key, value = self._mask_qkv(query, key, value, mask)
|
|
output = zigzag_ring_flash_attn_varlen_func(
|
|
query,
|
|
key,
|
|
value,
|
|
cu_seqlens=cu_seqlens,
|
|
max_seqlen=max_seqlen,
|
|
causal=module.is_causal,
|
|
dropout_p=kwargs.get('dropout', 0.0),
|
|
softmax_scale=kwargs.get('scaling', 0.0),
|
|
window_size=kwargs.get('sliding_window') or (-1, -1),
|
|
group=self.rp_group)
|
|
return output
|
|
else:
|
|
if 'cu_seq_lens_q' in kwargs:
|
|
position_ids = kwargs.get('position_ids')
|
|
if self.real_position_ids is not None:
|
|
position_ids = self.real_position_ids
|
|
position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
|
|
cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
|
|
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
|
|
assert query.shape[2] == cu_seqlens[-1]
|
|
kwargs['cu_seq_lens_q'] = cu_seqlens
|
|
kwargs['cu_seq_lens_k'] = cu_seqlens
|
|
kwargs['max_length_q'] = max_seqlen
|
|
kwargs['max_length_k'] = max_seqlen
|
|
return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args,
|
|
**kwargs)[0]
|
|
|
|
dist_attn.local_attn = _attention
|
|
|
|
return dist_attn(
|
|
query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
|
|
*args, **kwargs), None
|
|
|
|
def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
|
|
dist_attn, **kwargs):
|
|
# Bypass SP logic when world_size == 1 (SP disabled) or module not in text_model
|
|
if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
|
|
return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states,
|
|
attention_mask, *args, **kwargs)
|
|
if dist_attn.local_attn is None:
|
|
|
|
def _attention(query, key, value, *args, **kwargs):
|
|
query = query.transpose(1, 2)
|
|
key = key.transpose(1, 2)
|
|
value = value.transpose(1, 2)
|
|
if self.rp_world_size > 1:
|
|
raise NotImplementedError('SDPA does not support Ring attention.')
|
|
return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0]
|
|
|
|
dist_attn.local_attn = _attention
|
|
return dist_attn(
|
|
query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
|
|
*args, **kwargs), None
|
|
|
|
ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2']
|
|
ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa']
|
|
ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = partial(
|
|
local_flash_attn, dist_attn=DistributedAttention(None, self))
|
|
ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self))
|
|
|
|
def _prepare_forward_hook(self, base_model: torch.nn.Module):
|
|
|
|
def pre_forward_split_hook(_self, args, kwargs):
|
|
if self.world_size == 1:
|
|
return args, kwargs
|
|
input_ids = kwargs.get('input_ids', None)
|
|
inputs_embeds = kwargs.get('inputs_embeds', None)
|
|
position_ids = kwargs['position_ids']
|
|
attention_mask = kwargs.get('attention_mask', None)
|
|
if hasattr(_self, 'language_model'):
|
|
embed_tokens = getattr(_self.language_model, 'embed_tokens', None)
|
|
else:
|
|
embed_tokens = getattr(_self, 'embed_tokens', None)
|
|
input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs(
|
|
input_ids,
|
|
inputs_embeds,
|
|
None,
|
|
position_ids,
|
|
attention_mask,
|
|
None,
|
|
embed_tokens=embed_tokens,
|
|
real_position_ids=self.real_position_ids)
|
|
kwargs['input_ids'] = input_ids
|
|
kwargs['inputs_embeds'] = inputs_embeds
|
|
kwargs['position_ids'] = position_ids
|
|
kwargs['attention_mask'] = attention_mask
|
|
return args, kwargs
|
|
|
|
base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True)
|
|
|
|
def _prepare_moe_aux_loss(self, base_model: torch.nn.Module):
|
|
from .utils import GatherLoss
|
|
|
|
def moe_aux_loss_hook(module, args, kwargs, output):
|
|
router_logits = getattr(output, 'router_logits', None)
|
|
if router_logits is None:
|
|
return output
|
|
|
|
attention_mask = kwargs['attention_mask']
|
|
if attention_mask is None:
|
|
batch_size = 1
|
|
else:
|
|
batch_size = attention_mask.shape[0]
|
|
|
|
assert router_logits[0].shape[0] % batch_size == 0
|
|
seq_len = router_logits[0].shape[0] // batch_size
|
|
|
|
_gathered_logits = []
|
|
for i in range(batch_size):
|
|
_slice = slice(i * seq_len, (i + 1) * seq_len)
|
|
_bs_logits = [logit[_slice] for logit in router_logits]
|
|
compute_device = _bs_logits[0].device
|
|
_bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0)
|
|
_bs_logits, _ = GatherLoss.apply(_bs_logits, None, 1, self.real_position_ids)
|
|
_gathered_logits.append(_bs_logits)
|
|
router_logits = torch.stack(_gathered_logits, dim=0)
|
|
if self.real_position_ids is not None:
|
|
router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :]
|
|
output['router_logits'] = tuple(
|
|
[logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)])
|
|
return output
|
|
|
|
base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True)
|
|
|
|
def prepare(self, sp_size: int, model: torch.nn.Module, tokenizer: PreTrainedTokenizer, padding_free: bool):
|
|
from swift.model import get_llm_model
|
|
self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_key_value_heads')
|
|
if self.num_heads is None:
|
|
self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_attention_heads')
|
|
assert self.num_heads is not None, 'Cannot find num_heads config in config.json'
|
|
self.padding_free = padding_free
|
|
self.world_size = sp_size
|
|
|
|
llm_model = get_llm_model(model)
|
|
|
|
if hasattr(llm_model, 'language_model'):
|
|
if hasattr(llm_model.language_model, '_update_causal_mask'):
|
|
self.causal_mask_func = llm_model.language_model._update_causal_mask
|
|
else:
|
|
if hasattr(llm_model, '_update_causal_mask'):
|
|
self.causal_mask_func = llm_model._update_causal_mask
|
|
|
|
if not SequenceParallel._global_inited:
|
|
# these operations are global initializations and patches
|
|
self._init_device_mesh()
|
|
self._prepare_flash_attn(llm_model)
|
|
SequenceParallel._global_inited = True
|
|
|
|
self._prepare_forward_hook(llm_model)
|
|
|
|
if model.model_info.is_moe_model:
|
|
self._prepare_moe_aux_loss(llm_model)
|
|
|
|
self.model_dtype = next(model.parameters()).dtype
|
|
self.tokenizer = tokenizer
|
|
if self.rp_world_size > 1 and not self.padding_free:
|
|
raise NotImplementedError(
|
|
f'The world_size {self.world_size} needs ulysses/ring-attention, which needs --padding_free true')
|
|
|
|
def _mask_qkv(self, query, key, value, mask):
|
|
mask = mask.unsqueeze(2).unsqueeze(3)
|
|
query = query * mask
|
|
value = value * mask
|
|
mask = (~mask) * -1e5 # for bf16
|
|
key = key + mask.to(key.dtype)
|
|
return query, key, value
|
|
|
|
def pad(self, tensor, padding_value, position_ids=None, dim=1):
|
|
"""Pad tensor for sequence parallel"""
|
|
if self.rp_world_size > 1:
|
|
world_size = self.world_size * 2
|
|
else:
|
|
world_size = self.world_size
|
|
|
|
def _do_pad(tensor):
|
|
length = tensor.shape[dim]
|
|
pad_num = world_size - (length % world_size)
|
|
if pad_num == 0 or pad_num == world_size:
|
|
return tensor
|
|
if not isinstance(padding_value, torch.Tensor):
|
|
# ids
|
|
pad_shape = ((*tensor.shape[:dim], pad_num, *tensor.shape[dim + 1:]) if dim != -1 else
|
|
(*tensor.shape[:dim], pad_num))
|
|
pad = torch.full(pad_shape, padding_value, dtype=tensor.dtype, device=tensor.device)
|
|
tensor = torch.cat([tensor, pad], dim=dim)
|
|
else:
|
|
# For embeddings
|
|
tensor = torch.cat([tensor, padding_value.unsqueeze(0).repeat(tensor.shape[0], pad_num, 1)], dim=dim)
|
|
return tensor
|
|
|
|
if position_ids is not None and self.rp_world_size > 1:
|
|
cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
|
|
all_tensors = []
|
|
for i in range(len(cu_seqlens) - 1):
|
|
if dim == 1:
|
|
sub_tensor = tensor[:, cu_seqlens[i]:cu_seqlens[i + 1]]
|
|
elif dim == -1:
|
|
sub_tensor = tensor[..., cu_seqlens[i]:cu_seqlens[i + 1]]
|
|
else:
|
|
raise NotImplementedError()
|
|
all_tensors.append(_do_pad(sub_tensor))
|
|
tensor = torch.cat(all_tensors, dim=dim)
|
|
|
|
return _do_pad(tensor)
|
|
|
|
def gather(self, local_output, dim: int, position_ids=None):
|
|
"""Gather tensor for sequence parallel - reverse of split"""
|
|
if self.world_size == 1:
|
|
return local_output
|
|
|
|
if self.rp_world_size > 1:
|
|
input_dim = local_output.dim()
|
|
assert input_dim >= 2
|
|
|
|
if position_ids is not None:
|
|
position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
|
|
|
|
# Step 1: Gather from all sequence parallel ranks
|
|
# Each sp_rank has its own piece, we need to gather them first
|
|
gathered_sp = [torch.zeros_like(local_output) for _ in range(self.sp_world_size)]
|
|
torch.distributed.all_gather(gathered_sp, local_output.contiguous(), group=self.sp_group)
|
|
|
|
# Concatenate the sp pieces to form the complete chunk for this rp_rank
|
|
rp_chunk = torch.cat(gathered_sp, dim=dim)
|
|
|
|
# Step 2: Gather all rp chunks
|
|
gathered_rp = [torch.zeros_like(rp_chunk) for _ in range(self.rp_world_size)]
|
|
torch.distributed.all_gather(gathered_rp, rp_chunk, group=self.rp_group)
|
|
|
|
cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
|
|
all_tensor_length = []
|
|
for i in range(len(cu_seqlens) - 1):
|
|
length = cu_seqlens[i + 1] - cu_seqlens[i]
|
|
padding_length = math.ceil(length / (self.world_size * 2)) * (self.world_size * 2)
|
|
all_tensor_length.append(padding_length)
|
|
|
|
full_output = torch.zeros(
|
|
[local_output.shape[0], sum(all_tensor_length), *local_output.shape[2:]], device=local_output.device)
|
|
for idx_rp, rp_tensor in enumerate(gathered_rp): # rp world size
|
|
# re-group the zigzag to the correct order
|
|
accumulated_length = 0
|
|
for idx_seq, length in enumerate(all_tensor_length): # sequence number
|
|
local_length = length // self.rp_world_size
|
|
local_tensor = rp_tensor[:, accumulated_length:local_length + accumulated_length]
|
|
chunk_size = local_length // 2
|
|
left_idx = accumulated_length * self.rp_world_size + idx_rp * chunk_size
|
|
right_idx = accumulated_length * self.rp_world_size + (idx_rp + 1) * chunk_size
|
|
full_output[:, left_idx:right_idx] = local_tensor[:, :chunk_size]
|
|
left_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp
|
|
- 1) * chunk_size
|
|
right_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp) * chunk_size
|
|
full_output[:, left_idx:right_idx] = local_tensor[:, chunk_size:]
|
|
accumulated_length += local_length
|
|
|
|
return full_output.contiguous()
|
|
else:
|
|
gathered_sp = torch.empty(
|
|
[local_output.shape[0] * self.sp_world_size] + list(local_output.shape[1:]),
|
|
dtype=local_output.dtype,
|
|
device=local_output.device)
|
|
dist.all_gather_into_tensor(gathered_sp, local_output, group=self.sp_group)
|
|
gathered_sp = torch.cat(gathered_sp.split(local_output.shape[0], dim=0), dim=dim)
|
|
return gathered_sp.contiguous()
|
|
|
|
def _split_packed(self, value, cu_seqlens, dim=1):
|
|
"""Split and re-group in zigzag"""
|
|
local_values = []
|
|
for i in range(len(cu_seqlens) - 1):
|
|
start, end = cu_seqlens[i], cu_seqlens[i + 1]
|
|
if dim == 1:
|
|
sub_value = value[:, start:end]
|
|
elif dim == -1:
|
|
sub_value = value[..., start:end]
|
|
else:
|
|
raise NotImplementedError()
|
|
local_value = sub_value.chunk(2 * self.rp_world_size, dim=dim)
|
|
local_values.extend([
|
|
local_value[self.rp_rank],
|
|
local_value[2 * self.rp_world_size - 1 - self.rp_rank],
|
|
])
|
|
return torch.cat(local_values, dim=dim).contiguous()
|
|
|
|
def split(self, input, dim: int, position_ids=None):
|
|
"""Split tensor for sequence parallel"""
|
|
if self.world_size == 1:
|
|
return input
|
|
|
|
if self.rp_world_size > 1:
|
|
input_dim = input.dim()
|
|
assert input_dim >= 2
|
|
cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
|
|
assert torch.all(cu_seqlens % (2 * self.rp_world_size) == 0)
|
|
value_chunks = self._split_packed(input, cu_seqlens, dim=dim)
|
|
local_value = value_chunks.chunk(self.sp_world_size, dim=dim)[self.sp_rank].contiguous()
|
|
return local_value
|
|
else:
|
|
rank = self.sp_rank
|
|
dim_size = input.size(dim)
|
|
assert dim_size % self.sp_world_size == 0, (
|
|
f'The dimension to split ({dim_size}) is not a multiple of '
|
|
f'world size ({self.sp_world_size}), cannot split tensor evenly')
|
|
|
|
tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim)
|
|
output = tensor_list[rank].contiguous()
|
|
return output
|
|
|
|
def pad_and_split_mm_tokens(self, visual_mask, mm_embeds):
|
|
input_ids = self.extra_kwargs['input_ids']
|
|
empty_embeds = torch.empty(
|
|
(input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype)
|
|
empty_embeds[visual_mask] = mm_embeds
|
|
|
|
embeds = SimpleNamespace(weight=mm_embeds)
|
|
|
|
_, split_input_embeds, _, _, _, _, extra_values = self.pad_and_split_inputs(
|
|
None,
|
|
empty_embeds,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
embeds,
|
|
self.real_position_ids,
|
|
extra_split_values=[(visual_mask, 0, -1)])
|
|
visual_mask = extra_values[0]
|
|
return visual_mask, split_input_embeds[visual_mask]
|
|
|
|
def pad_and_split_inputs(self,
|
|
input_ids,
|
|
input_embeds,
|
|
labels,
|
|
position_ids,
|
|
attention_mask,
|
|
loss_scale,
|
|
embed_tokens=None,
|
|
real_position_ids=None,
|
|
extra_split_values=None):
|
|
"""Common implementation for padding and splitting inputs
|
|
|
|
When a sequence comes, it will be split into rp_world_size * 2 sub tensors, and group them as the
|
|
zigzag order. So we get rp_world_size tensors, then we split each tensor to sp_world_size ones.
|
|
So, we should first pad the original sequence to the length can be divided by 2 * world_size, then re-group it.
|
|
|
|
Only support padding_free for ring-attention, because non-padding_free mode needs another pad/split workflow
|
|
man, that's a lot of work...
|
|
|
|
Args:
|
|
input_ids: input_ids
|
|
input_embeds: input_embeds
|
|
labels: labels
|
|
position_ids: position_ids or, position_ids for mrope
|
|
attention_mask: attention_mask
|
|
loss_scale: loss_scale
|
|
embed_tokens: embed_tokens
|
|
real_position_ids: the real position_ids to represent the seq length information
|
|
extra_split_values: List of Tuples for extra split values, e.g.: (tensor, pad_value, split_dim)
|
|
"""
|
|
tokenizer = self.tokenizer
|
|
real_position_ids = real_position_ids if real_position_ids is not None else position_ids
|
|
extra_values = []
|
|
batch_size = input_ids.shape[
|
|
0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else None
|
|
if real_position_ids is not None and batch_size is not None and real_position_ids.shape[0] == batch_size:
|
|
# TODO clone everytime, but the position_ids is a small tensor
|
|
self.extra_kwargs['text_position_ids'] = real_position_ids.clone()
|
|
if input_ids is not None:
|
|
input_ids = self.pad(input_ids, padding_value=tokenizer.pad_token_id, position_ids=real_position_ids)
|
|
self.extra_kwargs['input_ids'] = input_ids.clone()
|
|
if input_embeds is not None:
|
|
pad_emb = torch.zeros(
|
|
(1, embed_tokens.weight.shape[-1])).to(embed_tokens.weight.device).to(embed_tokens.weight.dtype)
|
|
input_embeds = self.pad(input_embeds, padding_value=pad_emb, position_ids=real_position_ids)
|
|
batch_size = input_ids.shape[
|
|
0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else 1
|
|
if position_ids is not None:
|
|
position_ids = self.pad(position_ids, padding_value=-1, position_ids=real_position_ids, dim=-1)
|
|
if labels is not None:
|
|
labels = self.pad(labels, padding_value=-100, position_ids=real_position_ids)
|
|
if loss_scale is not None:
|
|
loss_scale = self.pad(loss_scale, padding_value=0., position_ids=real_position_ids)
|
|
if real_position_ids is not None:
|
|
real_position_ids = self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids)
|
|
if (input_ids is not None or input_embeds is not None) and batch_size > 1:
|
|
# not padding_free, so not ring-attention
|
|
inputs = input_ids if input_ids is not None else input_embeds
|
|
attn_shape = inputs.shape[1] # The sequence length
|
|
if attention_mask is None:
|
|
attention_mask = torch.ones_like(real_position_ids)
|
|
# no need position_ids here, because padding_free does not need attention_mask,
|
|
# so this is not ring-attention
|
|
attention_mask = self.pad(attention_mask, padding_value=0)
|
|
cache_position = torch.arange(0, attn_shape, device=inputs.device)
|
|
# pad attention mask to 4d to avoid calculation errors
|
|
if hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None:
|
|
attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), cache_position,
|
|
None, None)
|
|
if extra_split_values is not None:
|
|
for (tensor, pad_value, split_dim) in extra_split_values:
|
|
extra_values.append(
|
|
self.pad(tensor, padding_value=pad_value, position_ids=real_position_ids, dim=split_dim))
|
|
if input_ids is not None:
|
|
input_ids = self.split(input_ids, dim=1, position_ids=real_position_ids)
|
|
if input_embeds is not None:
|
|
input_embeds = self.split(input_embeds, dim=1, position_ids=real_position_ids)
|
|
if labels is not None:
|
|
labels = torch.roll(labels, shifts=-1, dims=-1)
|
|
labels = self.split(labels, dim=-1, position_ids=real_position_ids)
|
|
if loss_scale is not None:
|
|
loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1)
|
|
loss_scale = self.split(loss_scale, dim=-1, position_ids=real_position_ids)
|
|
|
|
if position_ids is not None:
|
|
position_ids = self.split(position_ids, dim=-1, position_ids=real_position_ids)
|
|
if extra_split_values is not None:
|
|
for i in range(len(extra_values)):
|
|
extra_values[i] = self.split(
|
|
extra_values[i], dim=extra_split_values[i][2], position_ids=real_position_ids)
|
|
return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale, extra_values
|
|
|
|
def _gather_object_dp(self, input_data):
|
|
"""Gather object for data parallel"""
|
|
input_data_list = [None] * self.dp_world_size
|
|
dist.all_gather_object(input_data_list, input_data, group=self.dp_group)
|
|
return [x for y in input_data_list for x in y]
|
|
|
|
def _init_device_mesh(self):
|
|
"""Initialize device mesh for sequence and ring parallel.
|
|
|
|
The logic is unified:
|
|
1. Determine the Sequence Parallel (SP) size first based on GCD to satisfy constraints.
|
|
2. Allocate all remaining model parallelism to Ring Parallel (RP).
|
|
"""
|
|
_, _, world_size, _ = get_dist_setting()
|
|
self.dp_world_size = world_size // self.world_size
|
|
# SP size is the GCD of num_heads and world_size, guaranteeing it divides both.
|
|
self.sp_world_size = math.gcd(self.num_heads, self.world_size)
|
|
# RP takes the remaining factor so all model-parallel GPUs are used.
|
|
self.rp_world_size = self.world_size // self.sp_world_size
|
|
|
|
if self.rp_world_size > 1:
|
|
mesh_shape = (self.dp_world_size, self.rp_world_size, self.sp_world_size)
|
|
mesh_dim_names = ('data', 'ring', 'sequence')
|
|
else:
|
|
mesh_shape = (self.dp_world_size, self.sp_world_size)
|
|
mesh_dim_names = ('data', 'sequence')
|
|
self.device_mesh = init_device_mesh(
|
|
get_device().split(':')[0], mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names)
|
|
|
|
def _dim_group(self, dim_name: str):
|
|
"""Return the process group of the given mesh dim, or None if absent."""
|
|
if not self.device_mesh or dim_name not in self.device_mesh.mesh_dim_names:
|
|
return None
|
|
return self.device_mesh[dim_name].get_group()
|
|
|
|
def _dim_rank(self, dim_name: str, default: int) -> int:
|
|
"""Return the rank within the given mesh dim, or `default` if absent."""
|
|
group = self._dim_group(dim_name)
|
|
return dist.get_rank(group) if group is not None else default
|
|
|
|
@property
|
|
def sp_group(self):
|
|
"""Return the sequence parallel group"""
|
|
return self._dim_group('sequence')
|
|
|
|
def enabled(self):
|
|
return self.world_size is not None and self.world_size > 1
|
|
|
|
@property
|
|
def sp_rank(self):
|
|
"""Return the sequence parallel rank"""
|
|
return self._dim_rank('sequence', default=0)
|
|
|
|
@property
|
|
def dp_group(self):
|
|
"""Return the data parallel group"""
|
|
return self._dim_group('data')
|
|
|
|
@property
|
|
def dp_rank(self):
|
|
"""Return the data parallel rank"""
|
|
return self._dim_rank('data', default=0)
|
|
|
|
@property
|
|
def rp_group(self):
|
|
"""Return the ring parallel group"""
|
|
return self._dim_group('ring')
|
|
|
|
@property
|
|
def rp_rank(self):
|
|
"""Return the ring parallel rank"""
|
|
return self._dim_rank('ring', default=-1)
|
|
|
|
def prepare_inputs(self, inputs):
|
|
"""Prepare inputs
|
|
|
|
1. set extra_kwargs['text_position_ids']
|
|
2. split labels
|
|
"""
|
|
position_ids = inputs.get('text_position_ids')
|
|
input_ids = inputs.get('input_ids')
|
|
if position_ids is None:
|
|
position_ids = inputs.get('position_ids')
|
|
if position_ids is not None and input_ids is not None and position_ids.shape[0] == input_ids.shape[0]:
|
|
self.extra_kwargs['text_position_ids'] = position_ids.clone()
|
|
if input_ids is not None:
|
|
self.extra_kwargs['input_ids'] = input_ids.clone()
|
|
if 'labels' in inputs:
|
|
labels = inputs['labels']
|
|
_, _, labels, _, _, _, _ = self.pad_and_split_inputs(
|
|
None, None, labels, None, None, None, real_position_ids=position_ids)
|
|
inputs['labels'] = labels
|
|
|
|
|
|
sequence_parallel = SequenceParallel()
|