[AutoTP] Enalbe 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$"],
+            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
@@ -48,6 +52,7 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
             layer_specs.append(TPLayerSpec(
                 patterns=[regex_pattern],
                 partition_type=partition_type,
+                gather_output=gather_output,
             ))
 
         return layer_specs

deepspeed/runtime/engine.py

@@ -135,7 +135,7 @@
 from ..git_version_info import version
 
 from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler
-from deepspeed.utils.logging import print_json_dist, print_configuration, set_log_level_from_string
+from deepspeed.utils.logging import print_dist, print_json_dist, print_configuration, set_log_level_from_string
 
 from deepspeed.accelerator import get_accelerator
 
@@ -703,6 +703,43 @@ def _apply_autotp_partitioning(self, model, tp_config):
         if hasattr(tp_config, "get_partition_config_object"):
             partition_config = tp_config.get_partition_config_object()
 
+        model_config = getattr(model, "config", None)
+        base_tp_plan = getattr(model_config, "base_model_tp_plan", None) if model_config is not None else None
+        class_tp_plan = getattr(type(model), "_tp_plan", None)
+        runtime_tp_plan = getattr(model, "__dict__", {}).get("_tp_plan")
+        from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
+        hf_tp_plan = _get_hf_tp_plan(model)
+
+        def lm_head_entries(tp_plan):
+            if not isinstance(tp_plan, dict):
+                return {}
+            return {
+                pattern: style
+                for pattern, style in tp_plan.items()
+                if any(part in ("lm_head", "embed_out") for part in pattern.split('.'))
+            }
+
+        lm_head_modules = [
+            name for name, _ in model.named_modules()
+            if name and any(part in ("lm_head", "embed_out") for part in name.split('.'))
+        ]
+        selected_route = "custom partition_config" if partition_config is not None else "HuggingFace tp_plan or AutoTP"
+        model_class = f"{type(model).__module__}.{type(model).__qualname__}"
+        print_dist(
+            f"AutoTP tp_plan diagnostics: model_class={model_class}; route={selected_route}; "
+            f"base_model_tp_plan={base_tp_plan!r}; type(model)._tp_plan={class_tp_plan!r}; "
+            f"instance_tp_plan={runtime_tp_plan!r}; "
+            f"effective_tp_plan={hf_tp_plan!r}",
+            ranks=[0],
+        )
+        print_dist(
+            f"AutoTP lm_head diagnostics: modules={lm_head_modules!r}; "
+            f"base_entries={lm_head_entries(base_tp_plan)!r}; class_entries={lm_head_entries(class_tp_plan)!r}; "
+            f"runtime_entries={lm_head_entries(runtime_tp_plan)!r}; "
+            f"effective_entries={lm_head_entries(hf_tp_plan)!r}",
+            ranks=[0],
+        )
+
         if partition_config is not None:
             autotp = AutoTP(module=model,
                             all_reduce_linears=(),
@@ -723,19 +760,22 @@ def _apply_autotp_partitioning(self, model, tp_config):
             setattr(model, "ds_autotp_parsed", True)
             return
 
-        model_config = getattr(model, "config", None)
         from deepspeed.module_inject import replace_transformer_layer
-
-        from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
-
-        hf_tp_plan = _get_hf_tp_plan(model)
         if hf_tp_plan:
             from deepspeed.module_inject.tp_plan_converter import TPPlanConverter
             from deepspeed.module_inject.autotp_config import AutoTPConfig
 
             layer_specs = TPPlanConverter.convert(hf_tp_plan)
             if layer_specs is not None:
-                logger.info(f"Using HuggingFace tp_plan with {len(layer_specs)} layer specifications")
+                gathered_output_patterns = [
+                    pattern for pattern, style in hf_tp_plan.items()
+                    if style.lower() in ("colwise_rep", "colwise_gather_output")
+                ]
+                print_dist(
+                    f"Using HuggingFace tp_plan with {len(layer_specs)} layer specifications; "
+                    f"gathered column output patterns={gathered_output_patterns}",
+                    ranks=[0],
+                )
                 tp_plan_config = AutoTPConfig(tp_size=tp_size, layer_specs=layer_specs)
                 autotp = AutoTP(
                     module=model,
@@ -752,6 +792,14 @@ def _apply_autotp_partitioning(self, model, tp_config):
                 autotp._replace_module(model)
                 setattr(model, "ds_autotp_parsed", True)
                 return
+            print_dist(
+                f"AutoTP: effective HuggingFace tp_plan could not be converted; falling back to heuristic AutoTP. "
+                f"styles={sorted(set(hf_tp_plan.values()))!r}",
+                ranks=[0],
+            )
+        else:
+            print_dist("AutoTP: no effective HuggingFace tp_plan was found; falling back to heuristic AutoTP.",
+                       ranks=[0])
 
         parser_dict = AutoTP.tp_parser(model)
         for client_module, injection_policy in parser_dict: