[AutoTP] Add support for vocab-parallel LM head and related configurations

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,7 +226,20 @@ 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._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)
@@ -372,6 +386,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 +453,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 +512,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 +528,60 @@ 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):
+        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 _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():
+                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}")
+
+        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):
         """Configure a replicated fallback for gathered output layers tied to embeddings."""
         if self._gathered_column_tie_fallbacks_configured:
@@ -538,6 +612,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),

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.
 

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)

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,16 +911,15 @@ 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,
                 )
                 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. "
@@ -916,13 +930,30 @@ def _apply_autotp_partitioning(self, model, tp_config):
             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)
+            vocab_head_autotp._resolve_vocab_parallel_lm_head()
+
         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)
 
-        setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model))
-        setattr(model, "ds_autotp_parsed", 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):
         try:

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

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)