[AutoTP] Enable HF colwise_gather_output to support lm_head replace

deepspeed/module_inject/auto_tp.py

@@ -16,6 +16,7 @@
 from .fusedqkv_utils import require_tp_fused_qkvw
 from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
 from deepspeed.utils import groups
+from deepspeed.utils.logging import print_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
 from .autotp_config import TPLayerSpec, AutoTPConfig, PartitionType
 
@@ -214,6 +215,8 @@ def __init__(self,
         self.linear_policies = None
         self.conv_linear_layer = False
         self.partition_config = partition_config
+        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)
 
     def in_module_list(module, module_list):
@@ -424,6 +427,17 @@ def _replace_with_config(self, child, name):
         model_type = self._get_model_type()
         spec = self.partition_config.find_matching_spec(param_name, model_type)
 
+        if any(part in ("lm_head", "embed_out") for part in name.split('.')):
+            spec_details = None
+            if spec is not None:
+                spec_details = {
+                    "patterns": spec.patterns,
+                    "partition_type": spec.partition_type.value,
+                    "gather_output": spec.gather_output,
+                }
+            print_dist(f"AutoTP lm_head spec match: parameter={param_name!r}; matched_spec={spec_details!r}",
+                       ranks=[0])
+
         if spec is None:
             # No matching spec found
             if self.partition_config.strict_mode:
