From 8cf638ae303073926f0b457a2453475047027093 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 8 Sep 2026 10:19:48 +0000 Subject: [PATCH 1/6] Add vocab-parallel LM loss Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/module_inject/auto_tp.py | 37 +- deepspeed/module_inject/layers.py | 18 +- deepspeed/module_inject/replace_module.py | 1 + deepspeed/runtime/engine.py | 29 +- deepspeed/runtime/tensor_parallel/config.py | 3 + deepspeed/sequence/__init__.py | 3 + deepspeed/sequence/cross_entropy.py | 357 ++++++++++++++++-- docs/_pages/config-json.md | 7 + docs/_tutorials/autotp-training.md | 53 ++- .../test_tp_partition_config_path.py | 89 ++++- .../test_vocab_parallel_cross_entropy.py | 292 ++++++++++++++ 11 files changed, 842 insertions(+), 47 deletions(-) create mode 100644 tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 31892edb9d19..e42a2afb9726 100755 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -207,6 +207,7 @@ def __init__(self, orig_layer_impl, keep_module_on_host=False, partition_config: Optional[AutoTPConfig] = None, + vocab_parallel_lm_head=False, model_config=None, tp_grain_size: int = 1, training_mode: bool = False): @@ -225,6 +226,7 @@ def __init__(self, self.linear_policies = None self.conv_linear_layer = False self.partition_config = partition_config + self.vocab_parallel_lm_head = vocab_parallel_lm_head self.training_mode = training_mode self._gathered_column_tie_fallbacks_configured = False self._tied_gathered_column_module_names = set() @@ -372,6 +374,9 @@ def _replace(self, child, name, conv_linear_layer): if getattr(child, "_is_autoep_layer", False): return child + if self._is_vocab_parallel_lm_head(child, name): + return self._create_vocab_parallel_layer(child, name) + weight_shape = child.weight.shape mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group) @@ -436,6 +441,9 @@ def _replace_with_config(self, child, name): if getattr(child, "replaced", False) == True: return child + if self._is_vocab_parallel_lm_head(child, name): + return self._create_vocab_parallel_layer(child, name) + # Build the full parameter name for pattern matching param_name = name + ".weight" if not name.endswith(".weight") else name @@ -492,7 +500,7 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str): gather_output=spec.gather_output, tp_meta=self.tp_meta) # Only use fused-QKV heuristics when no partition_config is provided. - elif self.partition_config is None and require_tp_fused_qkvw(name, self.mp_size): + if self.partition_config is None and require_tp_fused_qkvw(name, self.mp_size): # Check and handle fused qkv for TP return fused_LinearLayer(module, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta) if spec.shape is not None: @@ -508,6 +516,31 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str): ) return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output, tp_meta=self.tp_meta) + @staticmethod + def _is_lm_head_name(name): + # Only the final path segment may match, so auxiliary projections whose names + # merely contain "lm_head" (e.g. "lm_head_proj") are never captured. + return str(name).split('.')[-1] in ("lm_head", "embed_out") + + def _is_vocab_parallel_lm_head(self, child, name): + # VocabParallelLinear assumes an [vocab, hidden] nn.Linear weight; a Conv1D head + # stores [hidden, vocab] and would be cut on the wrong dimension. + return self.vocab_parallel_lm_head and isinstance(child, nn.Linear) and self._is_lm_head_name(name) + + def _create_vocab_parallel_layer(self, child, name): + self._validate_untied_vocab_head(child) + setattr(child, "replaced", True) + log_dist( + f"AutoTP: vocab_parallel_lm_head keeps '{name}' vocabulary-sharded and installs the " + f"distributed causal-LM loss", + ranks=[0]) + return VocabParallelLinear(child, self.mp_group, name=name, tp_meta=self.tp_meta) + + def _validate_untied_vocab_head(self, lm_head): + for _, module in self.module.named_modules(): + if isinstance(module, nn.Embedding) and getattr(module, "weight", None) is lm_head.weight: + raise ValueError("A no-gather vocab-parallel LM head requires untied embedding and output weights") + def _configure_gathered_column_tie_fallbacks(self): """Configure a replicated fallback for gathered output layers tied to embeddings.""" if self._gathered_column_tie_fallbacks_configured: @@ -724,7 +757,7 @@ def _replace_autoep_shared_experts(self, autoep_layer, autoep_name): self.update_mp_params(child, full_name) def _replace_module(self, r_module, prev_name='', prev_class_name=''): - if prev_name == '' and prev_class_name == '': + if prev_name == '' and prev_class_name == '' and not self.vocab_parallel_lm_head: self._configure_gathered_column_tie_fallbacks() for name, child in r_module.named_children(): diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py index 6c8a862dbb8d..4e9ac44dcc04 100644 --- a/deepspeed/module_inject/layers.py +++ b/deepspeed/module_inject/layers.py @@ -27,7 +27,7 @@ __all__ = [ "TensorParallel_Layer", "LinearAllreduce", "LinearLayer", "LmHeadLinearAllreduce", "Yuan_LinearAllreduce", "Yuan_LinearLayer", "GateUpPack_LinearLayer", "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer", - "SubParamLinearLayer", "SubParamLinearAllreduce" + "SubParamLinearLayer", "SubParamLinearAllreduce", "VocabParallelLinear" ] DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE @@ -960,6 +960,22 @@ def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=Non return cls(linear, skip_partition=True, gather_output=gather_output) +class VocabParallelLinear(LinearLayer): + """Column-parallel vocabulary projection that keeps rank-local logits.""" + + def __init__(self, module, mp_group=None, **kwargs): + super().__init__(module, mp_group, gather_output=False, **kwargs) + if min(self._partition_sizes) == 0: + # The shard-size list is identical on every TP rank, so all ranks raise here + # together instead of one rank failing into a collective hang at the loss. + raise ValueError(f"vocab_parallel_lm_head requires a vocabulary of at least tp_size=" + f"{self.tp_world_size} rows, but '{self.name}' has {self._orig_weight_shape[0]}") + self.is_vocab_parallel_lm_head = True + self.vocab_size = self._orig_weight_shape[0] + self.vocab_start_index = sum(self._partition_sizes[:self.tp_index]) + self.vocab_end_index = self.vocab_start_index + self._partition_sizes[self.tp_index] + + class SubParamColumnParallel(LinearLayer): """Column-parallel layer whose shard concatenates one piece of every sub-parameter. diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 0424285bfb9c..574ceb13b8d1 100644 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -299,6 +299,7 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None): orig_layer_impl, config.keep_module_on_host, partition_config=partition_config, + vocab_parallel_lm_head=getattr(config, "vocab_parallel_lm_head", False), model_config=meta_config, tp_grain_size=config.tensor_parallel.tp_grain_size, training_mode=training_mode) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..888675afece7 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -846,6 +846,22 @@ def _apply_autotp_partitioning(self, model, tp_config): from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan hf_tp_plan = _get_hf_tp_plan(model) + def finalize_autotp(autotp=None, attach_uc_metadata=False): + if autotp is not None: + autotp.register_replicated_grad_hooks(model) + + from deepspeed.module_inject.layers import VocabParallelLinear + vocab_parallel_heads = [module for module in model.modules() if isinstance(module, VocabParallelLinear)] + if len(vocab_parallel_heads) > 1: + raise ValueError("Unable to choose a loss for multiple no-gather vocab-parallel LM heads") + if vocab_parallel_heads: + from deepspeed.sequence.cross_entropy import configure_vocab_parallel_loss + configure_vocab_parallel_loss(model, vocab_parallel_heads[0]) + + if attach_uc_metadata: + setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model)) + setattr(model, "ds_autotp_parsed", True) + if partition_config is not None: autotp = AutoTP(module=model, all_reduce_linears=(), @@ -855,15 +871,14 @@ def _apply_autotp_partitioning(self, model, tp_config): orig_layer_impl=None, keep_module_on_host=tp_config.keep_module_on_host, partition_config=partition_config, + vocab_parallel_lm_head=tp_config.vocab_parallel_lm_head, model_config=model_config, tp_grain_size=tp_config.tensor_parallel.tp_grain_size, training_mode=True) autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group) autotp.update_linear_policies() autotp._replace_module(model) - autotp.register_replicated_grad_hooks(model) - setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model)) - setattr(model, "ds_autotp_parsed", True) + finalize_autotp(autotp, attach_uc_metadata=True) return if tp_size <= 1: @@ -896,6 +911,7 @@ def _apply_autotp_partitioning(self, model, tp_config): orig_layer_impl=None, keep_module_on_host=tp_config.keep_module_on_host, partition_config=tp_plan_config, + vocab_parallel_lm_head=tp_config.vocab_parallel_lm_head, model_config=model_config, tp_grain_size=tp_config.tensor_parallel.tp_grain_size, training_mode=True, @@ -903,9 +919,7 @@ def _apply_autotp_partitioning(self, model, tp_config): autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group) autotp.update_linear_policies() autotp._replace_module(model) - autotp.register_replicated_grad_hooks(model) - setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model)) - setattr(model, "ds_autotp_parsed", True) + finalize_autotp(autotp, attach_uc_metadata=True) return log_dist( f"AutoTP: effective HuggingFace tp_plan could not be converted; falling back to heuristic AutoTP. " @@ -921,8 +935,7 @@ def _apply_autotp_partitioning(self, model, tp_config): tp_config.injection_policy_tuple = injection_policy replace_transformer_layer(client_module, model, None, tp_config, model_config, training_mode=True) - setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model)) - setattr(model, "ds_autotp_parsed", True) + finalize_autotp(attach_uc_metadata=True) def __del__(self): try: diff --git a/deepspeed/runtime/tensor_parallel/config.py b/deepspeed/runtime/tensor_parallel/config.py index fb0ac10b9fcd..9d196eca63df 100644 --- a/deepspeed/runtime/tensor_parallel/config.py +++ b/deepspeed/runtime/tensor_parallel/config.py @@ -50,6 +50,9 @@ class TPTrainingConfig(DeepSpeedConfigModel): tp_overlap_comm: bool = False """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """ + vocab_parallel_lm_head: bool = False + """Keep an untied LM head vocabulary-sharded and install a compatible distributed loss.""" + tensor_parallel: TPConfig = Field({}, alias="tp") """ Configuration for tensor parallelism used to split the model across several diff --git a/deepspeed/sequence/__init__.py b/deepspeed/sequence/__init__.py index b76f944eff79..c61198be17c5 100644 --- a/deepspeed/sequence/__init__.py +++ b/deepspeed/sequence/__init__.py @@ -8,3 +8,6 @@ from deepspeed.sequence.autosp_fusion import (ModalityFusionSPAdapter, LlavaFusionAdapter, InternVLFusionAdapter, Qwen2VLFusionAdapter) from deepspeed.sequence.auto_sp import auto_wrap_model_for_sp +from deepspeed.sequence.cross_entropy import (VocabParallelCausalLMLoss, VocabParallelCrossEntropyLoss, + configure_vocab_parallel_loss, vocab_parallel_cross_entropy, + vocab_sequence_parallel_cross_entropy) diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py index baa7bc1ea7a8..f9275e84aab7 100644 --- a/deepspeed/sequence/cross_entropy.py +++ b/deepspeed/sequence/cross_entropy.py @@ -4,57 +4,346 @@ # DeepSpeed Team import torch +from torch import nn import deepspeed.comm as dist +from deepspeed.utils.logging import logger -class _VocabSequenceParallelCrossEntropy(torch.autograd.Function): +class _VocabParallelCrossEntropy(torch.autograd.Function): @staticmethod - def forward(ctx, vocab_seq_parallel_logits, target, sp_group): - # vocab_seq_parallel_logits: [S/P, B, V] - # target: [S/P, B] - # return: [S, B] + def forward(ctx, vocab_parallel_logits, target, tp_group, vocab_start_index, vocab_end_index, ignore_index): + target = target.to(dtype=torch.long) + logits = vocab_parallel_logits.float() + local_vocab_size = logits.shape[-1] - # Need softmax for backward - softmax = torch.nn.functional.softmax(vocab_seq_parallel_logits, dim=-1) - ctx.vocab_size = vocab_seq_parallel_logits.size(2) - loss = torch.nn.functional.nll_loss(softmax.log().view(-1, ctx.vocab_size), target.view(-1), reduction='none') + local_max = logits.amax(dim=-1) + global_max = local_max.clone() + if tp_group is not None and dist.get_world_size(tp_group) > 1: + dist.all_reduce(global_max, op=dist.ReduceOp.MAX, group=tp_group) - sp_world_size = dist.get_world_size(sp_group) - sp_rank = dist.get_rank(sp_group) - ctx.sp_world_size = sp_world_size - ctx.sp_rank = sp_rank - ctx.seqlen = vocab_seq_parallel_logits.size(0) * sp_world_size - batch_size = vocab_seq_parallel_logits.size(1) + exp_logits = torch.exp(logits - global_max.unsqueeze(-1)) + global_sum_exp = exp_logits.sum(dim=-1) + if tp_group is not None and dist.get_world_size(tp_group) > 1: + dist.all_reduce(global_sum_exp, op=dist.ReduceOp.SUM, group=tp_group) - loss_all = torch.empty(ctx.seqlen, - batch_size, - dtype=vocab_seq_parallel_logits.dtype, - device=vocab_seq_parallel_logits.device) - dist.all_gather_into_tensor(loss_all, loss, group=sp_group) + valid_target = target != ignore_index + target_in_partition = valid_target & (target >= vocab_start_index) & (target < vocab_end_index) + local_target = (target - vocab_start_index).clamp(min=0, max=local_vocab_size - 1) + target_logits = logits.gather(-1, local_target.unsqueeze(-1)).squeeze(-1) + target_logits = torch.where(target_in_partition, target_logits, torch.zeros_like(target_logits)) + if tp_group is not None and dist.get_world_size(tp_group) > 1: + dist.all_reduce(target_logits, op=dist.ReduceOp.SUM, group=tp_group) - ctx.save_for_backward(softmax, target) + loss = torch.log(global_sum_exp) + global_max - target_logits + loss = torch.where(valid_target, loss, torch.zeros_like(loss)) - return loss_all + ctx.save_for_backward(exp_logits, global_sum_exp, local_target, target_in_partition, valid_target) + ctx.logits_dtype = vocab_parallel_logits.dtype + return loss @staticmethod def backward(ctx, grad_output): - softmax, target = ctx.saved_tensors + exp_logits, global_sum_exp, local_target, target_in_partition, valid_target = ctx.saved_tensors - step_seqlen = ctx.seqlen // ctx.sp_world_size - sp_rank = ctx.sp_rank - grad_output_part = grad_output[step_seqlen * sp_rank:step_seqlen * (sp_rank + 1), :] + grad_logits = exp_logits / global_sum_exp.unsqueeze(-1) + grad_logits.scatter_add_(-1, local_target.unsqueeze(-1), + -target_in_partition.unsqueeze(-1).to(dtype=grad_logits.dtype)) + grad_logits *= valid_target.unsqueeze(-1).to(dtype=grad_logits.dtype) + grad_logits *= grad_output.to(dtype=grad_logits.dtype).unsqueeze(-1) - grad_input = softmax - grad_2d = grad_input.view(-1, ctx.vocab_size) - arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) + return grad_logits.to(dtype=ctx.logits_dtype), None, None, None, None, None - grad_2d[arange_1d, target.view(-1)] -= 1 - grad_input.mul_(grad_output_part.unsqueeze(dim=-1)) - return grad_input, None, None, None +class _GatherSequenceLoss(torch.autograd.Function): + @staticmethod + def forward(ctx, local_loss, sp_group): + ctx.sp_group = sp_group + ctx.local_sequence_size = local_loss.shape[0] + + output_shape = (ctx.local_sequence_size * dist.get_world_size(sp_group), *local_loss.shape[1:]) + gathered_loss = torch.empty(output_shape, dtype=local_loss.dtype, device=local_loss.device) + dist.all_gather_into_tensor(gathered_loss, local_loss.contiguous(), group=sp_group) + return gathered_loss + + @staticmethod + def backward(ctx, grad_output): + grad_input = torch.empty((ctx.local_sequence_size, *grad_output.shape[1:]), + dtype=grad_output.dtype, + device=grad_output.device) + dist.reduce_scatter_fn(grad_input, grad_output.contiguous(), group=ctx.sp_group) + return grad_input, None + + +def _global_sp_sum(local_value, sp_group): + if sp_group is None or dist.get_world_size(sp_group) == 1: + return local_value + + global_value = local_value.detach().clone() + dist.all_reduce(global_value, op=dist.ReduceOp.SUM, group=sp_group) + return local_value + (global_value - local_value.detach()) + + +def _validate_vocab_shard_bounds(local_vocab_size, vocab_start_index, vocab_end_index, tp_group, device): + tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1 + if tp_world_size == 1: + if vocab_start_index != 0 or vocab_end_index != local_vocab_size: + raise ValueError("Vocabulary shard bounds must cover the complete local vocabulary when TP is disabled") + return vocab_end_index + + local_metadata = torch.tensor([vocab_start_index, vocab_end_index, local_vocab_size], + dtype=torch.long, + device=device) + gathered_metadata = torch.empty(tp_world_size * local_metadata.numel(), dtype=local_metadata.dtype, device=device) + dist.all_gather_into_tensor(gathered_metadata, local_metadata, group=tp_group) + + expected_start = 0 + for rank, (shard_start, shard_end, shard_size) in enumerate(gathered_metadata.view(tp_world_size, 3).tolist()): + if shard_end - shard_start != shard_size: + raise ValueError(f"Vocabulary shard bounds for TP rank {rank} do not match its local vocabulary size") + if shard_size <= 0: + raise ValueError(f"TP rank {rank} received an empty vocabulary shard; the vocabulary must be at least " + f"as large as the tensor-parallel size") + if shard_start != expected_start: + raise ValueError("Vocabulary shard bounds must form a contiguous, non-overlapping partition starting at 0") + expected_start = shard_end + + return expected_start + + +_vocab_metadata_cache = {} + + +def _resolve_vocab_metadata(local_vocab_size, vocab_start_index, vocab_end_index, tp_group, device): + """Collectively validate the vocabulary shard layout once and cache the result. + + The shard geometry is fixed for the lifetime of the layers, so the repeated calls a + training loop makes every micro-batch must not re-run the validation collectives or + their host synchronizations. Decisions are made from identical all-gathered data, so + every TP rank raises together instead of diverging into a collective hang. + """ + key = (tp_group, local_vocab_size, vocab_start_index, vocab_end_index) + cached = _vocab_metadata_cache.get(key) + if cached is not None: + return cached + + if vocab_start_index is None: + tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1 + tp_rank = dist.get_rank(tp_group) if tp_group is not None else 0 + if tp_world_size > 1: + local_size = torch.tensor(local_vocab_size, device=device, dtype=torch.long) + min_local_size = local_size.clone() + max_local_size = local_size.clone() + dist.all_reduce(min_local_size, op=dist.ReduceOp.MIN, group=tp_group) + dist.all_reduce(max_local_size, op=dist.ReduceOp.MAX, group=tp_group) + if min_local_size.item() != max_local_size.item(): + raise ValueError("Explicit vocabulary shard bounds are required for uneven tensor-parallel shards") + vocab_start_index = tp_rank * local_vocab_size + vocab_end_index = vocab_start_index + local_vocab_size + global_vocab_size = _validate_vocab_shard_bounds(local_vocab_size, vocab_start_index, vocab_end_index, tp_group, + device) + + metadata = (vocab_start_index, vocab_end_index, global_vocab_size) + _vocab_metadata_cache[key] = metadata + return metadata + + +def vocab_parallel_cross_entropy(vocab_parallel_logits, + target, + tp_group=None, + sp_group=None, + vocab_start_index=None, + vocab_end_index=None, + ignore_index=-100, + reduction="mean", + gather_sequence_loss=False): + """Compute cross entropy over vocabulary-sharded logits. + + Tensor parallel ranks collectively own the last (vocabulary) dimension. Sequence + parallel ranks may independently own shards of the leading sequence dimension. + """ + if vocab_parallel_logits.shape[:-1] != target.shape: + raise ValueError("vocab_parallel_logits and target must have matching non-vocabulary dimensions") + # With tensor parallelism an empty shard is rejected from the all-gathered shard + # metadata so every rank fails together; only the single-process case can raise here. + if vocab_parallel_logits.shape[-1] == 0 and (tp_group is None or dist.get_world_size(tp_group) == 1): + raise ValueError("vocab_parallel_logits must contain at least one local vocabulary entry") + if reduction not in ("none", "sum", "mean"): + raise ValueError(f"Unsupported reduction: {reduction}") + if gather_sequence_loss and reduction != "none": + raise ValueError("gather_sequence_loss is only supported with reduction='none'") + + local_vocab_size = vocab_parallel_logits.shape[-1] + if (vocab_start_index is None) != (vocab_end_index is None): + raise ValueError("vocab_start_index and vocab_end_index must be provided together") + + vocab_start_index, vocab_end_index, global_vocab_size = _resolve_vocab_metadata( + local_vocab_size, vocab_start_index, vocab_end_index, tp_group, vocab_parallel_logits.device) + # Data-dependent, so it cannot be hoisted out of the training loop: an out-of-range + # target belongs to no shard and would otherwise silently contribute a wrong, finite loss. + invalid_target = (target != ignore_index) & ((target < 0) | (target >= global_vocab_size)) + if invalid_target.any().item(): + raise ValueError(f"Target is out of range for vocabulary size {global_vocab_size}") + + loss = _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, tp_group, vocab_start_index, + vocab_end_index, ignore_index) + if reduction == "none": + if gather_sequence_loss: + if sp_group is None: + raise ValueError("sp_group is required when gather_sequence_loss=True") + loss = _GatherSequenceLoss.apply(loss, sp_group) + return loss + + loss_sum = _global_sp_sum(loss.sum(), sp_group) + if reduction == "sum": + return loss_sum + + valid_tokens = (target != ignore_index).sum().to(dtype=loss.dtype) + if sp_group is not None and dist.get_world_size(sp_group) > 1: + dist.all_reduce(valid_tokens, op=dist.ReduceOp.SUM, group=sp_group) + return loss_sum / valid_tokens.clamp_min(1) + + +def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, + target, + sp_group, + tp_group=None, + vocab_start_index=None, + vocab_end_index=None, + ignore_index=-100, + reduction="none", + gather_sequence_loss=True): + """Sequence-parallel wrapper over :func:`vocab_parallel_cross_entropy`. + + Backward reduce-scatters the gradient over ``sp_group``, so each rank receives the + gradient of its own sequence shard with the other ranks' contributions already + summed in. Downstream code must therefore treat the returned loss as replicated + across the SP group and must not average SP gradients a second time. + """ + return vocab_parallel_cross_entropy(vocab_parallel_logits, + target, + tp_group=tp_group, + sp_group=sp_group, + vocab_start_index=vocab_start_index, + vocab_end_index=vocab_end_index, + ignore_index=ignore_index, + reduction=reduction, + gather_sequence_loss=gather_sequence_loss) + + +class VocabParallelCrossEntropyLoss(nn.Module): + + def __init__(self, + tp_group=None, + sp_group=None, + vocab_start_index=None, + vocab_end_index=None, + ignore_index=-100, + reduction="mean", + gather_sequence_loss=False): + super().__init__() + self.tp_group = tp_group + self.sp_group = sp_group + self.vocab_start_index = vocab_start_index + self.vocab_end_index = vocab_end_index + self.ignore_index = ignore_index + self.reduction = reduction + self.gather_sequence_loss = gather_sequence_loss + + def forward(self, vocab_parallel_logits, target): + return vocab_parallel_cross_entropy(vocab_parallel_logits, + target, + tp_group=self.tp_group, + sp_group=self.sp_group, + vocab_start_index=self.vocab_start_index, + vocab_end_index=self.vocab_end_index, + ignore_index=self.ignore_index, + reduction=self.reduction, + gather_sequence_loss=self.gather_sequence_loss) + + +class VocabParallelCausalLMLoss(nn.Module): + """Distributed causal-LM loss for a vocabulary-sharded (no-gather) LM head. + + ``sp_group`` must stay ``None`` under DeepSpeed's Ulysses sequence-parallel engine: + that engine aggregates the per-shard means itself, weighted by each shard's + valid-token count, so an additional SP reduction here would double-count tokens. + Pass ``sp_group`` only when this loss is the sole aggregation over a manually + constructed TP x SP process-group mesh. + """ + + def __init__(self, tp_group=None, sp_group=None, vocab_start_index=None, vocab_end_index=None, ignore_index=-100): + super().__init__() + self.tp_group = tp_group + self.sp_group = sp_group + self.vocab_start_index = vocab_start_index + self.vocab_end_index = vocab_end_index + self.ignore_index = ignore_index + + def forward(self, logits, labels=None, vocab_size=None, shift_labels=None, num_items_in_batch=None, **kwargs): + if shift_labels is None: + if labels is None: + raise ValueError("labels or shift_labels must be provided") + shift_labels = labels[..., 1:].contiguous() + logits = logits[..., :-1, :].contiguous() + else: + shift_labels = shift_labels.contiguous() + + if vocab_size is not None and self.vocab_start_index is not None and self.vocab_end_index is not None: + # The LM head's shard metadata is the source of truth; a mismatch usually means + # the embedding was resized, which gathered loss implementations tolerated. + _, _, global_vocab_size = _resolve_vocab_metadata(self.vocab_end_index - self.vocab_start_index, + self.vocab_start_index, self.vocab_end_index, + self.tp_group, logits.device) + if vocab_size != global_vocab_size: + logger.warning_once(f"Vocab-parallel LM head holds vocab_size={global_vocab_size}, but the caller " + f"described vocab_size={vocab_size}; the LM head's weights win") + + reduction = "sum" if num_items_in_batch is not None else "mean" + loss = vocab_parallel_cross_entropy(logits, + shift_labels, + tp_group=self.tp_group, + sp_group=self.sp_group, + vocab_start_index=self.vocab_start_index, + vocab_end_index=self.vocab_end_index, + ignore_index=self.ignore_index, + reduction=reduction) + if num_items_in_batch is not None: + denominator = torch.as_tensor(num_items_in_batch, device=loss.device, dtype=loss.dtype) + loss = loss / denominator.clamp_min(1) + return loss + + +def configure_vocab_parallel_loss(model, vocab_parallel_head, sp_group=None, ignore_index=-100): + """Install the causal-LM loss required by a no-gather vocabulary projection. + + Leave ``sp_group`` as ``None`` when running under DeepSpeed's Ulysses + sequence-parallel engine: the engine performs the token-count-weighted aggregation + across SP ranks itself and expects this loss to return the local shard's mean. + """ + if not hasattr(model, "loss_function"): + raise ValueError("A no-gather vocab-parallel LM head requires a writable loss_function hook; " + "use gather_output=True for models without one") + + loss_fn = VocabParallelCausalLMLoss(tp_group=vocab_parallel_head.mp_group, + sp_group=sp_group, + vocab_start_index=vocab_parallel_head.vocab_start_index, + vocab_end_index=vocab_parallel_head.vocab_end_index, + ignore_index=ignore_index) + # Keep the stock loss reachable so callers can restore it when tearing the head down. + if not hasattr(model, "_deepspeed_original_loss_function"): + model._deepspeed_original_loss_function = model.loss_function + + # Some model classes expose loss_function as a read-only property, in which case the + # assignment raises; the identity check below turns that into an actionable error + # instead of leaving the model silently computing loss on rank-local logits. + try: + model.loss_function = loss_fn + except AttributeError: + pass -def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, target, sp_group): - return _VocabSequenceParallelCrossEntropy.apply(vocab_parallel_logits, target, sp_group) + if model.loss_function is not loss_fn: + raise ValueError("Unable to install the vocab-parallel loss_function hook; use gather_output=True") + return model diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index cd1c2def9562..ec7209972bf3 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -785,6 +785,7 @@ When a HuggingFace model provides a built-in `tp_plan` (via `model.config.base_m "autotp_size": 4, "preset_model": "llama", "tp_overlap_comm": false, + "vocab_parallel_lm_head": false, "partition_config": { "use_default_specs": false, "layer_specs": [ @@ -820,6 +821,12 @@ When a HuggingFace model provides a built-in `tp_plan` (via `model.config.base_m | -------------------------------------------------------------------------------------------------------- | ------- | | Overlap tensor-parallel allreduce communication with computation (training only). | `false` | +***vocab_parallel_lm_head***: [boolean] + +| Description | Default | +| ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------- | +| Keep an untied `lm_head`/`embed_out` output vocabulary sharded and install DeepSpeed's pure-PyTorch vocab-parallel causal-LM loss instead of gathering logits. | `false` | + ***partition_config***: [dictionary] | Description | Default | diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md index a8c9304ba55c..19a3839691b3 100644 --- a/docs/_tutorials/autotp-training.md +++ b/docs/_tutorials/autotp-training.md @@ -9,6 +9,7 @@ This tutorial covers **Automatic Tensor Parallelism** for combining tensor paral - [Introduction](#introduction) - [Quick Start](#quick-start) - [HuggingFace tp_plan Support](#huggingface-tp_plan-support) +- [Vocabulary-parallel LM Loss](#vocabulary-parallel-lm-loss) - [Custom Layer Specifications](#custom-layer-specifications) - [Limitations](#limitations) @@ -142,6 +143,52 @@ If you need to override the model's built-in `tp_plan`, provide a `partition_config` in the DeepSpeed config -- it takes precedence. +## Vocabulary-parallel LM Loss + +Causal language models normally gather the complete `lm_head` output before +computing cross entropy. To keep an untied output vocabulary sharded, enable +`vocab_parallel_lm_head`: + +```json +{ + "train_micro_batch_size_per_gpu": 1, + "zero_optimization": { "stage": 2 }, + "tensor_parallel": { + "autotp_size": 4, + "vocab_parallel_lm_head": true + } +} +``` + +DeepSpeed then keeps `lm_head` (or `embed_out`) local to each TP rank and +installs a pure-PyTorch distributed causal-LM loss through the model's +`loss_function` hook. The loss computes a numerically stable distributed +log-sum-exp and target lookup without gathering vocabulary logits. Uneven +vocabulary shards are supported. + +This option requires an untied output head and a model with a writable +`loss_function` hook. Models that share the output weight with the input +embedding must continue using gathered output until coupled vocabulary-parallel +embedding support is available. The head's vocabulary must also be at least as +large as `autotp_size`; smaller vocabularies fail at startup instead of leaving +TP ranks with empty shards. The flag itself is the only trigger: a `colwise` +`lm_head` specification with local output keeps the previous behavior of +returning rank-local logits without installing the distributed loss. + +The lower-level `vocab_parallel_cross_entropy` API also accepts an explicit +sequence-parallel group. With `reduction="none"`, callers may return local token +losses or gather them along sequence dimension 0. `sum` and `mean` reduce over +the supplied SP group. TP and SP may be combined with explicit orthogonal +process groups, but AutoTP does not currently construct a combined TP x SP mesh +automatically. + +Under DeepSpeed's Ulysses sequence-parallel engine, the installed loss must keep +its default sequence-parallel settings (no `sp_group`): the engine aggregates +each shard's mean itself, weighted by the shard's valid-token count, so an +additional SP reduction would double-count tokens. Pass an `sp_group` only when +this loss is the sole aggregation over a manually constructed TP x SP mesh. + + ## Custom Patterns If you are training a custom model, define regex-based patterns and partition rules in `tensor_parallel.partition_config`: @@ -229,7 +276,11 @@ For Grouped Query Attention with different Q/K/V sizes: 1. **Ranks beyond the key/value head count stay idle**: Attention heads are distributed whole, and the distribution may be uneven -- 6 key/value heads over 4 ranks becomes 2/2/1/1, and a fused QKV weight is cut on the same head boundaries rather than inside a head. With more ranks than key/value heads, for example an 8-head model at `autotp_size=16`, the surplus ranks receive no attention weights at all. The result is still correct, because those ranks contribute zeros to the row-parallel all-reduce, but they do no attention work; AutoTP logs a warning instead of replicating heads to fill them. Hidden and vocabulary dimensions do not need to be divisible by the tensor parallel size: uneven shards are carried through save, conversion and restore via per-TP-rank shapes and widths. -2. **Cross-topology universal restore**: Loading a universal checkpoint back into a topology with a *different* tensor-parallel degree goes through DeepSpeed's Megatron-style model-state loader, which is not AutoTP-aware; prefer same-topology restore when changing world size. +2. **Vocabulary-parallel tied weights**: `vocab_parallel_lm_head` requires an untied output projection and a writable model `loss_function` hook, and a vocabulary at least as large as the tensor-parallel size. Coupled sharding of a tied input embedding and output projection is not yet implemented. + +3. **Combined TP and SP orchestration**: The vocab-parallel loss supports explicit orthogonal TP and SP groups, but AutoTP does not currently construct a combined TP x SP process mesh automatically. + +4. **Cross-topology universal restore**: Loading a universal checkpoint back into a topology with a *different* tensor-parallel degree goes through DeepSpeed's Megatron-style model-state loader, which is not AutoTP-aware; prefer same-topology restore when changing world size. ## See Also diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py index e43da01c6bad..9bd96ebbc1fe 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -12,8 +12,10 @@ import torch.nn as nn from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec -from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, set_autotp_mode +from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, VocabParallelLinear, + set_autotp_mode) from deepspeed.module_inject.tp_plan_converter import TPPlanConverter +from deepspeed.sequence.cross_entropy import VocabParallelCausalLMLoss, configure_vocab_parallel_loss class SubAttn(nn.Module): @@ -160,6 +162,25 @@ def _build_gathered_lm_head_autotp(model, mp_size=1): return autotp +def _build_local_lm_head_autotp(model, vocab_parallel_lm_head=True): + config = AutoTPConfig(layer_specs=[ + TPLayerSpec(patterns=[r".*lm_head\.weight$"], partition_type=PartitionType.COLUMN), + ]) + autotp = AutoTP( + module=model, + all_reduce_linears=[], + prefix="", + state_dict=None, + linear_layer_setting=None, + orig_layer_impl=None, + partition_config=config, + vocab_parallel_lm_head=vocab_parallel_lm_head, + ) + autotp.set_tensor_parallel_config(1, None) + autotp.update_linear_policies() + return autotp + + def _build_legacy_lm_head_autotp(model, training_mode=False): autotp = AutoTP( module=model, @@ -247,6 +268,72 @@ def test_gathered_lm_head_uses_column_parallel_layer_when_output_dim_is_uneven() assert model.lm_head.gather_output +def test_vocab_parallel_linear_exposes_vocab_metadata(): + layer = VocabParallelLinear(nn.Linear(32, 101, bias=False), mp_group=None, name="lm_head") + + assert layer.vocab_size == 101 + assert layer.vocab_start_index == 0 + assert layer.vocab_end_index == 101 + assert not layer.gather_output + + +def test_plain_colwise_lm_head_uses_vocab_parallel_layer(): + model = OutputModel(tied=False) + + _build_local_lm_head_autotp(model, vocab_parallel_lm_head=True)._replace_module(model) + + assert isinstance(model.lm_head, VocabParallelLinear) + + +def test_plain_colwise_lm_head_without_flag_keeps_local_logits(): + model = OutputModel(tied=False) + + _build_local_lm_head_autotp(model, vocab_parallel_lm_head=False)._replace_module(model) + + assert isinstance(model.lm_head, LinearLayer) + assert not isinstance(model.lm_head, VocabParallelLinear) + assert not model.lm_head.gather_output + + +@pytest.mark.parametrize("name,expected", [ + ("lm_head", True), + ("embed_out", True), + ("model.lm_head", True), + ("lm_head_proj", False), + ("model.lm_head_proj.weight", False), + ("inner_lm_head.block", False), +]) +def test_lm_head_name_matching_ignores_projections(name, expected): + assert AutoTP._is_lm_head_name(name) is expected + + +def test_plain_colwise_lm_head_rejects_tied_weights(): + model = OutputModel(tied=True) + + with pytest.raises(ValueError, match="requires untied"): + _build_local_lm_head_autotp(model)._replace_module(model) + + +def test_configure_vocab_parallel_loss_installs_and_preserves_hook(): + model = OutputModel(tied=False) + model.loss_function = nn.CrossEntropyLoss() + original_loss_function = model.loss_function + _build_local_lm_head_autotp(model)._replace_module(model) + + configure_vocab_parallel_loss(model, model.lm_head) + + assert isinstance(model.loss_function, VocabParallelCausalLMLoss) + assert model._deepspeed_original_loss_function is original_loss_function + + +def test_configure_vocab_parallel_loss_requires_hook(): + model = OutputModel(tied=False) + _build_local_lm_head_autotp(model)._replace_module(model) + + with pytest.raises(ValueError, match="requires a writable loss_function"): + configure_vocab_parallel_loss(model, model.lm_head) + + @pytest.mark.parametrize("head", ["lm_head", "embed_out"]) def test_legacy_output_head_defaults_to_column_parallel_during_training(head): model = OutputModel(tied=False) diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py new file mode 100644 index 000000000000..858683645d2f --- /dev/null +++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py @@ -0,0 +1,292 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +import deepspeed.comm as dist +import deepspeed.sequence.cross_entropy as cross_entropy +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.layers import VocabParallelLinear +from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list +from deepspeed.sequence.cross_entropy import VocabParallelCausalLMLoss, vocab_parallel_cross_entropy +from unit.common import DistributedTest + + +@pytest.mark.parametrize("reduction", ["none", "sum", "mean"]) +def test_vocab_parallel_cross_entropy_tp1_matches_torch(reduction): + torch.manual_seed(42) + logits = torch.randn(2, 3, 17, requires_grad=True) + reference_logits = logits.detach().clone().requires_grad_(True) + target = torch.tensor([[0, 8, -100], [16, 7, 9]]) + + expected = F.cross_entropy(reference_logits.view(-1, 17), target.view(-1), reduction=reduction, ignore_index=-100) + if reduction == "none": + expected = expected.view_as(target) + actual = vocab_parallel_cross_entropy(logits, target, reduction=reduction) + + torch.testing.assert_close(actual, expected) + actual.sum().backward() + expected.sum().backward() + torch.testing.assert_close(logits.grad, reference_logits.grad) + + +def test_vocab_parallel_cross_entropy_all_ignored_is_zero(): + logits = torch.randn(2, 3, 11, requires_grad=True) + target = torch.full((2, 3), -100) + + loss = vocab_parallel_cross_entropy(logits, target) + + torch.testing.assert_close(loss, torch.zeros_like(loss)) + loss.backward() + torch.testing.assert_close(logits.grad, torch.zeros_like(logits)) + + +def test_vocab_parallel_cross_entropy_validates_inputs(): + logits = torch.randn(2, 3, 11) + target = torch.zeros(2, 3, dtype=torch.long) + + with pytest.raises(ValueError, match="matching non-vocabulary dimensions"): + vocab_parallel_cross_entropy(logits, target[:, :-1]) + with pytest.raises(ValueError, match="Unsupported reduction"): + vocab_parallel_cross_entropy(logits, target, reduction="batchmean") + with pytest.raises(ValueError, match="only supported with reduction='none'"): + vocab_parallel_cross_entropy(logits, target, reduction="mean", gather_sequence_loss=True) + with pytest.raises(ValueError, match="Vocabulary shard bounds"): + vocab_parallel_cross_entropy(logits, target, vocab_start_index=0, vocab_end_index=10) + target[0, 0] = 11 + with pytest.raises(ValueError, match="out of range"): + vocab_parallel_cross_entropy(logits, target) + + +def test_causal_lm_loss_shifts_labels_and_matches_reference(): + torch.manual_seed(7) + vocab_size = 13 + logits = torch.randn(2, 5, vocab_size) + labels = torch.tensor([[0, 3, 12, -100, 7], [1, 2, 3, 4, -100]]) + loss_fn = VocabParallelCausalLMLoss(vocab_start_index=0, vocab_end_index=vocab_size) + + from_labels = loss_fn(logits=logits, labels=labels) + from_shift_labels = loss_fn(logits=logits[..., :-1, :], shift_labels=labels[..., 1:]) + expected = F.cross_entropy(logits[..., :-1, :].reshape(-1, vocab_size), + labels[..., 1:].reshape(-1), + ignore_index=-100) + + torch.testing.assert_close(from_labels, expected) + torch.testing.assert_close(from_shift_labels, expected) + + +def test_causal_lm_loss_divides_sum_by_num_items_in_batch(): + torch.manual_seed(8) + vocab_size = 9 + logits = torch.randn(2, 4, vocab_size) + shift_labels = torch.tensor([[0, 1, -100], [2, 3, 4]]) + loss_fn = VocabParallelCausalLMLoss(vocab_start_index=0, vocab_end_index=vocab_size) + + loss = loss_fn(logits=logits[..., :-1, :], shift_labels=shift_labels, num_items_in_batch=5) + + token_losses = F.cross_entropy(logits[..., :-1, :].reshape(-1, vocab_size), + shift_labels.reshape(-1), + reduction="none", + ignore_index=-100) + torch.testing.assert_close(loss, token_losses.sum() / 5) + + +def test_causal_lm_loss_vocab_size_mismatch_warns_instead_of_raising(): + torch.manual_seed(9) + vocab_size = 7 + logits = torch.randn(2, 3, vocab_size) + labels = torch.tensor([[0, 1, 2], [3, 4, 6]]) + loss_fn = VocabParallelCausalLMLoss(vocab_start_index=0, vocab_end_index=vocab_size) + + # A resized embedding leaves the caller's vocab_size stale; the head's weights win. + loss = loss_fn(logits=logits, labels=labels, vocab_size=vocab_size + 1) + + expected = F.cross_entropy(logits[..., :-1, :].reshape(-1, vocab_size), + labels[..., 1:].reshape(-1), + ignore_index=-100) + torch.testing.assert_close(loss, expected) + + +def test_vocab_metadata_validation_runs_once(monkeypatch): + torch.manual_seed(10) + # A vocab size unique to this test keeps the process-wide metadata cache from being + # warmed by the other tests, whatever order pytest runs them in. + vocab_size = 23 + logits = torch.randn(2, 3, vocab_size) + target = torch.zeros(2, 3, dtype=torch.long) + calls = [] + original_validate = cross_entropy._validate_vocab_shard_bounds + + def counting_validate(*args, **kwargs): + calls.append(1) + return original_validate(*args, **kwargs) + + monkeypatch.setattr(cross_entropy, "_validate_vocab_shard_bounds", counting_validate) + vocab_parallel_cross_entropy(logits, target, vocab_start_index=0, vocab_end_index=vocab_size) + vocab_parallel_cross_entropy(logits, target, vocab_start_index=0, vocab_end_index=vocab_size) + + assert len(calls) == 1 + + +class TestVocabParallelCrossEntropyTP(DistributedTest): + world_size = 2 + + def test_uneven_vocab_matches_torch(self): + device = torch.device(get_accelerator().current_device_name()) + rank = dist.get_rank() + partition_sizes = get_shard_size_list(17, self.world_size, AutoTPMeta(), "lm_head") + vocab_start_index = sum(partition_sizes[:rank]) + vocab_end_index = vocab_start_index + partition_sizes[rank] + + torch.manual_seed(123) + full_logits = torch.randn(2, 3, 17, device=device) + reference_logits = full_logits.detach().clone().requires_grad_(True) + local_logits = full_logits[..., vocab_start_index:vocab_end_index].detach().clone().requires_grad_(True) + target = torch.tensor([[0, 8, -100], [16, 7, 9]], device=device) + + expected = F.cross_entropy(reference_logits.view(-1, 17), target.view(-1)) + actual = vocab_parallel_cross_entropy(local_logits, + target, + tp_group=dist.get_world_group(), + vocab_start_index=vocab_start_index, + vocab_end_index=vocab_end_index) + + torch.testing.assert_close(actual, expected) + actual.backward() + expected.backward() + torch.testing.assert_close(local_logits.grad, reference_logits.grad[..., vocab_start_index:vocab_end_index]) + + def test_rejects_non_contiguous_vocab_shards(self): + device = torch.device(get_accelerator().current_device_name()) + rank = dist.get_rank() + target = torch.zeros(2, 3, dtype=torch.long, device=device) + + gap_start, gap_end = ((0, 8), (9, 17))[rank] + gap_logits = torch.randn(2, 3, gap_end - gap_start, device=device) + with pytest.raises(ValueError, match="contiguous, non-overlapping"): + vocab_parallel_cross_entropy(gap_logits, + target, + tp_group=dist.get_world_group(), + vocab_start_index=gap_start, + vocab_end_index=gap_end) + + overlap_start, overlap_end = ((0, 9), (8, 17))[rank] + overlap_logits = torch.randn(2, 3, overlap_end - overlap_start, device=device) + with pytest.raises(ValueError, match="contiguous, non-overlapping"): + vocab_parallel_cross_entropy(overlap_logits, + target, + tp_group=dist.get_world_group(), + vocab_start_index=overlap_start, + vocab_end_index=overlap_end) + + +class TestVocabParallelCrossEntropySP(DistributedTest): + world_size = 2 + + def test_sequence_loss_local_gathered_and_mean(self): + device = torch.device(get_accelerator().current_device_name()) + rank = dist.get_rank() + local_sequence_size = 2 + + torch.manual_seed(456) + full_logits = torch.randn(4, 2, 13, device=device) + target = torch.tensor([[0, 1], [2, -100], [12, 3], [4, 5]], device=device) + sequence_start = rank * local_sequence_size + sequence_end = sequence_start + local_sequence_size + local_logits = full_logits[sequence_start:sequence_end].detach().clone().requires_grad_(True) + local_target = target[sequence_start:sequence_end] + + expected_none = F.cross_entropy(full_logits.view(-1, 13), target.view(-1), reduction="none").view_as(target) + actual_local = vocab_parallel_cross_entropy(local_logits, local_target, reduction="none") + actual_gathered = vocab_parallel_cross_entropy(local_logits, + local_target, + sp_group=dist.get_world_group(), + reduction="none", + gather_sequence_loss=True) + actual_mean = vocab_parallel_cross_entropy(local_logits, + local_target, + sp_group=dist.get_world_group(), + reduction="mean") + + torch.testing.assert_close(actual_local, expected_none[sequence_start:sequence_end]) + torch.testing.assert_close(actual_gathered, expected_none) + torch.testing.assert_close(actual_mean, expected_none.sum() / (target != -100).sum()) + + def test_gathered_sequence_loss_backward_accumulates_all_ranks(self): + device = torch.device(get_accelerator().current_device_name()) + rank = dist.get_rank() + + torch.manual_seed(654) + local_logits = torch.randn(2, 3, 13, device=device, requires_grad=True) + reference_logits = local_logits.detach().clone().requires_grad_(True) + local_target = torch.tensor([[0, 1, 2], [3, 4, 5]], device=device) + + gathered_loss = vocab_parallel_cross_entropy(local_logits, + local_target, + sp_group=dist.get_world_group(), + reduction="none", + gather_sequence_loss=True) + rank_weight = rank + 1 + (gathered_loss * rank_weight).sum().backward() + + total_weight = sum(range(1, self.world_size + 1)) + reference_loss = F.cross_entropy(reference_logits.view(-1, 13), local_target.view(-1), reduction="sum") + (reference_loss * total_weight).backward() + torch.testing.assert_close(local_logits.grad, reference_logits.grad) + + +class TestVocabParallelCrossEntropyTPAndSP(DistributedTest): + world_size = 4 + + def test_orthogonal_groups_match_torch(self): + device = torch.device(get_accelerator().current_device_name()) + rank = dist.get_rank() + tp_groups = [dist.new_group(ranks=[0, 1]), dist.new_group(ranks=[2, 3])] + sp_groups = [dist.new_group(ranks=[0, 2]), dist.new_group(ranks=[1, 3])] + tp_group = tp_groups[rank // 2] + sp_group = sp_groups[rank % 2] + tp_rank = rank % 2 + sp_rank = rank // 2 + + partition_sizes = get_shard_size_list(17, 2, AutoTPMeta(), "lm_head") + vocab_start_index = sum(partition_sizes[:tp_rank]) + vocab_end_index = vocab_start_index + partition_sizes[tp_rank] + sequence_start = sp_rank * 2 + sequence_end = sequence_start + 2 + + torch.manual_seed(789) + full_logits = torch.randn(4, 2, 17, device=device) + reference_logits = full_logits.detach().clone().requires_grad_(True) + local_logits = full_logits[sequence_start:sequence_end, ..., + vocab_start_index:vocab_end_index].detach().clone().requires_grad_(True) + target = torch.tensor([[0, 1], [8, -100], [16, 3], [4, 9]], device=device) + local_target = target[sequence_start:sequence_end] + + expected = F.cross_entropy(reference_logits.view(-1, 17), target.view(-1)) + actual = vocab_parallel_cross_entropy(local_logits, + local_target, + tp_group=tp_group, + sp_group=sp_group, + vocab_start_index=vocab_start_index, + vocab_end_index=vocab_end_index) + + torch.testing.assert_close(actual, expected) + actual.backward() + expected.backward() + expected_grad = reference_logits.grad[sequence_start:sequence_end, ..., vocab_start_index:vocab_end_index] + torch.testing.assert_close(local_logits.grad, expected_grad) + + +class TestVocabParallelLinearRejectsEmptyShard(DistributedTest): + world_size = 2 + + def test_vocabulary_smaller_than_tp_size_raises_on_all_ranks(self): + # Both ranks derive the same shard-size list, so the failure must be raised + # everywhere at construction time instead of hanging in a later collective. + with pytest.raises(ValueError, match="at least tp_size"): + VocabParallelLinear(nn.Linear(4, 1, bias=False), mp_group=dist.get_world_group(), name="lm_head") From bf8b7e298470f17cf7df583d9c71e241e453e012 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 8 Sep 2026 11:11:50 +0000 Subject: [PATCH 2/6] Fix vocab-parallel integration Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/module_inject/auto_tp.py | 4 +- deepspeed/sequence/cross_entropy.py | 19 ++++--- .../test_tp_partition_config_path.py | 53 +++++++++++++++++++ 3 files changed, 67 insertions(+), 9 deletions(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index e42a2afb9726..3d97bd9e0380 100755 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -571,6 +571,8 @@ def _configure_gathered_column_tie_fallbacks(self): for module_name, module in named_modules: if not module_name or isinstance(module, nn.Embedding) or not hasattr(module, "weight"): continue + if self._is_vocab_parallel_lm_head(module, module_name): + continue tied_embedding_name = next( (embedding_name for embedding_name, embedding in embeddings if module.weight is embedding.weight), @@ -757,7 +759,7 @@ def _replace_autoep_shared_experts(self, autoep_layer, autoep_name): self.update_mp_params(child, full_name) def _replace_module(self, r_module, prev_name='', prev_class_name=''): - if prev_name == '' and prev_class_name == '' and not self.vocab_parallel_lm_head: + if prev_name == '' and prev_class_name == '': self._configure_gathered_column_tie_fallbacks() for name, child in r_module.named_children(): diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py index f9275e84aab7..11422a1bf845 100644 --- a/deepspeed/sequence/cross_entropy.py +++ b/deepspeed/sequence/cross_entropy.py @@ -264,7 +264,7 @@ def forward(self, vocab_parallel_logits, target): gather_sequence_loss=self.gather_sequence_loss) -class VocabParallelCausalLMLoss(nn.Module): +class VocabParallelCausalLMLoss: """Distributed causal-LM loss for a vocabulary-sharded (no-gather) LM head. ``sp_group`` must stay ``None`` under DeepSpeed's Ulysses sequence-parallel engine: @@ -275,14 +275,13 @@ class VocabParallelCausalLMLoss(nn.Module): """ def __init__(self, tp_group=None, sp_group=None, vocab_start_index=None, vocab_end_index=None, ignore_index=-100): - super().__init__() self.tp_group = tp_group self.sp_group = sp_group self.vocab_start_index = vocab_start_index self.vocab_end_index = vocab_end_index self.ignore_index = ignore_index - def forward(self, logits, labels=None, vocab_size=None, shift_labels=None, num_items_in_batch=None, **kwargs): + def __call__(self, logits, labels=None, vocab_size=None, shift_labels=None, num_items_in_batch=None, **kwargs): if shift_labels is None: if labels is None: raise ValueError("labels or shift_labels must be provided") @@ -332,18 +331,22 @@ def configure_vocab_parallel_loss(model, vocab_parallel_head, sp_group=None, ign vocab_start_index=vocab_parallel_head.vocab_start_index, vocab_end_index=vocab_parallel_head.vocab_end_index, ignore_index=ignore_index) - # Keep the stock loss reachable so callers can restore it when tearing the head down. - if not hasattr(model, "_deepspeed_original_loss_function"): - model._deepspeed_original_loss_function = model.loss_function + original_loss_function = model.loss_function + registered_loss_module = getattr(model, "_modules", {}).pop("loss_function", None) # Some model classes expose loss_function as a read-only property, in which case the # assignment raises; the identity check below turns that into an actionable error # instead of leaving the model silently computing loss on rank-local logits. try: model.loss_function = loss_fn - except AttributeError: - pass + except (AttributeError, TypeError): + if registered_loss_module is not None: + model.add_module("loss_function", registered_loss_module) if model.loss_function is not loss_fn: raise ValueError("Unable to install the vocab-parallel loss_function hook; use gather_output=True") + + # Keep the stock loss reachable so callers can restore it when tearing the head down. + if not hasattr(model, "_deepspeed_original_loss_function"): + model._deepspeed_original_loss_function = original_loss_function return model diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py index 9bd96ebbc1fe..d666347271de 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -10,6 +10,7 @@ import pytest import torch.nn as nn +from transformers import PreTrainedModel, PretrainedConfig from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, VocabParallelLinear, @@ -56,6 +57,15 @@ def __init__(self, tied): self.lm_head.weight = self.embed_tokens.weight +class HFOutputModel(PreTrainedModel): + config_class = PretrainedConfig + + def __init__(self): + super().__init__(PretrainedConfig()) + self.embed_tokens = nn.Embedding(100, 32) + self.lm_head = nn.Linear(32, 100, bias=False) + + def _build_config(): """Partition config that matches q_proj and o_proj via regex.""" return AutoTPConfig(layer_specs=[ @@ -326,6 +336,18 @@ def test_configure_vocab_parallel_loss_installs_and_preserves_hook(): assert model._deepspeed_original_loss_function is original_loss_function +def test_configure_vocab_parallel_loss_installs_on_huggingface_model(): + model = HFOutputModel() + original_loss_function = model.loss_function + _build_local_lm_head_autotp(model)._replace_module(model) + + configure_vocab_parallel_loss(model, model.lm_head) + + assert isinstance(model.loss_function, VocabParallelCausalLMLoss) + assert not isinstance(model.loss_function, nn.Module) + assert model._deepspeed_original_loss_function is original_loss_function + + def test_configure_vocab_parallel_loss_requires_hook(): model = OutputModel(tied=False) _build_local_lm_head_autotp(model)._replace_module(model) @@ -378,6 +400,37 @@ def test_legacy_tied_lm_head_stays_replicated_during_training(): assert model.lm_head.weight is tied_weight +def test_vocab_parallel_flag_preserves_nonstandard_tied_output_fallback(): + model = OutputModel(tied=False) + model.output_proj = model.lm_head + model.output_proj.weight = model.embed_tokens.weight + del model.lm_head + tied_weight = model.embed_tokens.weight + specs = TPPlanConverter.convert({ + "embed_tokens": "embedding_rowwise", + "output_proj": "colwise_gather_output", + }) + autotp = AutoTP( + module=model, + all_reduce_linears=[], + prefix="", + state_dict=None, + linear_layer_setting=None, + orig_layer_impl=None, + partition_config=AutoTPConfig(layer_specs=specs), + vocab_parallel_lm_head=True, + ) + autotp.set_tensor_parallel_config(2, None) + autotp.update_linear_policies() + + autotp._replace_module(model) + + assert isinstance(model.embed_tokens, nn.Embedding) + assert isinstance(model.output_proj, nn.Linear) + assert model.embed_tokens.weight is tied_weight + assert model.output_proj.weight is tied_weight + + def test_explicit_row_parallel_lm_head_is_not_overridden_by_its_name(): model = OutputModel(tied=False) _build_row_output_head_autotp(model, training_mode=True)._replace_module(model) From 8f03c27ea00dd74b651e28f6c7a688ea60138ca8 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 8 Sep 2026 11:35:32 +0000 Subject: [PATCH 3/6] Preserve sequence loss gradients Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/sequence/cross_entropy.py | 29 +++++++++++-------- .../test_vocab_parallel_cross_entropy.py | 18 +++++++++++- 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py index 11422a1bf845..326e6a47e0a8 100644 --- a/deepspeed/sequence/cross_entropy.py +++ b/deepspeed/sequence/cross_entropy.py @@ -59,9 +59,11 @@ def backward(ctx, grad_output): class _GatherSequenceLoss(torch.autograd.Function): @staticmethod - def forward(ctx, local_loss, sp_group): + def forward(ctx, local_loss, sp_group, sum_gradients): ctx.sp_group = sp_group ctx.local_sequence_size = local_loss.shape[0] + ctx.sum_gradients = sum_gradients + ctx.sp_rank = dist.get_rank(sp_group) output_shape = (ctx.local_sequence_size * dist.get_world_size(sp_group), *local_loss.shape[1:]) gathered_loss = torch.empty(output_shape, dtype=local_loss.dtype, device=local_loss.device) @@ -70,11 +72,15 @@ def forward(ctx, local_loss, sp_group): @staticmethod def backward(ctx, grad_output): + if not ctx.sum_gradients: + start = ctx.sp_rank * ctx.local_sequence_size + return grad_output.narrow(0, start, ctx.local_sequence_size), None, None + grad_input = torch.empty((ctx.local_sequence_size, *grad_output.shape[1:]), dtype=grad_output.dtype, device=grad_output.device) dist.reduce_scatter_fn(grad_input, grad_output.contiguous(), group=ctx.sp_group) - return grad_input, None + return grad_input, None, None def _global_sp_sum(local_value, sp_group): @@ -193,7 +199,7 @@ def vocab_parallel_cross_entropy(vocab_parallel_logits, if gather_sequence_loss: if sp_group is None: raise ValueError("sp_group is required when gather_sequence_loss=True") - loss = _GatherSequenceLoss.apply(loss, sp_group) + loss = _GatherSequenceLoss.apply(loss, sp_group, True) return loss loss_sum = _global_sp_sum(loss.sum(), sp_group) @@ -215,22 +221,21 @@ def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, ignore_index=-100, reduction="none", gather_sequence_loss=True): - """Sequence-parallel wrapper over :func:`vocab_parallel_cross_entropy`. + """Sequence-parallel wrapper preserving the legacy local-slice gradient.""" + if gather_sequence_loss and reduction != "none": + raise ValueError("gather_sequence_loss is only supported with reduction='none'") - Backward reduce-scatters the gradient over ``sp_group``, so each rank receives the - gradient of its own sequence shard with the other ranks' contributions already - summed in. Downstream code must therefore treat the returned loss as replicated - across the SP group and must not average SP gradients a second time. - """ - return vocab_parallel_cross_entropy(vocab_parallel_logits, + loss = vocab_parallel_cross_entropy(vocab_parallel_logits, target, tp_group=tp_group, sp_group=sp_group, vocab_start_index=vocab_start_index, vocab_end_index=vocab_end_index, ignore_index=ignore_index, - reduction=reduction, - gather_sequence_loss=gather_sequence_loss) + reduction=reduction) + if not gather_sequence_loss: + return loss + return _GatherSequenceLoss.apply(loss, sp_group, False) class VocabParallelCrossEntropyLoss(nn.Module): diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py index 858683645d2f..44ca1a4506c6 100644 --- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py +++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py @@ -13,7 +13,8 @@ from deepspeed.accelerator import get_accelerator from deepspeed.module_inject.layers import VocabParallelLinear from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list -from deepspeed.sequence.cross_entropy import VocabParallelCausalLMLoss, vocab_parallel_cross_entropy +from deepspeed.sequence.cross_entropy import (VocabParallelCausalLMLoss, vocab_parallel_cross_entropy, + vocab_sequence_parallel_cross_entropy) from unit.common import DistributedTest @@ -239,6 +240,21 @@ def test_gathered_sequence_loss_backward_accumulates_all_ranks(self): (reference_loss * total_weight).backward() torch.testing.assert_close(local_logits.grad, reference_logits.grad) + def test_legacy_sequence_loss_backward_keeps_local_gradient_scale(self): + device = torch.device(get_accelerator().current_device_name()) + local_logits = torch.randn(2, 3, 13, device=device, requires_grad=True) + reference_logits = local_logits.detach().clone().requires_grad_(True) + local_target = torch.tensor([[0, 1, 2], [3, 4, 5]], device=device) + + gathered_loss = vocab_sequence_parallel_cross_entropy(local_logits, + local_target, + sp_group=dist.get_world_group()) + gathered_loss.sum().backward() + + reference_loss = F.cross_entropy(reference_logits.view(-1, 13), local_target.view(-1), reduction="sum") + reference_loss.backward() + torch.testing.assert_close(local_logits.grad, reference_logits.grad) + class TestVocabParallelCrossEntropyTPAndSP(DistributedTest): world_size = 4 From 97078ce926276740509da43fe434fad683f7ea34 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 8 Sep 2026 11:54:53 +0000 Subject: [PATCH 4/6] Document sequence loss gradients Require an explicit sequence-parallel group when the legacy wrapper gathers losses, and document the distinct backward conventions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/sequence/cross_entropy.py | 2 ++ docs/_tutorials/autotp-training.md | 10 ++++++++++ .../test_vocab_parallel_cross_entropy.py | 2 ++ 3 files changed, 14 insertions(+) diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py index 326e6a47e0a8..d035f0bfc1b6 100644 --- a/deepspeed/sequence/cross_entropy.py +++ b/deepspeed/sequence/cross_entropy.py @@ -224,6 +224,8 @@ def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, """Sequence-parallel wrapper preserving the legacy local-slice gradient.""" if gather_sequence_loss and reduction != "none": raise ValueError("gather_sequence_loss is only supported with reduction='none'") + if gather_sequence_loss and sp_group is None: + raise ValueError("sp_group is required when gather_sequence_loss=True") loss = vocab_parallel_cross_entropy(vocab_parallel_logits, target, diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md index 19a3839691b3..d3d00eb3666a 100644 --- a/docs/_tutorials/autotp-training.md +++ b/docs/_tutorials/autotp-training.md @@ -182,6 +182,16 @@ the supplied SP group. TP and SP may be combined with explicit orthogonal process groups, but AutoTP does not currently construct a combined TP x SP mesh automatically. +Gathered sequence losses have two backward conventions. The general +`vocab_parallel_cross_entropy(..., gather_sequence_loss=True)` API sums the +gradient contributions from every SP rank before returning each rank's local +slice. The compatibility wrapper `vocab_sequence_parallel_cross_entropy` +preserves its legacy behavior and returns only the corresponding local slice of +the gathered loss gradient. New callers that consume or reduce the gathered +loss on every SP rank should use the general API; existing callers can retain +the wrapper without changing their gradient scale. Both gathered forms require +an explicit `sp_group`. + Under DeepSpeed's Ulysses sequence-parallel engine, the installed loss must keep its default sequence-parallel settings (no `sp_group`): the engine aggregates each shard's mean itself, weighted by the shard's valid-token count, so an diff --git a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py index 44ca1a4506c6..50cd1cfab49b 100644 --- a/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py +++ b/tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py @@ -57,6 +57,8 @@ def test_vocab_parallel_cross_entropy_validates_inputs(): vocab_parallel_cross_entropy(logits, target, reduction="batchmean") with pytest.raises(ValueError, match="only supported with reduction='none'"): vocab_parallel_cross_entropy(logits, target, reduction="mean", gather_sequence_loss=True) + with pytest.raises(ValueError, match="sp_group is required"): + vocab_sequence_parallel_cross_entropy(logits, target, sp_group=None) with pytest.raises(ValueError, match="Vocabulary shard bounds"): vocab_parallel_cross_entropy(logits, target, vocab_start_index=0, vocab_end_index=10) target[0, 0] = 11 From a3b4405236c60a7830b6889d328032a83516e2b8 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 8 Sep 2026 12:03:02 +0000 Subject: [PATCH 5/6] Fix vocab-parallel heuristic integration Replace the output head in heuristic AutoTP and preserve original tied-weight evidence before embedding traversal mutates parameters. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/module_inject/auto_tp.py | 28 ++++++++++ deepspeed/runtime/engine.py | 17 ++++++ .../test_tp_partition_config_path.py | 23 ++++++++ tests/unit/v1/autotp/test_autotp_training.py | 56 ++++++++++++++++++- 4 files changed, 123 insertions(+), 1 deletion(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 3d97bd9e0380..983f619ac9f3 100755 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -228,6 +228,15 @@ def __init__(self, self.partition_config = partition_config self.vocab_parallel_lm_head = vocab_parallel_lm_head self.training_mode = training_mode + embedding_weights = { + id(module.weight) + for module in self.module.modules() if isinstance(module, nn.Embedding) and hasattr(module, "weight") + } + self._originally_tied_vocab_head_ids = { + id(module) + for name, module in self.module.named_modules() + if self._is_vocab_parallel_lm_head(module, name) and id(module.weight) in embedding_weights + } self._gathered_column_tie_fallbacks_configured = False self._tied_gathered_column_module_names = set() TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host) @@ -537,10 +546,29 @@ def _create_vocab_parallel_layer(self, child, name): return VocabParallelLinear(child, self.mp_group, name=name, tp_meta=self.tp_meta) def _validate_untied_vocab_head(self, lm_head): + if id(lm_head) in self._originally_tied_vocab_head_ids: + raise ValueError("A no-gather vocab-parallel LM head requires untied embedding and output weights") for _, module in self.module.named_modules(): if isinstance(module, nn.Embedding) and getattr(module, "weight", None) is lm_head.weight: raise ValueError("A no-gather vocab-parallel LM head requires untied embedding and output weights") + def replace_vocab_parallel_lm_head(self): + candidates = [] + for parent_name, parent in self.module.named_modules(): + for child_name, child in parent.named_children(): + full_name = f"{parent_name}.{child_name}" if parent_name else child_name + if self._is_vocab_parallel_lm_head(child, full_name): + candidates.append((parent, child_name, child, full_name)) + + if not candidates: + raise ValueError("vocab_parallel_lm_head requires a supported nn.Linear named 'lm_head' or 'embed_out'") + if len(candidates) > 1: + names = [full_name for _, _, _, full_name in candidates] + raise ValueError(f"Unable to choose among multiple vocab-parallel LM heads: {names}") + + parent, child_name, child, full_name = candidates[0] + setattr(parent, child_name, self._create_vocab_parallel_layer(child, full_name)) + def _configure_gathered_column_tie_fallbacks(self): """Configure a replicated fallback for gathered output layers tied to embeddings.""" if self._gathered_column_tie_fallbacks_configured: diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 888675afece7..4e5caa462b83 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -930,11 +930,28 @@ def finalize_autotp(autotp=None, attach_uc_metadata=False): log_dist("AutoTP: no effective HuggingFace tp_plan was found; falling back to heuristic AutoTP.", ranks=[0]) + vocab_head_autotp = None + if tp_config.vocab_parallel_lm_head: + vocab_head_autotp = AutoTP(module=model, + all_reduce_linears=(), + prefix="", + state_dict=None, + linear_layer_setting=(torch.nn.Linear, torch.nn.Embedding), + orig_layer_impl=None, + keep_module_on_host=tp_config.keep_module_on_host, + vocab_parallel_lm_head=True, + model_config=model_config, + tp_grain_size=tp_config.tensor_parallel.tp_grain_size, + training_mode=True) + vocab_head_autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group) + parser_dict = AutoTP.tp_parser(model) for client_module, injection_policy in parser_dict: tp_config.injection_policy_tuple = injection_policy replace_transformer_layer(client_module, model, None, tp_config, model_config, training_mode=True) + if vocab_head_autotp is not None: + vocab_head_autotp.replace_vocab_parallel_lm_head() finalize_autotp(attach_uc_metadata=True) def __del__(self): diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py index d666347271de..ee6f19d55fe4 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -324,6 +324,29 @@ def test_plain_colwise_lm_head_rejects_tied_weights(): _build_local_lm_head_autotp(model)._replace_module(model) +def test_plain_colwise_lm_head_rejects_tie_before_embedding_is_sliced(): + model = OutputModel(tied=True) + specs = TPPlanConverter.convert({ + "embed_tokens": "embedding_rowwise", + "lm_head": "colwise", + }) + autotp = AutoTP( + module=model, + all_reduce_linears=[], + prefix="", + state_dict=None, + linear_layer_setting=None, + orig_layer_impl=None, + partition_config=AutoTPConfig(layer_specs=specs), + vocab_parallel_lm_head=True, + ) + autotp.set_tensor_parallel_config(2, None) + autotp.update_linear_policies() + + with pytest.raises(ValueError, match="requires untied"): + autotp._replace_module(model) + + def test_configure_vocab_parallel_loss_installs_and_preserves_hook(): model = OutputModel(tied=False) model.loss_function = nn.CrossEntropyLoss() diff --git a/tests/unit/v1/autotp/test_autotp_training.py b/tests/unit/v1/autotp/test_autotp_training.py index 1099951e73c4..64658cd23f7a 100644 --- a/tests/unit/v1/autotp/test_autotp_training.py +++ b/tests/unit/v1/autotp/test_autotp_training.py @@ -17,7 +17,8 @@ from contextlib import contextmanager from torch import nn from deepspeed.module_inject.auto_tp import AutoTP -from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, set_autotp_mode, is_autotp_training_mode +from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, VocabParallelLinear, set_autotp_mode, + is_autotp_training_mode) from deepspeed.module_inject.tp_shard import get_shard_size_list from unit.checkpoint.common import compare_lr_scheduler_states, compare_optimizer_states import os @@ -150,6 +151,59 @@ def forward(self, x): return self.lm_head(x) +@pytest.mark.sequential +class TestHeuristicVocabParallelLMHead(DistributedTest): + world_size = 2 + reuse_dist_env = False + + def test_training_path_replaces_head_and_uses_distributed_loss(self): + transformers = pytest.importorskip("transformers") + + class HeuristicLlamaForCausalLM(transformers.LlamaForCausalLM): + _tp_plan = None + + model_config = transformers.LlamaConfig(vocab_size=33, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + use_cache=False, + tie_word_embeddings=False) + model_config.base_model_tp_plan = None + model = HeuristicLlamaForCausalLM(model_config) + ds_config = { + "train_micro_batch_size_per_gpu": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-6, + "torch_adam": True, + }, + }, + "tensor_parallel": { + "autotp_size": self.world_size, + "vocab_parallel_lm_head": True, + }, + "zero_optimization": { + "stage": 0, + }, + } + + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config) + + assert isinstance(engine.module.lm_head, VocabParallelLinear) + device = torch.device(get_accelerator().current_device_name()) + input_ids = torch.randint(0, model_config.vocab_size, (1, 8), device=device) + dist.broadcast(input_ids, + src=groups.get_tensor_model_parallel_src_rank(), + group=groups.get_tensor_model_parallel_group()) + output = engine(input_ids=input_ids, labels=input_ids) + assert torch.isfinite(output.loss) + engine.backward(output.loss) + assert engine.module.lm_head.weight.grad is not None + + @contextmanager def should_assert_with_msg(expected_message): try: From a47d7abf5648f309eb88fc05683ec4db3c373cd2 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Tue, 8 Sep 2026 12:11:30 +0000 Subject: [PATCH 6/6] Validate vocab head before AutoTP Resolve and validate the heuristic output head before transformer partitioning, cache the candidate for replacement, and skip tied-weight scans when the feature is disabled. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- deepspeed/module_inject/auto_tp.py | 35 ++++++++++++++++++++---------- deepspeed/runtime/engine.py | 3 ++- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 983f619ac9f3..021119d6e109 100755 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -228,15 +228,18 @@ def __init__(self, self.partition_config = partition_config self.vocab_parallel_lm_head = vocab_parallel_lm_head self.training_mode = training_mode - embedding_weights = { - id(module.weight) - for module in self.module.modules() if isinstance(module, nn.Embedding) and hasattr(module, "weight") - } - self._originally_tied_vocab_head_ids = { - id(module) - for name, module in self.module.named_modules() - if self._is_vocab_parallel_lm_head(module, name) and id(module.weight) in embedding_weights - } + self._originally_tied_vocab_head_ids = set() + self._vocab_parallel_lm_head_candidate = None + if self.vocab_parallel_lm_head: + embedding_weights = { + id(module.weight) + for module in self.module.modules() if isinstance(module, nn.Embedding) and hasattr(module, "weight") + } + self._originally_tied_vocab_head_ids = { + id(module) + for name, module in self.module.named_modules() + if self._is_vocab_parallel_lm_head(module, name) and id(module.weight) in embedding_weights + } self._gathered_column_tie_fallbacks_configured = False self._tied_gathered_column_module_names = set() TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host) @@ -552,7 +555,10 @@ def _validate_untied_vocab_head(self, lm_head): if isinstance(module, nn.Embedding) and getattr(module, "weight", None) is lm_head.weight: raise ValueError("A no-gather vocab-parallel LM head requires untied embedding and output weights") - def replace_vocab_parallel_lm_head(self): + def _resolve_vocab_parallel_lm_head(self): + if self._vocab_parallel_lm_head_candidate is not None: + return self._vocab_parallel_lm_head_candidate + candidates = [] for parent_name, parent in self.module.named_modules(): for child_name, child in parent.named_children(): @@ -566,7 +572,14 @@ def replace_vocab_parallel_lm_head(self): names = [full_name for _, _, _, full_name in candidates] raise ValueError(f"Unable to choose among multiple vocab-parallel LM heads: {names}") - parent, child_name, child, full_name = candidates[0] + self._validate_untied_vocab_head(candidates[0][2]) + self._vocab_parallel_lm_head_candidate = candidates[0] + return self._vocab_parallel_lm_head_candidate + + def _replace_vocab_parallel_lm_head(self): + parent, child_name, child, full_name = self._resolve_vocab_parallel_lm_head() + if getattr(parent, child_name) is not child: + raise RuntimeError(f"Vocab-parallel LM head '{full_name}' changed during AutoTP partitioning") setattr(parent, child_name, self._create_vocab_parallel_layer(child, full_name)) def _configure_gathered_column_tie_fallbacks(self): diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 4e5caa462b83..82d0e1a9a059 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -944,6 +944,7 @@ def finalize_autotp(autotp=None, attach_uc_metadata=False): tp_grain_size=tp_config.tensor_parallel.tp_grain_size, training_mode=True) vocab_head_autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group) + vocab_head_autotp._resolve_vocab_parallel_lm_head() parser_dict = AutoTP.tp_parser(model) for client_module, injection_policy in parser_dict: @@ -951,7 +952,7 @@ def finalize_autotp(autotp=None, attach_uc_metadata=False): replace_transformer_layer(client_module, model, None, tp_config, model_config, training_mode=True) if vocab_head_autotp is not None: - vocab_head_autotp.replace_vocab_parallel_lm_head() + vocab_head_autotp._replace_vocab_parallel_lm_head() finalize_autotp(attach_uc_metadata=True) def __del__(self):