handle uneven gathered lm_head shards

Signed-off-by: iLeGend <824040212@qq.com>

Author: jinyouzhi

deepspeed/module_inject/auto_tp.py

@@ -475,6 +475,21 @@ 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, gather_output=spec.gather_output)
         # Only use fused-QKV heuristics when no partition_config is provided.

tests/unit/module_inject/test_tp_partition_config_path.py

@@ -12,7 +12,7 @@
 import torch.nn as nn
 
 from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec
-from deepspeed.module_inject.layers import LinearLayer
+from deepspeed.module_inject.layers import LinearLayer, LmHeadLinearAllreduce
 
 
 class SubAttn(nn.Module):
@@ -135,7 +135,7 @@ def test_nested_depth_correct():
         assert f"layers.{layer_idx}.self_attn.o_proj" in matched_names
 
 
-def _build_gathered_lm_head_autotp(model):
+def _build_gathered_lm_head_autotp(model, mp_size=1):
     config = AutoTPConfig(layer_specs=[
         TPLayerSpec(
             patterns=[r".*lm_head\.weight$"],
@@ -152,7 +152,7 @@ def _build_gathered_lm_head_autotp(model):
         orig_layer_impl=None,
         partition_config=config,
     )
-    autotp.set_tensor_parallel_config(1, None)
+    autotp.set_tensor_parallel_config(mp_size, None)
     autotp.update_linear_policies()
     return autotp
 
@@ -176,5 +176,14 @@ def test_gathered_lm_head_falls_back_for_runtime_parameter_tie():
     assert model.lm_head.weight is model.embed_tokens.weight
 
 
+def test_gathered_lm_head_falls_back_to_legacy_allreduce_when_output_dim_is_uneven():
+    model = OutputModel(tied=False)
+    model.lm_head = nn.Linear(32, 101, bias=False)
+
+    _build_gathered_lm_head_autotp(model, mp_size=2)._replace_module(model)
+
+    assert isinstance(model.lm_head, LmHeadLinearAllreduce)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])