mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-08 20:40:34 +00:00
[Feature] Zigzag Ring attention (#5905)
* halfway * fix cross-PP-stage position id length diff bug * fix typo * fix typo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * unified cross entropy func for all shardformer models * remove redundant lines * add basic ring attn; debug cross entropy * fwd bwd logic complete * fwd bwd logic complete; add experimental triton rescale * precision tests passed * precision tests passed * fix typos and remove misc files * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add sp_mode to benchmark; fix varlen interface * update softmax_lse shape by new interface * change tester name * remove buffer clone; support packed seq layout * add varlen tests * fix typo * all tests passed * add dkv_group; fix mask * remove debug statements --------- Co-authored-by: Edenzzzz <wtan45@wisc.edu> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Edenzzzz
GitHub
@@ -4,10 +4,15 @@ from torch.autograd import Function
|
||||
from torch.distributed import ProcessGroup
|
||||
from torch.nn import CrossEntropyLoss
|
||||
|
||||
from colossalai.shardformer.layer._operation import reduce_forward
|
||||
from colossalai.shardformer.shard import ShardConfig
|
||||
|
||||
from .utils import is_share_sp_tp
|
||||
|
||||
__all__ = ["DistCrossEntropy", "cross_entropy_1d", "dist_cross_entropy"]
|
||||
|
||||
_IGNORE_IDX = -100
|
||||
|
||||
|
||||
class DistCrossEntropy(Function):
|
||||
r"""
|
||||
@@ -26,11 +31,12 @@ class DistCrossEntropy(Function):
|
||||
process_group: ProcessGroup,
|
||||
vocab_size: int,
|
||||
dtype=torch.float32,
|
||||
mode="mean",
|
||||
):
|
||||
r"""
|
||||
Calculate the cross entropy loss before gather, the origin loss function is as follows:
|
||||
loss = -log(exp(x[class])/sum(exp(x[i]))
|
||||
and can be rewrite as:
|
||||
and can be rewriten as:
|
||||
loss = log(sum(exp(x[i])) - x[class]
|
||||
|
||||
To avoid the `nan` of log(sum(exp(x[i]))), we minus the max of x[i]
|
||||
@@ -44,12 +50,10 @@ class DistCrossEntropy(Function):
|
||||
Returns:
|
||||
:class:`torch.Tensor`: The cross entropy loss
|
||||
"""
|
||||
assert mode in ["mean", "sum"]
|
||||
# get the max
|
||||
logits_max = torch.max(vocab_logits, dim=-1)[0]
|
||||
dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=process_group)
|
||||
|
||||
# minus the max to avoid the result of sum of exp is too large and the log is nan
|
||||
vocab_logits = vocab_logits - logits_max.unsqueeze(dim=-1)
|
||||
handle = dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=process_group, async_op=True)
|
||||
|
||||
# mask the target in the local device
|
||||
rank = dist.get_rank(group=process_group)
|
||||
@@ -70,24 +74,25 @@ class DistCrossEntropy(Function):
|
||||
mask = (target < down_threshold) | (target >= up_threshold)
|
||||
masked_target = target.clone() - down_threshold
|
||||
masked_target[mask] = 0
|
||||
masked_target_1d = masked_target.view(-1).contiguous()
|
||||
|
||||
# minus the max to avoid the result of sum of exp is too large and the log is nan
|
||||
handle.wait()
|
||||
vocab_logits = vocab_logits - logits_max.unsqueeze(dim=-1)
|
||||
# reshape the logits and target
|
||||
# reshape the vocab_logits to [bath_size * seq_len, vocab_size]
|
||||
# reshape the labels to [bath_size * seq_len]
|
||||
self_vocab_size = vocab_logits.size()[-1]
|
||||
logits_2d = vocab_logits.view(-1, self_vocab_size)
|
||||
masked_target_1d = masked_target.view(-1)
|
||||
|
||||
# extract the x[class] and set the x[other device] to zero
|
||||
pred_logits_1d = logits_2d[
|
||||
torch.arange(start=0, end=logits_2d.shape[0], device=logits_2d.device), masked_target_1d
|
||||
]
|
||||
pred_logits_1d = pred_logits_1d.clone().contiguous()
|
||||
idx = torch.arange(start=0, end=logits_2d.shape[0], device=logits_2d.device)
|
||||
pred_logits_1d = logits_2d[idx, masked_target_1d].contiguous()
|
||||
pred_logits = pred_logits_1d.view_as(target)
|
||||
pred_logits[mask] = 0.0
|
||||
|
||||
# allreduce the get all x(i,y)
|
||||
dist.all_reduce(pred_logits, op=dist.ReduceOp.SUM, group=process_group)
|
||||
# all-reduce to get full x[i, y]
|
||||
handle = dist.all_reduce(pred_logits, op=dist.ReduceOp.SUM, group=process_group, async_op=True)
|
||||
exp_logits = vocab_logits
|
||||
torch.exp(vocab_logits, out=exp_logits)
|
||||
sum_exp_logits = torch.sum(exp_logits, dim=-1, dtype=torch.float32)
|
||||
@@ -95,23 +100,29 @@ class DistCrossEntropy(Function):
|
||||
|
||||
# calculate the loss
|
||||
# loss = log(sum(exp(x[i]))) - x[class]
|
||||
handle.wait()
|
||||
loss = torch.where(target == ignore_index, 0.0, torch.log(sum_exp_logits) - pred_logits)
|
||||
num_non_zero = torch.sum(loss != 0.0)
|
||||
ctx.inv_num_non_zero = 1.0 / num_non_zero
|
||||
loss = torch.sum(loss).div_(num_non_zero)
|
||||
if mode == "mean":
|
||||
num_non_zero = torch.sum(loss != 0.0)
|
||||
ctx.inv_num_non_zero = 1.0 / num_non_zero
|
||||
loss = torch.sum(loss).div_(num_non_zero)
|
||||
else:
|
||||
loss = torch.sum(loss)
|
||||
|
||||
# calculate the softmax
|
||||
exp_logits = exp_logits.div(sum_exp_logits.unsqueeze(dim=-1)).to(dtype)
|
||||
exp_logits[target == ignore_index] = 0.0
|
||||
ctx.save_for_backward(exp_logits, mask, masked_target_1d)
|
||||
ctx.dtype = dtype
|
||||
ctx.mode = mode
|
||||
|
||||
return loss
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
# retrieve the saved tensors
|
||||
grad_output = grad_output * ctx.inv_num_non_zero
|
||||
if ctx.mode == "mean":
|
||||
grad_output = grad_output * ctx.inv_num_non_zero
|
||||
exp_logits, mask, masked_target_1d = ctx.saved_tensors
|
||||
|
||||
# use exp logits as the input grad
|
||||
@@ -123,55 +134,113 @@ class DistCrossEntropy(Function):
|
||||
grad_logits_2d[torch.arange(0, grad_logits_2d.shape[0]), masked_target_1d] -= update
|
||||
|
||||
grad_logits.mul_(grad_output.unsqueeze(dim=-1))
|
||||
return grad_logits, None, None, None, None, None
|
||||
return grad_logits, None, None, None, None, None, None
|
||||
|
||||
|
||||
def cross_entropy_1d(
|
||||
vocab_logits: torch.Tensor,
|
||||
labels: torch.Tensor,
|
||||
ignore_index: int = -100,
|
||||
ignore_index: int = _IGNORE_IDX,
|
||||
process_group: ProcessGroup = None,
|
||||
vocab_size: int = None,
|
||||
dtype: torch.dtype = None,
|
||||
mode: str = "mean",
|
||||
) -> torch.Tensor:
|
||||
return DistCrossEntropy.apply(vocab_logits, labels, ignore_index, process_group, vocab_size, dtype)
|
||||
return DistCrossEntropy.apply(vocab_logits, labels, ignore_index, process_group, vocab_size, dtype, mode)
|
||||
|
||||
|
||||
def dist_cross_entropy(
|
||||
labels: torch.Tensor,
|
||||
logits: torch.Tensor,
|
||||
labels: torch.Tensor, # [B, S] or [B, S, Vocab_size]
|
||||
logits: torch.Tensor, # [B, S, Vocab_size]
|
||||
shard_config: ShardConfig,
|
||||
out_features: int,
|
||||
vocab_size: int,
|
||||
dtype: torch.dtype,
|
||||
seq_dim: int = 1,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Helper to compute cross entropy loss for most shardformer models,
|
||||
compatible with PP, TP and SP.
|
||||
Helper to compute cross entropy loss for most shardformer models supporting PP, TP and SP.
|
||||
"""
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = CrossEntropyLoss()
|
||||
shift_labels = shift_labels.view(-1)
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
|
||||
# Cross entropy with all-reduce for TP
|
||||
new_vocab_size = logits.shape[-1]
|
||||
shift_logits = shift_logits.view(-1, new_vocab_size)
|
||||
loss = cross_entropy_1d(
|
||||
shift_logits,
|
||||
shift_labels,
|
||||
process_group=shard_config.tensor_parallel_process_group,
|
||||
vocab_size=out_features,
|
||||
dtype=dtype,
|
||||
)
|
||||
else:
|
||||
# NOTE if use TP and not parallel_output, the output is gathered.
|
||||
# see VocabParallelLMHead1D
|
||||
shift_logits = shift_logits.view(-1, vocab_size)
|
||||
loss = loss_fct(shift_logits, shift_labels)
|
||||
# Split labels if not gather output
|
||||
sp_group = shard_config.sequence_parallel_process_group
|
||||
sp_rank = dist.get_rank(sp_group)
|
||||
sp_size = shard_config.sequence_parallel_size
|
||||
sp_mode = shard_config.sequence_parallelism_mode
|
||||
parallel_output = shard_config.parallel_output
|
||||
is_tp = shard_config.enable_tensor_parallelism
|
||||
is_packed = labels.dim() == 2
|
||||
if is_packed:
|
||||
bs, seq_len = labels.shape
|
||||
else:
|
||||
# padded sequence
|
||||
seq_len = labels.shape[-1]
|
||||
logits = logits.reshape(-1, *logits.shape[2:])
|
||||
seq_dim = 0
|
||||
|
||||
return loss
|
||||
# Shift labels to predict the next token, and remove the tail logit predicting <EOS>
|
||||
is_sp = sp_size > 1 and (not is_share_sp_tp(sp_mode))
|
||||
split_labels_here = seq_len // sp_size == logits.size(seq_dim) # ring attn splits labels before forward
|
||||
|
||||
if sp_mode == "ring_attn":
|
||||
# For Zigzag Ring Attention, labels should've been split and
|
||||
# shifted by RingAttention.prepare_varlen_batch()
|
||||
if sp_rank == 0:
|
||||
logits = logits[..., :-1, :]
|
||||
logits = torch.cat([logits, torch.full_like(logits[:, :1, :], _IGNORE_IDX)], dim=seq_dim)
|
||||
elif is_sp:
|
||||
# Shift only once: either before splitting or in the last rank without splitting
|
||||
if split_labels_here or (sp_rank == sp_size - 1):
|
||||
labels = labels[..., 1:]
|
||||
if split_labels_here:
|
||||
labels = labels.split(seq_len // sp_size, dim=-1)[sp_rank]
|
||||
|
||||
if sp_rank == sp_size - 1:
|
||||
logits = logits[..., :-1, :]
|
||||
# Pad logits and labels to the same shape across all ranks for TP all_reduce
|
||||
if is_tp and parallel_output:
|
||||
# If is packed sequence (label dim is 1), then each seq already has the end label token padded.
|
||||
# torch.cat is faster than F.pad...
|
||||
pad_shape = (logits.shape[0], 1, *logits.shape[2:]) if is_packed else (1, *logits.shape[1:])
|
||||
padding = torch.full(pad_shape, _IGNORE_IDX, dtype=logits.dtype, device=logits.device)
|
||||
logits = torch.cat([logits, padding], dim=seq_dim)
|
||||
pad_shape = (labels.shape[0], 1) if is_packed else (1,)
|
||||
padding = torch.full(pad_shape, _IGNORE_IDX, dtype=labels.dtype, device=labels.device)
|
||||
labels = torch.cat([labels, padding], dim=seq_dim)
|
||||
else:
|
||||
labels = labels[..., 1:]
|
||||
logits = logits[..., :-1, :]
|
||||
labels = labels.contiguous()
|
||||
logits = logits.contiguous()
|
||||
num_nonzero = (labels != _IGNORE_IDX).sum()
|
||||
assert labels.shape == logits.shape[:-1], f"label shape {labels.shape} does not match logit shape {logits.shape}"
|
||||
|
||||
# Flatten the tokens
|
||||
loss_fct = CrossEntropyLoss(ignore_index=_IGNORE_IDX, reduction="sum")
|
||||
labels = labels.view(-1)
|
||||
|
||||
if is_tp and parallel_output:
|
||||
# Cross entropy with all-reduce for TP
|
||||
new_vocab_size = logits.shape[-1]
|
||||
logits = logits.view(-1, new_vocab_size)
|
||||
loss = cross_entropy_1d(
|
||||
logits,
|
||||
labels,
|
||||
process_group=shard_config.tensor_parallel_process_group,
|
||||
vocab_size=out_features,
|
||||
dtype=dtype,
|
||||
mode="sum",
|
||||
)
|
||||
else:
|
||||
# NOTE if use TP and not parallel_output, the output is gathered in VocabParallelLMHead1D
|
||||
logits = logits.view(-1, vocab_size)
|
||||
loss = loss_fct(logits, labels)
|
||||
|
||||
# Reduce loss instead of gathering logits over seq dim for savings
|
||||
if split_labels_here or sp_mode == "ring_attn":
|
||||
# Get the global non-zero count
|
||||
loss = torch.stack((loss, num_nonzero))
|
||||
# Rescale to offset the grad / (DP * SP) in HybridParallelPlugin
|
||||
loss = reduce_forward(loss, sp_group, grad_scale=sp_size)
|
||||
loss, num_nonzero = loss[0], loss[1].detach()
|
||||
loss = (loss / num_nonzero).squeeze()
|
||||
return loss
|
||||
|
||||
Reference in New Issue
Block a user