@@ -461,21 +475,88 @@ def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
 
     def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create column-parallel layer (AllReduce in backward)."""
+        if spec.gather_output and self.mp_size is not None and self.mp_size > 1:
+            output_dim = module.weight.shape[0]
+            if output_dim % self.mp_size != 0:
+                if any(part in ("lm_head", "embed_out") for part in name.split('.')):
+                    print_dist(
+                        f"AutoTP: '{name}' uses gather_output with uneven output dim {output_dim} and tp_size="
+                        f"{self.mp_size}; falling back to legacy LmHeadLinearAllreduce for checkpoint-safe "
+                        "consolidation.",
+                        ranks=[0],
+                    )
+                    return LmHeadLinearAllreduce(module, self.mp_group)
+                raise NotImplementedError(
+                    f"AutoTP gather_output requires output dimension divisible by tp_size. Layer '{name}' has "
+                    f"output dim {output_dim} with tp_size={self.mp_size}.")
+
         if self.conv_linear_layer:
-            return conv_LinearLayer(module, self.mp_group, name=name)
+            return conv_LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
         # 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
             return fused_LinearLayer(module, self.mp_group, fused_module=self.module)
         if spec.shape is not None:
+            if spec.gather_output:
+                raise NotImplementedError("AutoTP gather_output does not yet support shaped sub-parameter layers.")
             return SubParamLinearLayer(
                 module,
                 self.mp_group,
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
             )
-        return LinearLayer(module, self.mp_group, name=name)
+        if spec.gather_output:
+            print_dist(f"AutoTP: replacing '{name}' with LinearLayer(gather_output=True)", ranks=[0])
+        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+
+    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:
+            return
+
+        named_modules = list(self.module.named_modules())
+        embeddings = [(name, module) for name, module in named_modules
+                      if isinstance(module, nn.Embedding) and hasattr(module, "weight")]
+        get_input_embeddings = getattr(self.module, "get_input_embeddings", None)
+        if callable(get_input_embeddings):
+            input_embedding = get_input_embeddings()
+            if input_embedding is not None and hasattr(input_embedding, "weight"):
+                input_embedding_name = next(
+                    (name for name, module in named_modules if module is input_embedding),
+                    None,
+                )
+                is_known_embedding = any(module is input_embedding for _, module in embeddings)
+                if input_embedding_name is not None and not is_known_embedding:
+                    embeddings.append((input_embedding_name, input_embedding))
+        if not embeddings:
+            self._gathered_column_tie_fallbacks_configured = True
+            return
+
+        model_type = self._get_model_type()
+        for module_name, module in named_modules:
+            if not module_name or isinstance(module, nn.Embedding) or not hasattr(module, "weight"):
+                continue
+
+            tied_embedding_name = next(
+                (embedding_name for embedding_name, embedding in embeddings if module.weight is embedding.weight),
+                None,
+            )
+            if tied_embedding_name is None:
+                continue
+
+            spec = self.partition_config.find_matching_spec(module_name + ".weight", model_type)
+            if spec is None or spec.partition_type != PartitionType.COLUMN or not spec.gather_output:
+                continue
+
+            self._tied_gathered_column_module_names.update((module_name, tied_embedding_name))
+            print_dist(
+                f"AutoTP: '{module_name}.weight' is tied to '{tied_embedding_name}.weight'; leaving both modules "
+                "replicated because coupled vocabulary-parallel embedding is not supported yet.",
+                ranks=[0],
+            )
+
+        self._gathered_column_tie_fallbacks_configured = True
 
     def _get_model_type(self) -> Optional[str]:
         """Extract model type from module config or class name."""
@@ -570,6 +651,9 @@ 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 == '':
+            self._configure_gathered_column_tie_fallbacks()
+
         for name, child in r_module.named_children():
             if getattr(child, "_is_autoep_layer", False):
                 full_name = prev_name + '.' + name if prev_name else name
@@ -595,7 +679,9 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
             # instead of linear_policies. This keeps all pattern logic centralized here.
             if self.partition_config is not None:
                 full_name = class_name + '.' + name if class_name else name
-                if isinstance(child, nn.Embedding):
+                if full_name in self._tied_gathered_column_module_names:
+                    continue
+                elif isinstance(child, nn.Embedding):
                     # Check if embedding matches any pattern
                     param_name = full_name + ".weight"
                     model_type = self._get_model_type()

deepspeed/module_inject/autotp_config.py

@@ -49,6 +49,13 @@ class TPLayerSpec:
             partition_type=PartitionType.COLUMN,
         )
 
+        # Column-parallel layer with replicated output (e.g., an untied LM head)
+        TPLayerSpec(
+            patterns=[".*lm_head\\.weight$"],
+            partition_type=PartitionType.COLUMN,
+            gather_output=True,
+        )
+
         # Fused QKV - GLM style [Q, K, V] concatenated on dim 0
         TPLayerSpec(
             patterns=[".*\\.query_key_value\\.weight$"],
@@ -117,6 +124,9 @@ class TPLayerSpec:
     # Optional: model type constraint (only apply for specific models)
     model_types: Optional[List[str]] = None
 
+    # Gather column-parallel output shards so every TP rank receives the full output
+    gather_output: bool = False
+
     def __post_init__(self):
         if isinstance(self.partition_type, str):
             self.partition_type = PartitionType(self.partition_type.lower())
@@ -285,6 +295,7 @@ def from_dict(cls, config_dict: dict) -> "AutoTPConfig":
                 TPLayerSpec(
                     patterns=spec_dict.get("patterns", []),
                     partition_type=partition_type,
+                    gather_output=spec_dict.get("gather_output", False),
                     shape=shape,
                     partition_dim=spec_dict.get("partition_dim"),
                     model_types=spec_dict.get("model_types"),

deepspeed/module_inject/layers.py

@@ -228,6 +228,48 @@ def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
         return None, grad_output
 
 
+class GatherFromTensorParallelRegion(torch.autograd.Function):
+    """Gather last-dimension shards while keeping the output replicated."""
+
+    @staticmethod
+    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor:
+        ctx.group = group
+        if group is None:
+            ctx.partition_sizes = (input.shape[-1], )
+            ctx.tp_index = 0
+            return input
+
+        tp_world_size = dist.get_world_size(group=group)
+        ctx.tp_index = dist.get_rank(group=group)
+        if tp_world_size == 1:
+            ctx.partition_sizes = (input.shape[-1], )
+            return input
+
+        local_size = torch.tensor([input.shape[-1]], dtype=torch.long, device=input.device)
+        gathered_sizes = [torch.empty_like(local_size) for _ in range(tp_world_size)]
+        dist.all_gather(gathered_sizes, local_size, group=group)
+        ctx.partition_sizes = tuple(int(size.item()) for size in gathered_sizes)
+
+        max_partition_size = max(ctx.partition_sizes)
+        if input.shape[-1] == max_partition_size:
+            input_padded = input.contiguous()
+        else:
+            padded_shape = (*input.shape[:-1], max_partition_size)
+            input_padded = input.new_zeros(padded_shape)
+            input_padded[..., :input.shape[-1]].copy_(input)
+
+        gathered = [torch.empty_like(input_padded) for _ in range(tp_world_size)]
+        dist.all_gather(gathered, input_padded, group=group)
+        return torch.cat([shard[..., :size] for shard, size in zip(gathered, ctx.partition_sizes)], dim=-1)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
+        shard_offset = sum(ctx.partition_sizes[:ctx.tp_index])
+        shard_size = ctx.partition_sizes[ctx.tp_index]
+        grad_input = grad_output.narrow(-1, shard_offset, shard_size).contiguous()
+        return None, grad_input
+
+
 class TensorParallel_Layer(nn.Module, ABC):
     """
     A base class for model layers with  tensor parallelism support.
