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

deepspeed/module_inject/auto_tp.py

@@ -205,7 +205,8 @@ def __init__(self,
                  linear_layer_setting,
                  orig_layer_impl,
                  keep_module_on_host=False,
-                 partition_config: Optional[AutoTPConfig] = None):
+                 partition_config: Optional[AutoTPConfig] = None,
+                 vocab_parallel_lm_head=False):
         self.module = module
         self.all_reduce_linears = all_reduce_linears
         self.prefix = prefix
@@ -218,6 +219,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._gathered_column_tie_fallbacks_configured = False
         self._tied_gathered_column_module_names = set()
         TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host)
@@ -364,6 +366,11 @@ def _replace(self, child, name, conv_linear_layer):
         if getattr(child, "_is_autoep_layer", False):
             return child
 
+        if self.vocab_parallel_lm_head and self._is_lm_head_name(name):
+            self._validate_untied_vocab_head(child)
+            setattr(child, "replaced", True)
+            return VocabParallelLinear(child, self.mp_group, name=name)
+
         weight_shape = child.weight.shape
         mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
 
@@ -423,6 +430,11 @@ def _replace_with_config(self, child, name):
         if getattr(child, "replaced", False) == True:
             return child
 
+        if self.vocab_parallel_lm_head and self._is_lm_head_name(name):
+            self._validate_untied_vocab_head(child)
+            setattr(child, "replaced", True)
+            return VocabParallelLinear(child, self.mp_group, name=name)
+
         # Build the full parameter name for pattern matching
         param_name = name + ".weight" if not name.endswith(".weight") else name
 
@@ -469,6 +481,9 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create column-parallel layer (AllReduce in backward)."""
         if self.conv_linear_layer:
             return conv_LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+        if self._is_lm_head_name(name) and not spec.gather_output:
+            self._validate_untied_vocab_head(module)
+            return VocabParallelLinear(module, self.mp_group, name=name)
         # 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):
             # Check and handle fused qkv for TP
@@ -485,6 +500,15 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
             )
         return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
 
+    @staticmethod
+    def _is_lm_head_name(name):
+        return any(part in ("lm_head", "embed_out") for part in str(name).split('.'))
+
+    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 or self.partition_config is None:
@@ -687,7 +711,7 @@ def _replace_autoep_shared_experts(self, autoep_layer, autoep_name):
                 self._replace_module(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():

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
@@ -939,6 +939,17 @@ 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)
+        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

@@ -286,7 +286,8 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
                          linear_layer_setting,
                          orig_layer_impl,
                          config.keep_module_on_host,
-                         partition_config=partition_config)
+                         partition_config=partition_config,
+                         vocab_parallel_lm_head=getattr(config, "vocab_parallel_lm_head", False))
 
         # 2. Set the tensor parallelism config
         _autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group)

deepspeed/runtime/engine.py

@@ -791,6 +791,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=(),
@@ -799,13 +815,12 @@ def _apply_autotp_partitioning(self, model, tp_config):
                             linear_layer_setting=(torch.nn.Linear, torch.nn.Embedding),
                             orig_layer_impl=None,
                             keep_module_on_host=tp_config.keep_module_on_host,
-                            partition_config=partition_config)
+                            partition_config=partition_config,
+                            vocab_parallel_lm_head=tp_config.vocab_parallel_lm_head)
             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:
@@ -838,13 +853,12 @@ 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,
                 )
                 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. "
@@ -860,8 +874,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)
 
-        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:

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)