diff --git a/colossalai/shardformer/layer/_operation.py b/colossalai/shardformer/layer/_operation.py index 0252f90e1c27..9aa1a44fdfeb 100644 --- a/colossalai/shardformer/layer/_operation.py +++ b/colossalai/shardformer/layer/_operation.py @@ -993,18 +993,33 @@ class _SplitForwardGatherBackward(torch.autograd.Function): @staticmethod def forward(ctx, input_, dim, process_group, grad_scale=None, fp8_communication=False): ctx.process_group = process_group - ctx.dim = dim + # Keep the unpadded extent so the gradient can be trimmed after the + # equal-sized collective. Sequence-parallel inputs are not always + # divisible by the group size (for example, a final packed batch). + # Normalising the dimension here also makes ``narrow`` work with + # negative dimensions in the backward pass. + ctx.dim = dim if dim >= 0 else input_.dim() + dim + ctx.input_dim_size = input_.size(ctx.dim) ctx.grad_scale = grad_scale ctx.fp8_communication = fp8_communication - return _split(input_, dim, process_group) + return _split(input_, ctx.dim, process_group, pad_to_world_size=True) @staticmethod def backward(ctx, grad_output): if ctx.grad_scale is not None: grad_output = grad_output * ctx.grad_scale + grad_input = _gather(grad_output, ctx.dim, ctx.process_group, ctx.fp8_communication, fp8_format="e5m2") + # The collective operates on the padded extent, but the input to this + # autograd function did not contain those synthetic tokens. Trim + # before returning so padding can never contribute to a valid input + # gradient (even when a downstream kernel produced a non-zero padded + # gradient). + if grad_input.size(ctx.dim) != ctx.input_dim_size: + grad_input = grad_input.narrow(ctx.dim, 0, ctx.input_dim_size).contiguous() + return ( - _gather(grad_output, ctx.dim, ctx.process_group, ctx.fp8_communication, fp8_format="e5m2"), + grad_input, None, None, None, @@ -1065,18 +1080,50 @@ class _GatherForwardSplitBackward(torch.autograd.Function): """ @staticmethod - def forward(ctx, input_, dim, process_group, grad_scale=None, fp8_communication=False): + def forward(ctx, input_, dim, process_group, grad_scale=None, fp8_communication=False, output_dim_size=None): ctx.process_group = process_group ctx.dim = dim ctx.grad_scale = grad_scale - - return _gather(input_, dim, process_group, fp8_communication=fp8_communication, fp8_format="e4m3") + ctx.output_dim_size = output_dim_size + ctx.input_dim_size = input_.size(dim) + + output = _gather(input_, dim, process_group, fp8_communication=fp8_communication, fp8_format="e4m3") + if output_dim_size is not None: + if output_dim_size < 0: + output_dim_size += output.size(dim) + assert output_dim_size <= output.size(dim), ( + f"The requested output extent ({output_dim_size}) cannot exceed the gathered extent " + f"({output.size(dim)})" + ) + if output_dim_size != output.size(dim): + output = output.narrow(dim, 0, output_dim_size).contiguous() + return output @staticmethod def backward(ctx, grad_output): if ctx.grad_scale is not None: grad_output = grad_output * ctx.grad_scale - return _split(grad_output, ctx.dim, ctx.process_group), None, None, None, None + + # A trimmed forward output may have an extent which is not divisible by + # the process-group size. Restore the collective extent before + # splitting the gradient; the synthetic tail has no corresponding + # source value and is therefore discarded by the input-side trim. + if ctx.output_dim_size is not None: + padded_dim_size = ctx.input_dim_size * dist.get_world_size(ctx.process_group) + if grad_output.size(ctx.dim) != padded_dim_size: + pad_shape = list(grad_output.shape) + pad_shape[ctx.dim] = padded_dim_size - grad_output.size(ctx.dim) + grad_output = torch.cat((grad_output, grad_output.new_zeros(pad_shape)), dim=ctx.dim) + return ( + _split(grad_output, ctx.dim, ctx.process_group, pad_to_world_size=True), + None, + None, + None, + None, + None, + ) + + return _split(grad_output, ctx.dim, ctx.process_group), None, None, None, None, None class _AllToAll(torch.autograd.Function): @@ -1188,7 +1235,62 @@ def _reduce(input_, process_group, fp8_communication=False, fp8_format="e5m2"): return input_ -def _split(input_, dim=-1, process_group=None): +def _pad_sequence_tensor(tensor, target_length, dim, value=0): + """Pad a tensor along one sequence dimension without changing its dtype.""" + if tensor is None or tensor.size(dim) >= target_length: + return tensor + + pad_shape = list(tensor.shape) + pad_shape[dim] = target_length - tensor.size(dim) + padding = tensor.new_zeros(pad_shape) if value == 0 else tensor.new_full(pad_shape, value) + return torch.cat((tensor, padding), dim=dim) + + +def pad_sequence_parallel_inputs(hidden_states, attention_mask, position_ids, cache_position, target_length): + """Align sequence metadata with a padded sequence-parallel hidden state. + + Equal-sized communication buffers are needed by sequence-parallel + collectives. This helper pads the hidden state and its positional/mask + inputs together; callers can retain the pre-padding length to trim outputs. + """ + hidden_states = _pad_sequence_tensor(hidden_states, target_length, dim=1) + + if attention_mask is not None: + if attention_mask.dim() <= 2: + attention_mask = _pad_sequence_tensor(attention_mask, target_length, dim=-1) + else: + mask_pad_value = torch.finfo(attention_mask.dtype).min if attention_mask.is_floating_point() else 0 + attention_mask = _pad_sequence_tensor(attention_mask, target_length, dim=-1, value=mask_pad_value) + attention_mask = _pad_sequence_tensor(attention_mask, target_length, dim=-2, value=mask_pad_value) + + if position_ids is not None: + position_ids = _pad_sequence_tensor(position_ids, target_length, dim=-1) + + if cache_position is not None: + old_length = cache_position.size(-1) + if old_length < target_length: + if old_length == 0: + next_position = torch.arange(target_length, device=cache_position.device, dtype=cache_position.dtype) + else: + next_position = cache_position[..., -1:] + torch.arange( + 1, + target_length - old_length + 1, + device=cache_position.device, + dtype=cache_position.dtype, + ) + cache_position = torch.cat((cache_position, next_position), dim=-1) + + return hidden_states, attention_mask, position_ids, cache_position + + +def _split(input_, dim=-1, process_group=None, pad_to_world_size=False): + """Return the rank-local chunk, optionally padding for an even split. + + Padding is opt-in because callers other than the sequence-parallel + autograd operation rely on the existing divisibility check. When enabled, + zeros are appended only to the collective buffer; the corresponding + autograd wrapper trims the gathered gradient back to the input extent. + """ # skip if only one rank involved world_size = dist.get_world_size(process_group) if world_size == 1: @@ -1196,10 +1298,18 @@ def _split(input_, dim=-1, process_group=None): # Split along last dimension. dim_size = input_.size(dim) - assert dim_size % world_size == 0, ( - f"The dimension to split ({dim_size}) is not a multiple of world size ({world_size}), " - f"cannot split tensor evenly" - ) + if dim_size % world_size != 0: + if not pad_to_world_size: + raise AssertionError( + f"The dimension to split ({dim_size}) is not a multiple of world size ({world_size}), " + f"cannot split tensor evenly" + ) + + padded_dim_size = ((dim_size + world_size - 1) // world_size) * world_size + pad_shape = list(input_.shape) + pad_shape[dim] = padded_dim_size - dim_size + input_ = torch.cat((input_, input_.new_zeros(pad_shape)), dim=dim) + dim_size = padded_dim_size tensor_list = torch.split(input_, dim_size // world_size, dim=dim) rank = dist.get_rank(process_group) @@ -1355,8 +1465,10 @@ def matmul_gather_forward_reducescatter_backward( ) -def gather_forward_split_backward(input_, dim, process_group, grad_scale=None, fp8_communication=False): - return _GatherForwardSplitBackward.apply(input_, dim, process_group, grad_scale, fp8_communication) +def gather_forward_split_backward( + input_, dim, process_group, grad_scale=None, fp8_communication=False, output_dim_size=None +): + return _GatherForwardSplitBackward.apply(input_, dim, process_group, grad_scale, fp8_communication, output_dim_size) def split_forward_gather_backward(input_, dim, process_group, grad_scale=None, fp8_communication=False): @@ -1375,7 +1487,7 @@ def all_to_all_comm(input_, process_group=None, scatter_dim=2, gather_dim=1, fp8 return _AllToAll.apply(input_, process_group, scatter_dim, gather_dim, fp8_communication) -def gather_sp_output(hidden_states, shard_config, sp_dim=1): +def gather_sp_output(hidden_states, shard_config, sp_dim=1, original_dim_size=None): """ Gather the output of the last layer for cross entropy computation """ @@ -1388,6 +1500,11 @@ def gather_sp_output(hidden_states, shard_config, sp_dim=1): # Rescale grad (HybridParallelPlugin applies ZeRO grad averaging on the DP * SP group) scale = None if is_share_sp_tp(sp_mode) else dist.get_world_size(sp_group) hidden_states = gather_forward_split_backward( - hidden_states, sp_dim, sp_group, grad_scale=scale, fp8_communication=fp8_comm + hidden_states, + sp_dim, + sp_group, + grad_scale=scale, + fp8_communication=fp8_comm, + output_dim_size=original_dim_size, ) return hidden_states diff --git a/colossalai/shardformer/layer/loss.py b/colossalai/shardformer/layer/loss.py index a9bb76fc7d6b..a6307c531e89 100644 --- a/colossalai/shardformer/layer/loss.py +++ b/colossalai/shardformer/layer/loss.py @@ -288,7 +288,19 @@ def dist_cross_entropy( # Shift labels to predict the next token, and remove the tail logit predicting 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 + local_seq_len = (seq_len + sp_size - 1) // sp_size + split_labels_here = is_sp and local_seq_len == logits.size(seq_dim) # ring attn splits labels before forward + + # Sequence-parallel collectives use an equal-sized padded buffer. Extend + # labels with ignored targets so synthetic tail logits never contribute to + # the loss or its gradient. + if split_labels_here: + padded_seq_len = local_seq_len * sp_size + if labels.size(-1) < padded_seq_len: + pad_shape = list(labels.shape) + pad_shape[-1] = padded_seq_len - labels.size(-1) + padding = torch.full(pad_shape, _IGNORE_IDX, dtype=labels.dtype, device=labels.device) + labels = torch.cat((labels, padding), dim=-1) if sp_mode == "ring_attn": # For Zigzag Ring Attention, labels should've been split and @@ -296,12 +308,12 @@ def dist_cross_entropy( if sp_rank == 0: logits = logits[..., :-1, :] logits = torch.cat([logits, torch.full_like(logits[:, :1, :], _IGNORE_IDX)], dim=seq_dim) - elif is_sp: + elif split_labels_here: # 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] + labels = labels.split(local_seq_len, dim=-1)[sp_rank] if sp_rank == sp_size - 1: logits = logits[..., :-1, :] @@ -315,7 +327,16 @@ def dist_cross_entropy( 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) + elif is_sp: + # A gathered sequence-parallel output may retain a synthetic tail. + # Trim it before applying the regular next-token shift. + if logits.size(seq_dim) > seq_len: + logits = logits.narrow(seq_dim, 0, seq_len).contiguous() + labels = labels[..., 1:] + logits = logits[..., :-1, :] else: + if logits.size(seq_dim) > seq_len: + logits = logits.narrow(seq_dim, 0, seq_len).contiguous() labels = labels[..., 1:] logits = logits[..., :-1, :] labels = labels.contiguous() diff --git a/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py index fe102eecf25a..b35d5793d929 100644 --- a/colossalai/shardformer/modeling/llama.py +++ b/colossalai/shardformer/modeling/llama.py @@ -24,7 +24,12 @@ from transformers.utils import logging from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer._operation import all_to_all_comm, gather_sp_output, split_forward_gather_backward +from colossalai.shardformer.layer._operation import ( + all_to_all_comm, + gather_sp_output, + pad_sequence_parallel_inputs, + split_forward_gather_backward, +) from colossalai.shardformer.layer.utils import is_share_sp_tp, split_batch_zigzag from colossalai.shardformer.shard import ShardConfig @@ -97,6 +102,27 @@ def llama_model_forward( sp_mode = shard_config.sequence_parallelism_mode sp_group = shard_config.sequence_parallel_process_group sp_size = shard_config.sequence_parallel_size + logical_seq_length = seq_length + padded_seq_length = seq_length + split_input = disable_pp or stage_manager.is_first_stage() + + # Sequence-parallel collectives require equal-sized buffers. Pad the + # hidden state and every sequence-side argument together, so rotary + # embeddings and attention masks observe the same extent as the + # communication buffer. The original extent is retained for trimming + # the gathered result and for loss alignment. + if split_input and sp_mode in ("all_to_all", "split_gather", "ring") and sp_size > 1: + padded_seq_length = ((seq_length + sp_size - 1) // sp_size) * sp_size + if padded_seq_length != seq_length: + hidden_states, attention_mask, position_ids, cache_position = pad_sequence_parallel_inputs( + hidden_states, + attention_mask, + position_ids, + cache_position, + padded_seq_length, + ) + seq_length = padded_seq_length + # Generating full positions ids for modes that gather sequence before attn if stage_manager and (sp_mode != "ring_attn" and not stage_manager.is_first_stage()): seq_length *= sp_size @@ -145,7 +171,6 @@ def llama_model_forward( ) # Support SP + PP. Later stages have already received the split input. - split_input = disable_pp or stage_manager.is_first_stage() if split_input: # Ring Attention zigzag batch processing if sp_mode == "ring_attn": @@ -229,7 +254,11 @@ def llama_model_forward( if disable_pp or stage_manager.is_last_stage(): hidden_states = self.norm(hidden_states) if (not shard_config.parallel_output) or force_sp_gather or is_share_sp_tp(sp_mode): # noqa - hidden_states = gather_sp_output(hidden_states, shard_config) + hidden_states = gather_sp_output( + hidden_states, + shard_config, + original_dim_size=(logical_seq_length if padded_seq_length != logical_seq_length else None), + ) # add hidden states from the last decoder layer if output_hidden_states: diff --git a/tests/test_shardformer/test_layer/test_uneven_sequence_split.py b/tests/test_shardformer/test_layer/test_uneven_sequence_split.py new file mode 100644 index 000000000000..945a3e4c3083 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_uneven_sequence_split.py @@ -0,0 +1,148 @@ +import torch +import torch.distributed as dist +from torch import nn +from transformers import LlamaConfig, LlamaModel + +from colossalai.shardformer.layer._operation import ( + gather_forward_split_backward, + pad_sequence_parallel_inputs, + split_forward_gather_backward, +) +from colossalai.shardformer.layer.loss import dist_cross_entropy +from colossalai.shardformer.modeling.llama import LlamaPipelineForwards, get_llama_flash_attention_forward +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +def _check_uneven_sequence_split(rank, world_size, port): + dist.init_process_group("gloo", rank=rank, world_size=world_size, init_method=f"tcp://127.0.0.1:{port}") + try: + # Five tokens cannot be evenly split over two sequence-parallel ranks. + input_ = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3) + input_.requires_grad_() + + local = split_forward_gather_backward(input_, dim=1, process_group=dist.group.WORLD) + assert local.shape == (2, 3, 3) + + gathered = [torch.empty_like(local) for _ in range(world_size)] + dist.all_gather(gathered, local) + gathered = torch.cat(gathered, dim=1) + expected = torch.cat((input_.detach(), torch.zeros(2, 1, 3)), dim=1) + torch.testing.assert_close(gathered, expected) + + # A downstream operation may touch padded positions. Those positions + # must not change the shape or values of the input gradient. Rank 1's + # final element is padding, while rank 0's corresponding element is a + # real token, so use different weights to distinguish the two. + weight = torch.ones_like(local) + weight[:, -1] = 10 + rank + (local * weight).sum().backward() + expected_grad = torch.ones_like(input_) + expected_grad[:, 2] = 10 + torch.testing.assert_close(input_.grad, expected_grad) + + # Keep the model-side metadata on the same padded extent as hidden + # states, then trim the gathered result back to the logical length. + hidden = torch.randn(2, 5, 3, requires_grad=True) + attention_mask = torch.ones(2, 5, dtype=torch.long) + position_ids = torch.arange(5, dtype=torch.long).unsqueeze(0) + cache_position = torch.arange(5, dtype=torch.long) + hidden, attention_mask, position_ids, cache_position = pad_sequence_parallel_inputs( + hidden, attention_mask, position_ids, cache_position, target_length=6 + ) + hidden.retain_grad() + assert hidden.shape == (2, 6, 3) + assert attention_mask.shape == (2, 6) + assert attention_mask[:, -1].eq(0).all() + assert position_ids.shape == (1, 6) + assert cache_position.tolist() == [0, 1, 2, 3, 4, 5] + + local_hidden = split_forward_gather_backward(hidden, dim=1, process_group=dist.group.WORLD) + logical_hidden = gather_forward_split_backward( + local_hidden, dim=1, process_group=dist.group.WORLD, output_dim_size=5 + ) + assert logical_hidden.shape == (2, 5, 3) + logical_hidden.sum().backward() + assert hidden.grad is not None and hidden.grad.shape == hidden.shape + assert hidden.grad[:, -1].eq(0).all() + + # The language-model loss consumes the local padded logits and the + # original (uneven) labels. Ignored targets must mask the synthetic + # tail while the reduced loss still match the unsharded reference. + vocab_size = 7 + global_logits = ( + torch.arange(2 * 5 * vocab_size, dtype=torch.float32).reshape(2, 5, vocab_size) / vocab_size + ).requires_grad_() + local_logits = split_forward_gather_backward(global_logits, dim=1, process_group=dist.group.WORLD) + loss_config = type("ShardConfig", (), {})() + loss_config.sequence_parallel_process_group = dist.group.WORLD + loss_config.sequence_parallel_size = world_size + loss_config.sequence_parallelism_mode = "all_to_all" + loss_config.parallel_output = True + loss_config.enable_tensor_parallelism = False + labels = torch.tensor([[1, 2, 3, 4, 5], [2, 3, 4, 5, 6]]) + loss = dist_cross_entropy(labels, local_logits, loss_config, vocab_size, global_logits.dtype) + reference = nn.functional.cross_entropy( + global_logits.detach()[:, :-1].reshape(-1, vocab_size), + labels[:, 1:].reshape(-1), + reduction="mean", + ) + torch.testing.assert_close(loss, reference) + loss.backward() + assert global_logits.grad is not None and global_logits.grad.shape == global_logits.shape + assert global_logits.grad[:, -1].abs().sum() == 0 + + # Exercise the Llama model path: all sequence-side metadata must use + # the padded extent while the public output keeps the original length. + llama_config = LlamaConfig( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + ) + llama_config._attn_implementation = "eager" + llama = LlamaModel(llama_config) + attention = llama.layers[0].self_attn + # Model-parallel projections expose a head-sharded shape to the + # all-to-all path, matching the dimensions used after tensor sharding. + attention.num_heads = 1 + attention.num_key_value_heads = 1 + attention.num_key_value_groups = 1 + attention.head_dim = 4 + attention.q_proj = nn.Linear(8, 8, bias=False) + attention.k_proj = nn.Linear(8, 8, bias=False) + attention.v_proj = nn.Linear(8, 8, bias=False) + attention.o_proj = nn.Linear(8, 8, bias=False) + + shard_config = type("ShardConfig", (), {})() + shard_config.sequence_parallelism_mode = "all_to_all" + shard_config.sequence_parallel_process_group = dist.group.WORLD + shard_config.sequence_parallel_size = world_size + shard_config.enable_sequence_parallelism = True + shard_config.enable_flash_attention = False + shard_config.fp8_communication = False + shard_config.parallel_output = False + shard_config.gradient_checkpoint_config = None + attention.forward = get_llama_flash_attention_forward( + shard_config, + sp_mode="all_to_all", + sp_size=world_size, + sp_group=dist.group.WORLD, + ).__get__(attention, type(attention)) + + model_output = LlamaPipelineForwards.llama_model_forward( + llama, + input_ids=torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]), + attention_mask=torch.ones(2, 5, dtype=torch.long), + shard_config=shard_config, + return_dict=True, + ) + assert model_output.last_hidden_state.shape == (2, 5, 8) + model_output.last_hidden_state.square().mean().backward() + finally: + dist.destroy_process_group() + + +@rerun_if_address_is_in_use() +def test_uneven_sequence_split_cpu(): + spawn(_check_uneven_sequence_split, nprocs=2)