@@ -677,10 +719,11 @@ def _mark_uc_metadata(self):
 #remove kwargs from partition.
 class LinearLayer(TensorParallel_Layer):
 
-    def __init__(self, module, mp_group=None, skip_partition=False, **kwargs):
+    def __init__(self, module, mp_group=None, skip_partition=False, gather_output=False, **kwargs):
         super(LinearLayer, self).__init__(mp_group, **kwargs)
         self.weight = module.weight
         self.bias = module.bias
+        self.gather_output = gather_output
         if not skip_partition and self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
         self.support_training = True
@@ -699,6 +742,9 @@ def forward(self, input):
         else:
             output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)
 
+        if self.gather_output:
+            output = GatherFromTensorParallelRegion.apply(self.mp_group, output)
+
         return output
 
     @torch.no_grad()
@@ -765,7 +811,7 @@ def _mark_uc_metadata(self):
 
     # for bwc
     @classmethod
-    def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None):
+    def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None, gather_output=False):
         if weight is not None:
             in_features = weight.shape[1]
             out_features = weight.shape[0]
@@ -777,7 +823,7 @@ def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=Non
             in_features = weight_shape[1]
             out_features = weight_shape[0]
             linear = nn.Linear(in_features, out_features, bias=(bias is not None))
-        return cls(linear, skip_partition=True)
+        return cls(linear, skip_partition=True, gather_output=gather_output)
 
 
 class FusedModuleWrapper:

deepspeed/module_inject/tp_plan_converter.py

@@ -9,7 +9,8 @@
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_STYLES = {"colwise", "rowwise"}
+SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise"}
+# `colwise_rep` was renamed to `colwise_gather_output` in huggingface/transformers#42809.
 
 
 class TPPlanConverter:
@@ -33,10 +34,13 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
 
         for pattern, partition in hf_tp_plan.items():
             regex_pattern = TPPlanConverter._wildcard_to_regex(pattern)
+            partition_style = partition.lower()
+            gather_output = False
 
-            if partition.lower() == "colwise":
+            if partition_style in ("colwise", "colwise_rep", "colwise_gather_output"):
                 partition_type = PartitionType.COLUMN
-            elif partition.lower() == "rowwise":
+                gather_output = partition_style != "colwise"
+            elif partition_style == "rowwise":
                 partition_type = PartitionType.ROW
 
             # Only add .weight suffix if not already present
@@ -45,10 +49,12 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
             else:
                 regex_pattern += r"$"
 
-            layer_specs.append(TPLayerSpec(
-                patterns=[regex_pattern],
-                partition_type=partition_type,
-            ))
+            layer_specs.append(
+                TPLayerSpec(
+                    patterns=[regex_pattern],
+                    partition_type=partition_type,
+                    gather_output=gather_output,
+                ))
 
         return layer_specs