From 2936b726f0acedfa467a53fb393ad8df125f676f Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 27 Jul 2026 12:34:25 +0000
Subject: [PATCH 01/10] Enhance AutoTP Uneven Column Parallel Gathered Output

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/checkpoint/universal_checkpoint.py  |   7 ++
 deepspeed/module_inject/auto_tp.py            |  15 ---
 deepspeed/module_inject/layers.py             |  72 ++++++++-----
 docs/_tutorials/autotp-training.md            |   5 +-
 docs/code-docs/source/training.rst            |  11 +-
 .../checkpoint/test_autotp_uc_checkpoint.py   |  40 +++++++
 .../model_parallelism/test_autotp_training.py | 100 ++++++++++++++++--
 .../test_tp_plan_real_models.py               |  52 +++++++++
 .../test_tp_partition_config_path.py          |   7 +-
 .../test_autotp_universal_checkpoint.py       |  23 ++++
 10 files changed, 271 insertions(+), 61 deletions(-)

diff --git a/deepspeed/checkpoint/universal_checkpoint.py b/deepspeed/checkpoint/universal_checkpoint.py
index f057393ecdfc..ed8534c75814 100644
--- a/deepspeed/checkpoint/universal_checkpoint.py
+++ b/deepspeed/checkpoint/universal_checkpoint.py
@@ -40,6 +40,7 @@ def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank,
     logical_shape = meta.get('logical_shape')
     sub_param_shape = meta.get('sub_param_shape')
     sub_param_sizes = meta.get('sub_param_sizes')
+    partition_sizes = meta.get('partition_sizes')
     replicated = meta.get('replicated', False)
 
     if replicated:
@@ -92,6 +93,12 @@ def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank,
         slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
         return slice_tensor.flatten()
 
+    if partition_sizes is not None:
+        shard_offset = sum(partition_sizes[:tp_rank])
+        shard_size = partition_sizes[tp_rank]
+        slice_tensor = full_view.narrow(partition_dim, shard_offset, shard_size)
+        return slice_tensor.flatten()
+
     slice_tensor = full_view.chunk(tp_world_size, dim=partition_dim)[tp_rank]
     return slice_tensor.flatten()
 
diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index 783a00d08ab6..01069ed84b38 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/deepspeed/module_inject/auto_tp.py
@@ -475,21 +475,6 @@ 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.
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 33b1fbe3dbd0..c5d4bc1db6de 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -64,6 +64,7 @@ def _build_param_uc_restore_meta(*,
                                  output_shape=None,
                                  sub_param_shape=None,
                                  sub_param_sizes=None,
+                                 partition_sizes=None,
                                  target_partition_shape=None,
                                  original_shape=None,
                                  is_bias=False,
@@ -86,6 +87,8 @@ def _build_param_uc_restore_meta(*,
         _normalize_uc_shape(sub_param_shape),
         'sub_param_sizes':
         _normalize_uc_shape(sub_param_sizes),
+        'partition_sizes':
+        _normalize_uc_shape(partition_sizes),
         'target_partition_shape':
         _normalize_uc_shape(target_partition_shape),
         'original_shape':
@@ -389,6 +392,7 @@ def _set_param_uc_meta(self,
                            output_shape=None,
                            sub_param_shape=None,
                            sub_param_sizes=None,
+                           partition_sizes=None,
                            target_partition_shape=None,
                            original_shape=None,
                            is_bias=False,
@@ -403,6 +407,7 @@ def _set_param_uc_meta(self,
                                          output_shape=output_shape,
                                          sub_param_shape=sub_param_shape,
                                          sub_param_sizes=sub_param_sizes,
+                                         partition_sizes=partition_sizes,
                                          target_partition_shape=target_partition_shape,
                                          original_shape=original_shape,
                                          is_bias=is_bias,
@@ -724,6 +729,8 @@ def __init__(self, module, mp_group=None, skip_partition=False, gather_output=Fa
         self.weight = module.weight
         self.bias = module.bias
         self.gather_output = gather_output
+        self._orig_weight_shape = tuple(module.weight.shape)
+        self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
         if not skip_partition and self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
         self.support_training = True
@@ -749,31 +756,40 @@ def forward(self, input):
 
     @torch.no_grad()
     def gather_params(self, params_list):
-        #  Does not support uneven shard.
         for idx, param in enumerate(params_list):
+            if param is None:
+                continue
 
             params_list[idx].data_partition = param.data
-            output_param = torch.empty((self.tp_world_size * param.shape[0], *param.shape[1:]),
-                                       dtype=param.dtype,
-                                       device=param.device)
-            dist.all_gather_into_tensor(output_param, param, group=self.mp_group)
+            if self.mp_group is None or self.tp_world_size == 1:
+                output_param = param.data.contiguous()
+            else:
+                local_size = torch.tensor([param.shape[0]], dtype=torch.long, device=param.device)
+                gathered_sizes = [torch.empty_like(local_size) for _ in range(self.tp_world_size)]
+                dist.all_gather(gathered_sizes, local_size, group=self.mp_group)
+                partition_sizes = tuple(int(size.item()) for size in gathered_sizes)
+                max_partition_size = max(partition_sizes)
+
+                if param.shape[0] == max_partition_size:
+                    param_padded = param.contiguous()
+                else:
+                    padded_shape = (max_partition_size, *param.shape[1:])
+                    param_padded = param.new_zeros(padded_shape)
+                    param_padded[:param.shape[0]].copy_(param)
+
+                gathered = [torch.empty_like(param_padded) for _ in range(self.tp_world_size)]
+                dist.all_gather(gathered, param_padded, group=self.mp_group)
+                output_param = torch.cat([shard[:size] for shard, size in zip(gathered, partition_sizes)], dim=0)
+
+            if param is self.weight:
+                output_param = output_param.reshape(self._orig_weight_shape)
+            elif param is self.bias and self._orig_bias_shape is not None:
+                output_param = output_param.reshape(self._orig_bias_shape)
             params_list[idx].data = output_param.contiguous()
 
     @torch.no_grad()
     def _tp_partition(self, params_list):
-
-        if not self.is_training_mode():
-            self.uneven_partition(params_list)
-            return
-        for idx, param in enumerate(params_list):
-            if param is None:
-                return
-            #split bias if provide
-            _partition = torch.chunk(param, self.tp_world_size, dim=0)[self.tp_index]
-
-            _partition = self.move(_partition).detach()
-
-            params_list[idx].data = _partition
+        self.uneven_partition(params_list)
 
     def uneven_partition(self, params_list):
 
@@ -781,7 +797,6 @@ def uneven_partition(self, params_list):
             if param is None:
                 #split bias if provide
                 return
-            assert self.name is not None, "The module name must be provided in the initialization."
             _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size,
                                                                     self.name),
                                                 dim=0)[self.tp_index]
@@ -791,22 +806,25 @@ def uneven_partition(self, params_list):
             params_list[idx].data = _partition
 
     def _mark_uc_metadata(self):
-        original_out_dim = self.weight.shape[0] * self.tp_world_size
-        original_weight_shape = (original_out_dim, self.weight.shape[1])
+        original_out_dim = self._orig_weight_shape[0]
+        partition_sizes = get_shard_size_list(original_out_dim, self.tp_world_size, self.name)
         self._set_param_uc_meta(self.weight,
                                 partition_type='column',
                                 partition_dim=0,
-                                logical_shape=original_weight_shape,
+                                logical_shape=self._orig_weight_shape,
                                 output_shape=(original_out_dim, ),
-                                original_shape=original_weight_shape)
+                                partition_sizes=partition_sizes,
+                                target_partition_shape=tuple(self.weight.shape),
+                                original_shape=self._orig_weight_shape)
         if self.bias is not None:
-            original_bias_shape = (self.bias.shape[0] * self.tp_world_size, )
             self._set_param_uc_meta(self.bias,
                                     partition_type='column',
                                     partition_dim=0,
-                                    logical_shape=original_bias_shape,
-                                    output_shape=original_bias_shape,
-                                    original_shape=original_bias_shape,
+                                    logical_shape=self._orig_bias_shape,
+                                    output_shape=self._orig_bias_shape,
+                                    partition_sizes=partition_sizes,
+                                    target_partition_shape=tuple(self.bias.shape),
+                                    original_shape=self._orig_bias_shape,
                                     is_bias=True)
 
     # for bwc
diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md
index a98ac3d7093e..a4c475d4fe30 100644
--- a/docs/_tutorials/autotp-training.md
+++ b/docs/_tutorials/autotp-training.md
@@ -121,7 +121,10 @@ DeepSpeed will read the model's `tp_plan` at initialization and convert it to
 internal partition rules. The supported types are `colwise`, `rowwise`,
 and `colwise_gather_output`(`colwise_rep`). The gathered column styles shard
 the linear weight along its output dimension and AllGather the local output
-shards so every tensor-parallel rank receives the complete output.
+shards so every tensor-parallel rank receives the complete output. For untied
+output layers, the output dimension does not need to be divisible by
+`autotp_size`; DeepSpeed uses uneven local shards and gathers back to the
+original logical output size.
 
 Gathered column parallelism currently supports untied output layers. If an
 output layer such as `lm_head` shares the same runtime `Parameter` object with
diff --git a/docs/code-docs/source/training.rst b/docs/code-docs/source/training.rst
index 0e4c43a5309b..8420c993f939 100644
--- a/docs/code-docs/source/training.rst
+++ b/docs/code-docs/source/training.rst
@@ -428,10 +428,13 @@ The resolution priority is:
 2. HuggingFace ``tp_plan`` (from model config)
 3. AutoTP heuristics / ``preset_model`` (lowest priority)
 
-Currently only ``colwise`` and ``rowwise`` partition types from the HuggingFace
-``tp_plan`` are supported. Other types (``colwise_rep``, ``local_colwise``,
-``local_rowwise``, ``local_packed_rowwise``, ``gather``, ``sequence_parallel``)
-are not yet handled and will raise an error.
+Currently ``colwise``, ``rowwise``, and gathered column styles
+(``colwise_gather_output`` / ``colwise_rep``) from the HuggingFace ``tp_plan``
+are supported. Gathered column output layers may have output dimensions that are
+not divisible by ``autotp_size``; DeepSpeed uses uneven local shards and gathers
+back to the original logical output size. Other types (``local_colwise``,
+``local_rowwise``, ``local_packed_rowwise``, ``gather``,
+``sequence_parallel``) are not yet handled and will raise an error.
 
 Heuristic rules
 ^^^^^^^^^^^^^^^
diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
index 4a23e5b43716..5f9d1e008bff 100644
--- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
+++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
@@ -140,6 +140,46 @@ def test_resolve_autotp_partition_subparam_sizes_uneven_gqa_like():
     assert torch.equal(slice_flat, expected)
 
 
+def test_resolve_autotp_partition_uses_uneven_partition_sizes():
+    full_hp_param = torch.arange(101 * 4, dtype=torch.float32).view(101, 4)
+    param = _make_param(
+        (50, 4), {
+            'partition_type': 'column',
+            'partition_dim': 0,
+            'logical_shape': (101, 4),
+            'output_shape': (101, ),
+            'partition_sizes': (51, 50),
+            'original_shape': (101, 4),
+            'is_bias': False,
+            'replicated': False,
+        })
+
+    slice_flat = _resolve_autotp_partition(param, {PARAM: full_hp_param}, full_hp_param, tp_rank=1, tp_world_size=2)
+
+    expected = full_hp_param.narrow(0, 51, 50).flatten()
+    assert torch.equal(slice_flat, expected)
+
+
+def test_resolve_autotp_partition_uses_uneven_partition_sizes_for_bias():
+    full_hp_param = torch.arange(101, dtype=torch.float32)
+    param = _make_param(
+        (50, ), {
+            'partition_type': 'column',
+            'partition_dim': 0,
+            'logical_shape': (101, ),
+            'output_shape': (101, ),
+            'partition_sizes': (51, 50),
+            'original_shape': (101, ),
+            'is_bias': True,
+            'replicated': False,
+        })
+
+    slice_flat = _resolve_autotp_partition(param, {PARAM: full_hp_param}, full_hp_param, tp_rank=1, tp_world_size=2)
+
+    expected = full_hp_param.narrow(0, 51, 50).flatten()
+    assert torch.equal(slice_flat, expected)
+
+
 def test_resolve_autotp_partition_replicated_bias():
     full_hp_param = torch.arange(8, dtype=torch.float32)
     param = _make_param(
diff --git a/tests/unit/model_parallelism/test_autotp_training.py b/tests/unit/model_parallelism/test_autotp_training.py
index 07b3a1e8c796..f20c1136f9e7 100644
--- a/tests/unit/model_parallelism/test_autotp_training.py
+++ b/tests/unit/model_parallelism/test_autotp_training.py
@@ -16,8 +16,8 @@
 from deepspeed.utils import groups
 from contextlib import contextmanager
 from torch import nn
-from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, set_autotp_mode,
-                                            is_autotp_training_mode)
+from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, set_autotp_mode, is_autotp_training_mode
+from deepspeed.module_inject.tp_shard import get_shard_size_list
 from unit.checkpoint.common import compare_lr_scheduler_states, compare_optimizer_states
 import os
 from deepspeed.runtime.utils import is_model_parallel_parameter
@@ -430,10 +430,12 @@ def test(self, tp_size: int):
             model(batch[0], batch[1])
 
 
-def process_linear_layer(hidden_dim, input):
+def process_linear_layer(hidden_dim, input, output_dim=None):
     torch.manual_seed(42)
+    if output_dim is None:
+        output_dim = hidden_dim
     torch_linear = nn.Linear(hidden_dim,
-                             hidden_dim,
+                             output_dim,
                              dtype=preferred_dtype(),
                              device=get_accelerator().current_device())
     torch_out = torch_linear(input)
@@ -442,7 +444,12 @@ def process_linear_layer(hidden_dim, input):
     return torch_linear, torch_out
 
 
-def run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel, use_tp_model_init=False, gather_output=False):
+def run_tp_layer_fwd_bwd(tp_size,
+                         tp_overlap_comm,
+                         column_parallel,
+                         use_tp_model_init=False,
+                         gather_output=False,
+                         output_dim=None):
     skip_on_device()
     hidden_dim = 128
     batch_size_per_device = 1
@@ -499,7 +506,7 @@ def run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel, use_tp_model
 
     # Note: correctness checks below use standalone TP wrappers and do not
     # rely on the model's AutoTP-partitioned parameters.
-    torch_linear, torch_out = process_linear_layer(hidden_dim, input)
+    torch_linear, torch_out = process_linear_layer(hidden_dim, input, output_dim=output_dim)
     if column_parallel:
         linear = LinearLayer(deepcopy(torch_linear),
                              groups.get_tensor_model_parallel_group(),
@@ -509,10 +516,13 @@ def run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel, use_tp_model
         loss.backward()
 
         expected_out = torch_out
+        output_partition_sizes = get_shard_size_list(torch_out.shape[-1], tp_size, linear.name)
+        tp_rank = groups.get_tensor_model_parallel_rank()
         if not gather_output:
-            expected_out = torch.chunk(torch_out, tp_size, dim=-1)[groups.get_tensor_model_parallel_rank()]
-        torch_grad = torch.chunk(torch_linear.weight.grad, tp_size, dim=0)[groups.get_tensor_model_parallel_rank()]
-        torch_bias_grad = torch.chunk(torch_linear.bias.grad, tp_size, dim=0)[groups.get_tensor_model_parallel_rank()]
+            shard_offset = sum(output_partition_sizes[:tp_rank])
+            expected_out = torch_out.narrow(-1, shard_offset, output_partition_sizes[tp_rank])
+        torch_grad = torch_linear.weight.grad.split(output_partition_sizes, dim=0)[tp_rank]
+        torch_bias_grad = torch_linear.bias.grad.split(output_partition_sizes, dim=0)[tp_rank]
 
         torch.testing.assert_close(linear.bias.grad,
                                    torch_bias_grad.to(get_accelerator().current_device()),
@@ -562,6 +572,9 @@ def testColumnParallel(self, tp_size: int, tp_overlap_comm: bool):
     def testGatheredColumnParallel(self, tp_size: int, tp_overlap_comm: bool):
         run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=True, gather_output=True)
 
+    def testUnevenGatheredColumnParallel(self, tp_size: int, tp_overlap_comm: bool):
+        run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=True, gather_output=True, output_dim=129)
+
 
 # @pytest.mark.sequential
 class TestParamsGather(DistributedTest):
@@ -650,6 +663,68 @@ def test(self, layer_type):
 
         assert expected_tp_params == tp_params2
 
+    def test_uneven_linear_gather_params(self):
+        skip_on_device()
+        tp_size = 4
+        hidden_dim = 128
+        output_dim = 129
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": tp_size,
+                "partition_config": {
+                    "use_default_specs": False,
+                    "layer_specs": [{
+                        "patterns": [".*\\.weight$"],
+                        "partition_type": "skip",
+                    }],
+                }
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        torch.manual_seed(42)
+        model = SequentialLinearModel(hidden_dim=hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        torch_linear = nn.Linear(hidden_dim, output_dim, dtype=preferred_dtype(), device="cpu")
+        total_params = sum(p.numel() for p in torch_linear.parameters())
+        tp_layer = LinearLayer(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
+        tp_rank = groups.get_tensor_model_parallel_rank()
+        output_partition_sizes = get_shard_size_list(output_dim, tp_size, tp_layer.name)
+        expected_tp_params = output_partition_sizes[tp_rank] * (hidden_dim + 1)
+
+        assert expected_tp_params == sum(p.numel() for p in tp_layer.parameters())
+
+        for name, param in tp_layer.named_parameters(recurse=False):
+            if is_model_parallel_parameter(param):
+                param.gather_params([param])
+
+        torch_linear = torch_linear.to(get_accelerator().current_device())
+        is_same_weights = all(
+            torch.equal(param1, param2) for param1, param2 in zip(tp_layer.parameters(), torch_linear.parameters()))
+
+        assert is_same_weights
+        assert total_params == sum(p.numel() for p in tp_layer.parameters())
+
+        for name, param in tp_layer.named_parameters(recurse=False):
+            if is_model_parallel_parameter(param):
+                param._tp_partition([param])
+
+        assert expected_tp_params == sum(p.numel() for p in tp_layer.parameters())
+
 
 def dummy_init_engine(config):
     # This is a dummy initialization function for the DeepSpeed engine.
@@ -716,8 +791,11 @@ def test_consolidated_checkpoint(self):
 
         engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
 
-        assert isinstance(engine.module.lm_head, LmHeadLinearAllreduce)
-        assert engine.module.lm_head.weight.shape == (vocab_size, hidden_dim // self.world_size)
+        tp_rank = groups.get_tensor_model_parallel_rank()
+        output_partition_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        assert isinstance(engine.module.lm_head, LinearLayer)
+        assert engine.module.lm_head.gather_output
+        assert engine.module.lm_head.weight.shape == (output_partition_sizes[tp_rank], hidden_dim)
 
         checkpoint = engine._consolidated_16bit_state_dict()
 
diff --git a/tests/unit/model_parallelism/test_tp_plan_real_models.py b/tests/unit/model_parallelism/test_tp_plan_real_models.py
index 0268ed4cfb44..d65b09a63204 100644
--- a/tests/unit/model_parallelism/test_tp_plan_real_models.py
+++ b/tests/unit/model_parallelism/test_tp_plan_real_models.py
@@ -9,6 +9,7 @@
 import deepspeed
 from deepspeed.accelerator import get_accelerator
 from deepspeed.module_inject.layers import LinearLayer
+from deepspeed.module_inject.tp_shard import get_shard_size_list
 from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
 from deepspeed.utils import groups
 from unit.common import DistributedTest
@@ -93,6 +94,57 @@ def test_qwen2_tp_plan_with_zero2(self):
 
             assert not torch.isnan(outputs.loss)
 
+    def test_qwen2_tp_plan_with_uneven_vocab(self):
+        """Test an untied Qwen2 LM head with an uneven gathered vocabulary size."""
+        skip_on_device()
+
+        try:
+            from transformers import AutoModelForCausalLM, Qwen2Config
+        except ImportError:
+            pytest.skip("transformers not installed")
+
+        config = Qwen2Config(
+            vocab_size=1001,
+            hidden_size=128,
+            intermediate_size=256,
+            num_hidden_layers=1,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            tie_word_embeddings=False,
+        )
+
+        model = AutoModelForCausalLM.from_config(config)
+        assert model.lm_head.weight is not model.model.embed_tokens.weight
+
+        ds_config = {
+            "train_micro_batch_size_per_gpu": 1,
+            "tensor_parallel": {
+                "autotp_size": 2
+            },
+            "zero_optimization": {
+                "stage": 0
+            },
+        }
+
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
+
+        tp_rank = groups.get_tensor_model_parallel_rank()
+        output_partition_sizes = get_shard_size_list(config.vocab_size, 2, "lm_head")
+        assert engine.autotp_size() == 2
+        assert isinstance(model.lm_head, LinearLayer)
+        assert model.lm_head.gather_output
+        assert model.lm_head.weight.shape == (output_partition_sizes[tp_rank], config.hidden_size)
+        assert model.model.embed_tokens.weight.shape == (config.vocab_size, config.hidden_size)
+
+        input_ids = torch.randint(0, config.vocab_size, (1, 16)).to(get_accelerator().current_device_name())
+        dist.broadcast(
+            input_ids,
+            src=groups.get_tensor_model_parallel_src_rank(),
+            group=groups.get_tensor_model_parallel_group(),
+        )
+        outputs = engine(input_ids)
+        assert outputs.logits.shape == (1, 16, config.vocab_size)
+
     def test_qwen2_tied_lm_head_falls_back_to_replicated(self):
         """Test that an actual Qwen2 Parameter tie remains replicated."""
         skip_on_device()
diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py
index 3e952e439b87..c86dc71afa18 100644
--- a/tests/unit/module_inject/test_tp_partition_config_path.py
+++ b/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, LmHeadLinearAllreduce
+from deepspeed.module_inject.layers import LinearLayer
 
 
 class SubAttn(nn.Module):
@@ -176,13 +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():
+def test_gathered_lm_head_uses_column_parallel_layer_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)
+    assert isinstance(model.lm_head, LinearLayer)
+    assert model.lm_head.gather_output
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
index f0f7315562f2..636f28bff776 100644
--- a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
+++ b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
@@ -80,6 +80,26 @@ def test_subparam_layer_marks_standardized_param_metadata():
     assert tuple(bias_meta["target_partition_shape"]) == tuple(layer.bias.shape)
 
 
+def test_linear_layer_marks_uneven_column_metadata():
+    layer = LinearLayer(torch.nn.Linear(8, 101, bias=True), mp_group=None, name="lm_head")
+    layer.tp_world_size = 2
+    layer.weight.data = layer.weight.data[:51].contiguous()
+    layer.bias.data = layer.bias.data[:51].contiguous()
+    layer._mark_uc_metadata()
+
+    weight_meta = getattr(layer.weight, DS_AUTOTP_UC_META)
+    bias_meta = getattr(layer.bias, DS_AUTOTP_UC_META)
+
+    assert weight_meta["logical_shape"] == (101, 8)
+    assert weight_meta["output_shape"] == (101, )
+    assert weight_meta["partition_sizes"] == (51, 50)
+    assert weight_meta["target_partition_shape"] == (51, 8)
+    assert weight_meta["original_shape"] == (101, 8)
+    assert bias_meta["logical_shape"] == (101, )
+    assert bias_meta["partition_sizes"] == (51, 50)
+    assert bias_meta["target_partition_shape"] == (51, )
+
+
 def test_universal_checkpoint_info_excludes_param_level_recovery_fields():
     layer = SubParamLinearLayer(torch.nn.Linear(12, 12, bias=True),
                                 mp_group=None,
@@ -96,6 +116,7 @@ def test_universal_checkpoint_info_excludes_param_level_recovery_fields():
     assert "partition_dim" in subparam_entry
     assert "patterns" in subparam_entry
     assert "sub_param_sizes" not in subparam_entry
+    assert "partition_sizes" not in subparam_entry
     assert "target_partition_shape" not in subparam_entry
 
 
@@ -126,6 +147,7 @@ def test_param_uc_restore_builder_normalizes_shapes_and_nests_conversion_view():
                                                 output_shape=[12],
                                                 sub_param_shape=[3, -1],
                                                 sub_param_sizes=[4, 4, 4],
+                                                partition_sizes=[6, 6],
                                                 target_partition_shape=torch.Size([4, 8]),
                                                 original_shape=torch.Size([12, 8]),
                                                 is_bias=False,
@@ -135,6 +157,7 @@ def test_param_uc_restore_builder_normalizes_shapes_and_nests_conversion_view():
     assert restore_meta["output_shape"] == (12, )
     assert restore_meta["sub_param_shape"] == (3, -1)
     assert restore_meta["sub_param_sizes"] == (4, 4, 4)
+    assert restore_meta["partition_sizes"] == (6, 6)
     assert restore_meta["target_partition_shape"] == (4, 8)
     assert restore_meta["original_shape"] == (12, 8)
     assert restore_meta["conversion"] == {

From 847c4e4460aa3ef82e1a19f1a0c00150611cc58b Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 27 Jul 2026 14:25:03 +0000
Subject: [PATCH 02/10] Complete AutoTP uneven sharding for row parallelism and
 checkpoints

Making column-parallel layers uneven-aware left the row-parallel side on
the old even-split assumption. Because a column layer's output dimension
and the following row layer's input dimension are the same physical
dimension, the two must agree per rank. They no longer did.

With num_kv_heads set (the heuristic AutoTP path), hidden=384 and tp=4,
q_proj was sharded [128, 128, 64, 64] by get_shard_size_list while o_proj
was still sharded [96, 96, 96, 96] by torch.chunk, so the forward pass
died with:

    RuntimeError: mat1 and mat2 shapes cannot be multiplied
                  (2x128 and 96x384)

get_shard_size_list is the correct splitter here: update_mp_params derives
each rank's num_attention_heads from the same function, so weights must be
split the same way to stay consistent with the head metadata. torch.chunk
cannot express this (it never pads, front-loads the remainder, and may even
return fewer than tp_world_size chunks).

This commit:

* Makes LinearAllreduce uneven-aware. _tp_partition now always uses
  uneven_partition, dropping the training-only torch.chunk branch that
  existed solely because gather_params could not handle uneven shards.
  _mark_uc_metadata records the true original shape and partition sizes
  instead of deriving them as shape[1] * tp_world_size.

* Adds TensorParallel_Layer._all_gather_shards, shared by both the row and
  column paths. Partition sizes are recomputed locally from the same
  deterministic split rather than discovered with an extra collective, and
  uneven shards are zero padded to a common size so the faster uniform
  all_gather_into_tensor stays usable.

* Teaches ds_to_universal about uneven shards. main() collapsed every tp
  rank's PARAM_SHAPES into one flat dict, so _merge_zero_shards reshaped
  every rank's slice to a single shape and conversion failed with:

      RuntimeError: shape '[50, 12]' is invalid for input of size 612

  Shapes are now kept per tp rank. The concatenation itself was already
  uneven-safe; only the reshape was wrong.

* Skips the legacy vocabulary padding in load_hp_checkpoint_state when
  AutoTP restore metadata is present. That path derives the padded size as
  shape[0] * tp_world_size, which contradicts an uneven partition that
  _resolve_autotp_partition already describes exactly.

* Asserts in get_shard_size_list that shard sizes sum to the dimension
  size. tp_grain_size quantization silently violates this today, e.g.
  get_shard_size_list(1001, 2) returns [512, 448] with tp_grain_size=64.

Removing the transposes that row-side gathering previously needed also
makes it faster, and the column path returns to its original cost:

    tp=4, bf16, 16384x16384      before      after
    column, even shards         +6% (regr)   +0.1% over comm floor
    row, even shards            baseline     -8%

Tested with 64 AutoTP unit tests plus a non-AutoTP universal checkpoint
subset, including new end-to-end save/convert/load coverage for an uneven
lm_head (vocab 101, tp=2) and uneven GQA attention (hidden 384, tp=4).

Signed-off-by: iLeGend <824040212@qq.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 deepspeed/checkpoint/ds_to_universal.py       |  50 ++--
 deepspeed/checkpoint/universal_checkpoint.py  |   5 +-
 deepspeed/module_inject/layers.py             | 126 ++++----
 deepspeed/module_inject/tp_shard.py           |   7 +
 .../checkpoint/test_autotp_uc_checkpoint.py   | 281 +++++++++++++++++-
 5 files changed, 381 insertions(+), 88 deletions(-)

diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py
index d0573a85906d..99bba403f6de 100755
--- a/deepspeed/checkpoint/ds_to_universal.py
+++ b/deepspeed/checkpoint/ds_to_universal.py
@@ -212,7 +212,30 @@ def dump_param_fragment(dir, tp_index, dp_index, state_name, state_flat_tensor,
     _save_checkpoint(path, state_flat_tensor)
 
 
-def _merge_zero_shards(param_base_path, state, tp_degree, slice_shape=None):
+def _collect_slice_shapes(ds_checkpoint):
+    """Collect each parameter's per-tp-rank slice shape, ordered by tp rank.
+
+    AutoTP may shard a dimension unevenly, so tp ranks cannot be assumed to share a single
+    slice shape. Model state files are named by model-parallel rank, which enumerates the
+    (pp, tp) grid with tp varying fastest, so reading them in order yields tp order.
+    """
+    slice_shapes = {}
+    for mp_rank_file in ds_checkpoint.mp_rank_files:
+        mp_sd = torch.load(mp_rank_file, map_location=torch.device('cpu'), weights_only=False)
+        for sub_group_shapes in mp_sd[PARAM_SHAPES]:
+            for param_name, param_shape in sub_group_shapes.items():
+                slice_shapes.setdefault(param_name, []).append(param_shape)
+
+    for param_name, shapes in slice_shapes.items():
+        assert len(shapes) == ds_checkpoint.tp_degree, (
+            f"Expected {ds_checkpoint.tp_degree} slice shapes for {param_name} (one per tp rank), "
+            f"but found {len(shapes)}. The checkpoint layout does not match tp_degree="
+            f"{ds_checkpoint.tp_degree}, pp_degree={ds_checkpoint.pp_degree}.")
+
+    return slice_shapes
+
+
+def _merge_zero_shards(param_base_path, state, tp_degree, slice_shapes=None):
     slices = []
     for tp_index in range(tp_degree):
         prefix_path = os.path.join(param_base_path, str(tp_index), f"{state}")
@@ -237,18 +260,18 @@ def _merge_zero_shards(param_base_path, state, tp_degree, slice_shape=None):
             assert all(v == shards[0] for v in shards), "All shards must have the same step value"
             slice = shards[0]
         else:
-            if slice_shape is None:
-                slice = torch.cat(shards, dim=0)
-            else:
-                slice = torch.cat(shards, dim=0).reshape(slice_shape)
+            slice = torch.cat(shards, dim=0)
+            if slice_shapes is not None:
+                # AutoTP may shard a dimension unevenly, so each tp rank has its own slice shape.
+                slice = slice.reshape(slice_shapes[tp_index])
 
         slices.append(slice)
     return slices
 
 
-def merge_tp_slices(ds_checkpoint, dir, slice_dir, tp_degree, name_and_shape):
+def merge_tp_slices(ds_checkpoint, dir, slice_dir, tp_degree, name_and_shapes):
 
-    name, shape = name_and_shape
+    name, slice_shapes = name_and_shapes
     slice_base_path = os.path.join(slice_dir, name)
     param_base_path = os.path.join(dir, name)
 
@@ -284,15 +307,14 @@ def get_matched_sub_params_pattern(name_):
 
     matched_sub_params_shape = get_matched_sub_params_pattern(name)
 
-    step_merged = _merge_zero_shards(slice_base_path, "step", tp_degree, shape)
+    step_merged = _merge_zero_shards(slice_base_path, "step", tp_degree, slice_shapes)
     if step_merged:
         _save_checkpoint(os.path.join(param_base_path, "step.pt"), step_merged[0])
 
     for state in ("fp32", "exp_avg", "exp_avg_sq"):
-        slices = _merge_zero_shards(slice_base_path, state, tp_degree, shape)
+        slices = _merge_zero_shards(slice_base_path, state, tp_degree, slice_shapes)
         final_path = os.path.join(param_base_path, f"{state}.pt")
 
-        #print(f"Expected shape: {shape}")
         #print(f"Fragment sizes:", list(frag.shape for frag in slices))
         ckpt_dict = {}
         if get_matched_pattern(replicated_parameters, name):
@@ -779,13 +801,7 @@ def main(args):
         checkpoint_paths = _create_checkpoint_paths(args.output_folder, iteration, ds_checkpoint.tp_degree,
                                                     ds_checkpoint.pp_degree)
 
-        slice_shapes = []
-        for mp_rank_file in ds_checkpoint.mp_rank_files:
-            mp_sd = torch.load(mp_rank_file, map_location=torch.device('cpu'), weights_only=False)
-            slice_shapes += mp_sd[PARAM_SHAPES]
-
-        # fix back to normal flat dict, merge duplicates for tp>1
-        slice_shapes = dict((k, v) for d in slice_shapes for k, v in d.items())
+        slice_shapes = _collect_slice_shapes(ds_checkpoint)
         temp_dir = os.path.join(args.output_folder, 'tmp')
 
         print('*** 1. Extracting ZeRO fragments')
diff --git a/deepspeed/checkpoint/universal_checkpoint.py b/deepspeed/checkpoint/universal_checkpoint.py
index ed8534c75814..487b5f632e18 100644
--- a/deepspeed/checkpoint/universal_checkpoint.py
+++ b/deepspeed/checkpoint/universal_checkpoint.py
@@ -163,7 +163,10 @@ def load_hp_checkpoint_state(self, folder, tp_rank, tp_world_size, ep_rank=0, ep
         # the converter to universal currently strips the original padding completely so the saved
         # weight is padding-free and we just need to add new padding depending on the target TP
         # degree
-        is_vocab_tensor = ckpt_dict.get(VOCAB_TENSOR, False) and not is_expert_param
+        # AutoTP restore metadata already describes the exact (possibly uneven) partition layout,
+        # so the legacy tp-degree-derived vocab padding must not be applied on top of it.
+        has_autotp_meta = _get_param_uc_restore_meta(self) is not None
+        is_vocab_tensor = ckpt_dict.get(VOCAB_TENSOR, False) and not is_expert_param and not has_autotp_meta
         if is_vocab_tensor:
             # In the absence of data passed from the user wrt new padded vocab specific to tp degree
             # we can again derive that data by reverse engineering the target shapes like so:
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index c5d4bc1db6de..efd1954b40c5 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -426,6 +426,39 @@ def is_training_mode(self):
         global DEEPSPEED_AUTOTP_MODE
         return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING
 
+    @torch.no_grad()
+    def _all_gather_shards(self, shard, partition_sizes, dim):
+        """Reassemble a parameter from its tensor parallel shards along ``dim``.
+
+        ``partition_sizes`` is derived locally from the same deterministic split used to
+        create the shards, so no extra collective is needed to discover the remote sizes.
+        Uneven shards are zero padded to a common size, which keeps the uniform (and
+        faster) ``all_gather_into_tensor`` collective usable, and are then trimmed back.
+        """
+        world_size = len(partition_sizes)
+        assert partition_sizes[self.tp_index] == shard.shape[dim], (
+            f"Rank {self.tp_index} holds {shard.shape[dim]} elements along dim {dim} of "
+            f"{self.name}, but the partition scheme expects {partition_sizes[self.tp_index]}.")
+
+        max_size = max(partition_sizes)
+        padded_shape = list(shard.shape)
+        padded_shape[dim] = max_size
+        if shard.shape[dim] == max_size:
+            padded = shard.contiguous()
+        else:
+            padded = shard.new_zeros(padded_shape)
+            padded.narrow(dim, 0, shard.shape[dim]).copy_(shard)
+
+        buffer = shard.new_empty((world_size * padded_shape[0], *padded_shape[1:]))
+        dist.all_gather_into_tensor(buffer, padded, group=self.mp_group)
+
+        if dim == 0 and min(partition_sizes) == max_size:
+            # Shards are uniform and concatenated along dim 0, so the flat buffer is the result.
+            return buffer
+
+        shards = buffer.view(world_size, *padded_shape)
+        return torch.cat([shards[i].narrow(dim, 0, size) for i, size in enumerate(partition_sizes)], dim=dim)
+
     def __deepcopy__(self, memo):
         # This function is designed for
         # 'mp_group' (a 'ProcessGroup') cannot be pickled during deepcopy in some usage.
@@ -631,6 +664,7 @@ def __init__(self, module, mp_group, **kwargs):
         super(LinearAllreduce, self).__init__(mp_group, **kwargs)
         self.weight = module.weight
         self.bias = module.bias
+        self._orig_weight_shape = tuple(module.weight.shape)
 
         if self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
@@ -650,51 +684,34 @@ def forward(self, input):
 
     @torch.no_grad()
     def gather_params(self, params_list):
-
-        for idx, param in enumerate(params_list):
-            if param is None or idx > 0:
-                # don't gather bias
-                return
-            params_list[idx].data_partition = param.data
-            param = param.transpose(0, 1).contiguous()
-
-            output_param = torch.empty(self.tp_world_size * param.shape[0],
-                                       param.shape[1],
-                                       dtype=param.dtype,
-                                       device=param.device)
-            dist.all_gather_into_tensor(output_param, param, group=self.mp_group)
-            params_list[idx].data = output_param.transpose(0, 1).contiguous()
-        return
-
-    @torch.no_grad()
-    def _tp_partition(self, params_list):
-
-        if not self.is_training_mode():
-            self.uneven_partition(params_list)
+        # Row parallelism only shards the weight; the bias is replicated across ranks.
+        weight = params_list[0]
+        if weight is None:
             return
 
-        else:
-            for idx, param in enumerate(params_list):
-                if param is None:
-                    # don't slipt bias
-                    return
-                if idx > 0:  # move bias to device at initialization
-                    _partition = self.move(param).detach()
-                    params_list[idx].data = _partition
-                    return
+        weight.data_partition = weight.data
+        if self.mp_group is None or self.tp_world_size == 1:
+            weight.data = weight.data.contiguous()
+            return
 
-                _partition = torch.chunk(param, self.tp_world_size, dim=-1)[self.tp_index]
+        partition_sizes = get_shard_size_list(self._orig_weight_shape[1], self.tp_world_size, self.name)
+        weight.data = self._all_gather_shards(weight, partition_sizes, dim=1).contiguous()
 
-                _partition = self.move(_partition).detach()
+    @torch.no_grad()
+    def _tp_partition(self, params_list):
+        # Row parallelism shards the weight's input dimension; the bias stays replicated.
+        self.uneven_partition(params_list)
 
-                params_list[idx].data = _partition
+        bias = params_list[1] if len(params_list) > 1 else None
+        if bias is not None and self.is_training_mode():
+            # Training materializes the replicated bias on the target device.
+            bias.data = self.move(bias).detach()
 
     def uneven_partition(self, params_list):
         for idx, param in enumerate(params_list):
             if param is None or idx > 0:
                 # don't slipt bias
                 return
-            assert self.name is not None, "The module name must be provided in the initialization."
             _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size,
                                                                     self.name),
                                                 dim=1)[self.tp_index]
@@ -703,13 +720,15 @@ def uneven_partition(self, params_list):
             params_list[idx].data = _partition
 
     def _mark_uc_metadata(self):
-        original_weight_shape = (self.weight.shape[0], self.weight.shape[1] * self.tp_world_size)
+        partition_sizes = get_shard_size_list(self._orig_weight_shape[1], self.tp_world_size, self.name)
         self._set_param_uc_meta(self.weight,
                                 partition_type='row',
                                 partition_dim=1,
-                                logical_shape=original_weight_shape,
-                                output_shape=(original_weight_shape[0], ),
-                                original_shape=original_weight_shape)
+                                logical_shape=self._orig_weight_shape,
+                                output_shape=(self._orig_weight_shape[0], ),
+                                partition_sizes=partition_sizes,
+                                target_partition_shape=tuple(self.weight.shape),
+                                original_shape=self._orig_weight_shape)
         if self.bias is not None:
             self._set_param_uc_meta(self.bias,
                                     partition_type='row',
@@ -762,30 +781,13 @@ def gather_params(self, params_list):
 
             params_list[idx].data_partition = param.data
             if self.mp_group is None or self.tp_world_size == 1:
-                output_param = param.data.contiguous()
-            else:
-                local_size = torch.tensor([param.shape[0]], dtype=torch.long, device=param.device)
-                gathered_sizes = [torch.empty_like(local_size) for _ in range(self.tp_world_size)]
-                dist.all_gather(gathered_sizes, local_size, group=self.mp_group)
-                partition_sizes = tuple(int(size.item()) for size in gathered_sizes)
-                max_partition_size = max(partition_sizes)
-
-                if param.shape[0] == max_partition_size:
-                    param_padded = param.contiguous()
-                else:
-                    padded_shape = (max_partition_size, *param.shape[1:])
-                    param_padded = param.new_zeros(padded_shape)
-                    param_padded[:param.shape[0]].copy_(param)
-
-                gathered = [torch.empty_like(param_padded) for _ in range(self.tp_world_size)]
-                dist.all_gather(gathered, param_padded, group=self.mp_group)
-                output_param = torch.cat([shard[:size] for shard, size in zip(gathered, partition_sizes)], dim=0)
-
-            if param is self.weight:
-                output_param = output_param.reshape(self._orig_weight_shape)
-            elif param is self.bias and self._orig_bias_shape is not None:
-                output_param = output_param.reshape(self._orig_bias_shape)
-            params_list[idx].data = output_param.contiguous()
+                params_list[idx].data = param.data.contiguous()
+                continue
+
+            # Column parallelism shards dim 0 of both the weight and the bias, so gathering
+            # along dim 0 restores the original shape.
+            partition_sizes = get_shard_size_list(self._orig_weight_shape[0], self.tp_world_size, self.name)
+            params_list[idx].data = self._all_gather_shards(param, partition_sizes, dim=0).contiguous()
 
     @torch.no_grad()
     def _tp_partition(self, params_list):
diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index f1dbaae43ec9..992ad1cd9da5 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -76,4 +76,11 @@ def get_shard_size_list(total_size, mp_size, name=None):
     shard_sizes = []
     for i in range(mp_size):
         shard_sizes.append(get_shard_size(total_size, mp_size, name, i))
+    # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
+    # reconstruct the original tensor. tp_grain_size quantization can violate this when the
+    # dimension is not a multiple of the grain size.
+    assert sum(shard_sizes) == total_size, (
+        f"AutoTP shard sizes {shard_sizes} for layer '{name}' do not sum to the dimension size {total_size} "
+        f"with tp_size={mp_size} and tp_grain_size={tp_grain_size}. Choose a tp_grain_size that divides "
+        f"{total_size}.")
     return shard_sizes
diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
index 5f9d1e008bff..9abaa41ba02f 100644
--- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
+++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
@@ -3,21 +3,29 @@
 
 # DeepSpeed Team
 
+import os
 import types
 from types import SimpleNamespace
 
+import pytest
 import torch
 
-from deepspeed.checkpoint.constants import (CAT_DIM, FP32_WEIGHT_KEY, PARAM, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
-                                            PARAMETER_WITH_SUB_PARAMS, SUB_PARAM_SHAPE,
-                                            TP_REPLICATED_PARAMETER_PATTERNS, UNIVERSAL_CHECKPOINT_INFO)
+import deepspeed
+import deepspeed.comm as dist
+from deepspeed.checkpoint.ds_to_universal import main as convert_to_universal
+from deepspeed.checkpoint.constants import (CAT_DIM, FP32_WEIGHT_KEY, PARAM, PARAM_SHAPES,
+                                            PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_SUB_PARAMS,
+                                            SUB_PARAM_SHAPE, TP_REPLICATED_PARAMETER_PATTERNS,
+                                            UNIVERSAL_CHECKPOINT_INFO)
 from deepspeed.checkpoint.universal_checkpoint import SubparamShape as CheckpointSubparamShape
-from deepspeed.checkpoint.ds_to_universal import merge_tp_slices
+from deepspeed.checkpoint.ds_to_universal import _collect_slice_shapes, merge_tp_slices
 from deepspeed.checkpoint.universal_checkpoint import (_get_param_uc_restore_meta, _resolve_autotp_partition,
                                                        load_hp_checkpoint_state)
 from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
 from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
 
+from unit.common import DistributedTest
+
 
 class _DummyAddress:
 
@@ -270,7 +278,8 @@ def test_merge_tp_slices_emits_subparam_shape_metadata(tmp_path):
     ds_checkpoint = SimpleNamespace(
         get_checkpoint_info=lambda key: uc_info if key == UNIVERSAL_CHECKPOINT_INFO else {})
 
-    unmatched = merge_tp_slices(ds_checkpoint, str(output_dir), str(slice_dir), 2, (param_name, torch.Size([3, 4])))
+    unmatched = merge_tp_slices(ds_checkpoint, str(output_dir), str(slice_dir), 2,
+                                (param_name, [torch.Size([3, 4]), torch.Size([3, 4])]))
 
     ckpt = torch.load(output_dir / param_name / "fp32.pt", weights_only=False)
     assert not unmatched
@@ -283,8 +292,9 @@ def test_merge_tp_slices_uses_row_parallel_cat_dim(tmp_path):
     output_dir = tmp_path / "out"
     param_name = "module.proj.weight"
 
-    tp0 = torch.arange(16, dtype=torch.float32).view(4, 4)
-    tp1 = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
+    # Uneven row-parallel shards: rank 0 owns 3 input columns, rank 1 owns 2.
+    tp0 = torch.arange(12, dtype=torch.float32).view(4, 3)
+    tp1 = torch.arange(12, 20, dtype=torch.float32).view(4, 2)
     _write_tp_states(slice_dir, param_name, 0, tp0)
     _write_tp_states(slice_dir, param_name, 1, tp1)
 
@@ -297,7 +307,8 @@ def test_merge_tp_slices_uses_row_parallel_cat_dim(tmp_path):
     ds_checkpoint = SimpleNamespace(
         get_checkpoint_info=lambda key: uc_info if key == UNIVERSAL_CHECKPOINT_INFO else {})
 
-    merge_tp_slices(ds_checkpoint, str(output_dir), str(slice_dir), 2, (param_name, torch.Size([4, 4])))
+    merge_tp_slices(ds_checkpoint, str(output_dir), str(slice_dir), 2,
+                    (param_name, [torch.Size([4, 3]), torch.Size([4, 2])]))
 
     ckpt = torch.load(output_dir / param_name / "fp32.pt", weights_only=False)
     assert ckpt[CAT_DIM] == 1
@@ -350,3 +361,257 @@ def test_get_param_uc_restore_meta_returns_top_level_restore_schema():
 
     assert restore_meta["partition_dim"] == 1
     assert restore_meta["conversion"]["partition_dim"] == 999
+
+
+CP_TAG = "uneven_tp"
+UNIVERSAL_TAG = f"{CP_TAG}_universal"
+
+
+class UnevenVocabLmHeadModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, vocab_size):
+        super().__init__()
+        self.lm_head = torch.nn.Linear(hidden_dim, vocab_size)
+
+    def forward(self, x):
+        return self.lm_head(x).sum()
+
+
+class GQAAttentionModel(torch.nn.Module):
+    """Column-parallel q/k/v feeding a row-parallel o_proj, sharded on kv-head boundaries."""
+
+    class Config:
+
+        def __init__(self, hidden_dim, num_heads):
+            self.hidden_size = hidden_dim
+            self.num_attention_heads = num_heads
+            self.num_key_value_heads = num_heads
+
+    class Attention(torch.nn.Module):
+
+        def __init__(self, hidden_dim):
+            super().__init__()
+            self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
+            self.k_proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
+            self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
+            self.o_proj = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        def forward(self, x):
+            return self.o_proj(self.q_proj(x) + self.k_proj(x) + self.v_proj(x))
+
+    class Layer(torch.nn.Module):
+
+        def __init__(self, hidden_dim):
+            super().__init__()
+            self.self_attn = GQAAttentionModel.Attention(hidden_dim)
+
+        def forward(self, x):
+            return self.self_attn(x)
+
+    def __init__(self, hidden_dim, num_heads):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([GQAAttentionModel.Layer(hidden_dim)])
+        self.config = GQAAttentionModel.Config(hidden_dim, num_heads)
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x.sum()
+
+
+def _convert_to_universal(checkpoint_dir, universal_dir):
+    convert_to_universal(
+        SimpleNamespace(input_folder=checkpoint_dir,
+                        output_folder=universal_dir,
+                        num_extract_workers=1,
+                        num_merge_workers=1,
+                        keep_temp_folder=False,
+                        strict=True,
+                        inject_missing_state=False))
+
+
+def _train_steps(engine, hidden_dim, steps=3):
+    for _ in range(steps):
+        batch = torch.randn(2, hidden_dim, device=engine.device)
+        dist.broadcast(batch, src=0)
+        engine.backward(engine(batch))
+        engine.step()
+
+
+def _save_and_convert(engine, tmpdir):
+    engine.save_checkpoint(tmpdir, tag=CP_TAG, client_state={"iteration": 3})
+    dist.barrier()
+    if dist.get_rank() == 0:
+        _convert_to_universal(os.path.join(tmpdir, CP_TAG), os.path.join(tmpdir, UNIVERSAL_TAG))
+    dist.barrier()
+
+
+class TestUnevenColumnUniversalCheckpoint(DistributedTest):
+    world_size = 2
+    reuse_dist_env = False
+
+    def test_save_convert_load_uneven_lm_head(self, tmpdir):
+        hidden_dim = 12
+        vocab_size = 101  # Not divisible by the two TP ranks, giving shards of 51 and 50.
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": self.world_size,
+                "partition_config": {
+                    "use_default_specs":
+                    False,
+                    "layer_specs": [{
+                        "patterns": [r".*lm_head\.weight$"],
+                        "partition_type": "column",
+                        "gather_output": True,
+                    }],
+                },
+            },
+            "zero_optimization": {
+                "stage": 1
+            },
+        }
+
+        torch.manual_seed(42)
+        model = UnevenVocabLmHeadModel(hidden_dim, vocab_size)
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert engine.module.lm_head.weight.shape[0] in (51, 50)
+
+        _train_steps(engine, hidden_dim)
+        expected_weight = engine.module.lm_head.weight.detach().cpu().clone()
+        expected_bias = engine.module.lm_head.bias.detach().cpu().clone()
+        _save_and_convert(engine, tmpdir)
+
+        if dist.get_rank() == 0:
+            merged = torch.load(os.path.join(tmpdir, UNIVERSAL_TAG, "zero", "lm_head.weight", "fp32.pt"),
+                                weights_only=False)
+            assert tuple(merged[PARAM].shape) == (vocab_size, hidden_dim)
+
+        config_dict["checkpoint"] = {"load_universal": True}
+        torch.manual_seed(123)
+        restored = UnevenVocabLmHeadModel(hidden_dim, vocab_size)
+        restored_engine, _, _, _ = deepspeed.initialize(model=restored,
+                                                        model_parameters=restored.parameters(),
+                                                        config=config_dict)
+        restored_engine.load_checkpoint(tmpdir, tag=UNIVERSAL_TAG, load_optimizer_states=True)
+
+        torch.testing.assert_close(restored_engine.module.lm_head.weight.detach().cpu(), expected_weight)
+        torch.testing.assert_close(restored_engine.module.lm_head.bias.detach().cpu(), expected_bias)
+
+        # The optimizer must be usable after the restore.
+        _train_steps(restored_engine, hidden_dim, steps=1)
+
+
+class TestUnevenRowUniversalCheckpoint(DistributedTest):
+    world_size = 4
+    reuse_dist_env = False
+
+    def test_save_convert_load_uneven_row_parallel(self, tmpdir):
+        hidden_dim = 384
+        num_heads = 6  # Not divisible by the four TP ranks, giving shards of 128/128/64/64.
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": self.world_size
+            },
+            "zero_optimization": {
+                "stage": 1
+            },
+        }
+
+        torch.manual_seed(42)
+        model = GQAAttentionModel(hidden_dim, num_heads)
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        attn = engine.module.layers[0].self_attn
+        # Column and row parallelism must shard the same dimension identically.
+        assert attn.q_proj.weight.shape[0] == attn.o_proj.weight.shape[1]
+
+        _train_steps(engine, hidden_dim)
+        expected_q = attn.q_proj.weight.detach().cpu().clone()
+        expected_o = attn.o_proj.weight.detach().cpu().clone()
+        _save_and_convert(engine, tmpdir)
+
+        if dist.get_rank() == 0:
+            merged = torch.load(os.path.join(tmpdir, UNIVERSAL_TAG, "zero", "layers.0.self_attn.o_proj.weight",
+                                             "fp32.pt"),
+                                weights_only=False)
+            assert tuple(merged[PARAM].shape) == (hidden_dim, hidden_dim)
+
+        config_dict["checkpoint"] = {"load_universal": True}
+        torch.manual_seed(123)
+        restored = GQAAttentionModel(hidden_dim, num_heads)
+        restored_engine, _, _, _ = deepspeed.initialize(model=restored,
+                                                        model_parameters=restored.parameters(),
+                                                        config=config_dict)
+        restored_engine.load_checkpoint(tmpdir, tag=UNIVERSAL_TAG, load_optimizer_states=True)
+
+        restored_attn = restored_engine.module.layers[0].self_attn
+        torch.testing.assert_close(restored_attn.q_proj.weight.detach().cpu(), expected_q)
+        torch.testing.assert_close(restored_attn.o_proj.weight.detach().cpu(), expected_o)
+
+        _train_steps(restored_engine, hidden_dim, steps=1)
+
+
+def _write_mp_rank_file(dir_path, mp_rank, param_shapes):
+    os.makedirs(dir_path, exist_ok=True)
+    path = os.path.join(dir_path, f"mp_rank_{mp_rank:02d}_model_states.pt")
+    torch.save({PARAM_SHAPES: param_shapes}, path)
+    return path
+
+
+def test_collect_slice_shapes_keeps_uneven_shapes_in_tp_order(tmp_path):
+    # tp=2 with a column dimension of 5, giving shards of 3 and 2.
+    files = [
+        _write_mp_rank_file(tmp_path, 0, [{
+            "lm_head.weight": torch.Size([3, 4])
+        }]),
+        _write_mp_rank_file(tmp_path, 1, [{
+            "lm_head.weight": torch.Size([2, 4])
+        }]),
+    ]
+    ds_checkpoint = SimpleNamespace(mp_rank_files=files, tp_degree=2, pp_degree=1)
+
+    shapes = _collect_slice_shapes(ds_checkpoint)
+
+    assert shapes["lm_head.weight"] == [torch.Size([3, 4]), torch.Size([2, 4])]
+
+
+def test_collect_slice_shapes_pipeline_parallel_layout(tmp_path):
+    # Model-parallel ranks enumerate the (pp, tp) grid with tp varying fastest, and each
+    # pipeline stage only owns its own parameters.
+    stage0 = {"layers.0.weight": torch.Size([3, 4])}, {"layers.0.weight": torch.Size([2, 4])}
+    stage1 = {"layers.1.weight": torch.Size([3, 4])}, {"layers.1.weight": torch.Size([2, 4])}
+    files = [
+        _write_mp_rank_file(tmp_path, 0, [stage0[0]]),
+        _write_mp_rank_file(tmp_path, 1, [stage0[1]]),
+        _write_mp_rank_file(tmp_path, 2, [stage1[0]]),
+        _write_mp_rank_file(tmp_path, 3, [stage1[1]]),
+    ]
+    ds_checkpoint = SimpleNamespace(mp_rank_files=files, tp_degree=2, pp_degree=2)
+
+    shapes = _collect_slice_shapes(ds_checkpoint)
+
+    # Every parameter is collected once per tp rank, in tp order, despite spanning two stages.
+    assert shapes["layers.0.weight"] == [torch.Size([3, 4]), torch.Size([2, 4])]
+    assert shapes["layers.1.weight"] == [torch.Size([3, 4]), torch.Size([2, 4])]
+
+
+def test_collect_slice_shapes_rejects_unexpected_rank_count(tmp_path):
+    files = [_write_mp_rank_file(tmp_path, 0, [{"lm_head.weight": torch.Size([3, 4])}])]
+    ds_checkpoint = SimpleNamespace(mp_rank_files=files, tp_degree=2, pp_degree=1)
+
+    with pytest.raises(AssertionError, match="one per tp rank"):
+        _collect_slice_shapes(ds_checkpoint)

From 8acd22af657dfa6348419d279360b11b1140e9a0 Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Tue, 28 Jul 2026 03:19:06 +0000
Subject: [PATCH 03/10] Fix tied parameter shape collection in universal
 checkpoints

Collect parameter shapes by explicit TP rank and deduplicate replicas
across pipeline stages. Validate that replicated shapes agree before
keeping one shape per TP rank, preventing tied parameters from exceeding
the expected TP degree during universal checkpoint conversion.

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/checkpoint/ds_to_universal.py       | 32 +++++++++---
 .../checkpoint/test_autotp_uc_checkpoint.py   | 50 +++++++++++++++++--
 2 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py
index 99bba403f6de..52d793e15138 100755
--- a/deepspeed/checkpoint/ds_to_universal.py
+++ b/deepspeed/checkpoint/ds_to_universal.py
@@ -212,19 +212,37 @@ def dump_param_fragment(dir, tp_index, dp_index, state_name, state_flat_tensor,
     _save_checkpoint(path, state_flat_tensor)
 
 
+def _collect_tp_rank_slice_shapes(ds_checkpoint, tp_index):
+    """Collect one slice shape per parameter owned by a single tp rank.
+
+    Parameters tied across pipeline stages are replicated in every stage's model state file,
+    so the replicas are merged into a single entry after checking that they agree.
+    """
+    shapes_of_tp_rank = {}
+    for pp_index in range(ds_checkpoint.pp_degree):
+        for mp_rank_file in ds_checkpoint.get_2d_parallel_files(tp_index=tp_index, pp_index=pp_index):
+            mp_sd = torch.load(mp_rank_file, map_location=torch.device('cpu'), weights_only=False)
+            for sub_group_shapes in mp_sd[PARAM_SHAPES]:
+                for param_name, param_shape in sub_group_shapes.items():
+                    replicated_shape = shapes_of_tp_rank.setdefault(param_name, param_shape)
+                    assert replicated_shape == param_shape, (
+                        f"Pipeline replicas of {param_name} on tp rank {tp_index} disagree on "
+                        f"shape: {replicated_shape} vs {param_shape}.")
+
+    return shapes_of_tp_rank
+
+
 def _collect_slice_shapes(ds_checkpoint):
     """Collect each parameter's per-tp-rank slice shape, ordered by tp rank.
 
     AutoTP may shard a dimension unevenly, so tp ranks cannot be assumed to share a single
-    slice shape. Model state files are named by model-parallel rank, which enumerates the
-    (pp, tp) grid with tp varying fastest, so reading them in order yields tp order.
+    slice shape. Model state files are gathered per tp rank so that the collected shapes
+    follow tp order regardless of how model-parallel ranks are laid out on disk.
     """
     slice_shapes = {}
-    for mp_rank_file in ds_checkpoint.mp_rank_files:
-        mp_sd = torch.load(mp_rank_file, map_location=torch.device('cpu'), weights_only=False)
-        for sub_group_shapes in mp_sd[PARAM_SHAPES]:
-            for param_name, param_shape in sub_group_shapes.items():
-                slice_shapes.setdefault(param_name, []).append(param_shape)
+    for tp_index in range(ds_checkpoint.tp_degree):
+        for param_name, param_shape in _collect_tp_rank_slice_shapes(ds_checkpoint, tp_index).items():
+            slice_shapes.setdefault(param_name, []).append(param_shape)
 
     for param_name, shapes in slice_shapes.items():
         assert len(shapes) == ds_checkpoint.tp_degree, (
diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
index 9abaa41ba02f..a85d31666c03 100644
--- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
+++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
@@ -572,6 +572,22 @@ def _write_mp_rank_file(dir_path, mp_rank, param_shapes):
     return path
 
 
+class _FakeDSCheckpoint:
+    """Minimal stand-in for DeepSpeedCheckpoint that maps (pp, tp) to model state files."""
+
+    def __init__(self, mp_rank_files, tp_degree, pp_degree):
+        self.mp_rank_files = mp_rank_files
+        self.tp_degree = tp_degree
+        self.pp_degree = pp_degree
+
+    def get_2d_parallel_files(self, tp_index, pp_index):
+        # Model-parallel ranks enumerate the (pp, tp) grid with tp varying fastest.
+        mp_rank = pp_index * self.tp_degree + tp_index
+        if mp_rank >= len(self.mp_rank_files):
+            return []
+        return [self.mp_rank_files[mp_rank]]
+
+
 def test_collect_slice_shapes_keeps_uneven_shapes_in_tp_order(tmp_path):
     # tp=2 with a column dimension of 5, giving shards of 3 and 2.
     files = [
@@ -582,7 +598,7 @@ def test_collect_slice_shapes_keeps_uneven_shapes_in_tp_order(tmp_path):
             "lm_head.weight": torch.Size([2, 4])
         }]),
     ]
-    ds_checkpoint = SimpleNamespace(mp_rank_files=files, tp_degree=2, pp_degree=1)
+    ds_checkpoint = _FakeDSCheckpoint(files, tp_degree=2, pp_degree=1)
 
     shapes = _collect_slice_shapes(ds_checkpoint)
 
@@ -600,7 +616,7 @@ def test_collect_slice_shapes_pipeline_parallel_layout(tmp_path):
         _write_mp_rank_file(tmp_path, 2, [stage1[0]]),
         _write_mp_rank_file(tmp_path, 3, [stage1[1]]),
     ]
-    ds_checkpoint = SimpleNamespace(mp_rank_files=files, tp_degree=2, pp_degree=2)
+    ds_checkpoint = _FakeDSCheckpoint(files, tp_degree=2, pp_degree=2)
 
     shapes = _collect_slice_shapes(ds_checkpoint)
 
@@ -611,7 +627,35 @@ def test_collect_slice_shapes_pipeline_parallel_layout(tmp_path):
 
 def test_collect_slice_shapes_rejects_unexpected_rank_count(tmp_path):
     files = [_write_mp_rank_file(tmp_path, 0, [{"lm_head.weight": torch.Size([3, 4])}])]
-    ds_checkpoint = SimpleNamespace(mp_rank_files=files, tp_degree=2, pp_degree=1)
+    ds_checkpoint = _FakeDSCheckpoint(files, tp_degree=2, pp_degree=1)
 
     with pytest.raises(AssertionError, match="one per tp rank"):
         _collect_slice_shapes(ds_checkpoint)
+
+
+def test_collect_slice_shapes_dedupes_pipeline_tied_parameters(tmp_path):
+    # A tied embedding is replicated in the first and last pipeline stages, so it appears in
+    # every stage's model state file but must still yield exactly one shape per tp rank.
+    tied_tp0 = {"tied_modules.embed.word_embeddings.weight": torch.Size([3, 4])}
+    tied_tp1 = {"tied_modules.embed.word_embeddings.weight": torch.Size([2, 4])}
+    files = [
+        _write_mp_rank_file(tmp_path, 0, [{
+            **tied_tp0, "layers.0.weight": torch.Size([3, 4])
+        }]),
+        _write_mp_rank_file(tmp_path, 1, [{
+            **tied_tp1, "layers.0.weight": torch.Size([2, 4])
+        }]),
+        _write_mp_rank_file(tmp_path, 2, [{
+            **tied_tp0, "layers.1.weight": torch.Size([3, 4])
+        }]),
+        _write_mp_rank_file(tmp_path, 3, [{
+            **tied_tp1, "layers.1.weight": torch.Size([2, 4])
+        }]),
+    ]
+    ds_checkpoint = _FakeDSCheckpoint(files, tp_degree=2, pp_degree=2)
+
+    shapes = _collect_slice_shapes(ds_checkpoint)
+
+    assert shapes["tied_modules.embed.word_embeddings.weight"] == [torch.Size([3, 4]), torch.Size([2, 4])]
+    assert shapes["layers.0.weight"] == [torch.Size([3, 4]), torch.Size([2, 4])]
+    assert shapes["layers.1.weight"] == [torch.Size([3, 4]), torch.Size([2, 4])]

From be4c6d7c0d03c2a62032780916284b6a16880b1c Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Wed, 29 Jul 2026 04:09:54 +0000
Subject: [PATCH 04/10] make GatherFromTensorParallelRegion partition sizes
 recomputed locally

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/module_inject/layers.py | 37 ++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index efd1954b40c5..cb3a0fca7073 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -235,7 +235,15 @@ 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:
+    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, total_size: int,
+                name: Optional[str]) -> torch.Tensor:
+        """Gather the shards of a column parallel output whose full width is ``total_size``.
+
+        The shard widths follow the same deterministic split used to partition the weight, so
+        they are derived locally instead of being discovered with an extra collective. Uneven
+        shards are zero padded to a common width, which keeps the uniform (and faster)
+        ``all_gather_into_tensor`` collective usable, and are then trimmed back.
+        """
         ctx.group = group
         if group is None:
             ctx.partition_sizes = (input.shape[-1], )
@@ -248,29 +256,32 @@ def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Te
             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)
+        ctx.partition_sizes = tuple(get_shard_size_list(total_size, tp_world_size, name))
+        local_size = ctx.partition_sizes[ctx.tp_index]
+        assert local_size == input.shape[-1], (
+            f"Rank {ctx.tp_index} produced {input.shape[-1]} output features for {name}, but the "
+            f"partition scheme for a width of {total_size} expects {local_size}.")
 
         max_partition_size = max(ctx.partition_sizes)
-        if input.shape[-1] == max_partition_size:
+        if local_size == 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)
+            input_padded[..., :local_size].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)
+        buffer = input.new_empty((tp_world_size * input_padded.shape[0], *input_padded.shape[1:]))
+        dist.all_gather_into_tensor(buffer, input_padded, group=group)
+
+        shards = buffer.view(tp_world_size, *input_padded.shape)
+        return torch.cat([shards[i].narrow(-1, 0, size) for i, size in enumerate(ctx.partition_sizes)], dim=-1)
 
     @staticmethod
-    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
+    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]:
         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
+        return None, grad_input, None, None
 
 
 class TensorParallel_Layer(nn.Module, ABC):
@@ -769,7 +780,7 @@ def forward(self, input):
             output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)
 
         if self.gather_output:
-            output = GatherFromTensorParallelRegion.apply(self.mp_group, output)
+            output = GatherFromTensorParallelRegion.apply(self.mp_group, output, self._orig_weight_shape[0], self.name)
 
         return output
 

From 790dab9509ef98ef4ba0f656a94447e5aa67a20d Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 3 Aug 2026 01:22:01 +0800
Subject: [PATCH 05/10] Remove the diagnostics or change as logger.info

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/module_inject/auto_tp.py | 12 +++++----
 deepspeed/runtime/engine.py        | 43 ++++--------------------------
 2 files changed, 12 insertions(+), 43 deletions(-)

diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index 01069ed84b38..17c5c6338894 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/deepspeed/module_inject/auto_tp.py
@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 # Automatic Tensor Parallelism
+import logging
 import re
 
 from torch import nn
@@ -16,7 +17,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.utils.logging import log_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
 from .autotp_config import TPLayerSpec, AutoTPConfig, PartitionType
 
@@ -435,8 +436,8 @@ def _replace_with_config(self, child, name):
                     "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])
+            log_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
@@ -492,7 +493,7 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 name=name,
             )
         if spec.gather_output:
-            print_dist(f"AutoTP: replacing '{name}' with LinearLayer(gather_output=True)", ranks=[0])
+            log_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):
@@ -535,10 +536,11 @@ def _configure_gathered_column_tie_fallbacks(self):
                 continue
 
             self._tied_gathered_column_module_names.update((module_name, tied_embedding_name))
-            print_dist(
+            log_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],
+                level=logging.WARNING,
             )
 
         self._gathered_column_tie_fallbacks_configured = True
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index b1b99c305f92..e802c8b1d39a 100755
--- a/deepspeed/runtime/engine.py
+++ b/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_dist, print_json_dist, print_configuration, set_log_level_from_string
+from deepspeed.utils.logging import print_json_dist, print_configuration, set_log_level_from_string
 
 from deepspeed.accelerator import get_accelerator
 
@@ -708,42 +708,9 @@ def _apply_autotp_partitioning(self, model, tp_config):
             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=(),
@@ -775,7 +742,7 @@ def lm_head_entries(tp_plan):
                     pattern for pattern, style in hf_tp_plan.items()
                     if style.lower() in ("colwise_rep", "colwise_gather_output")
                 ]
-                print_dist(
+                log_dist(
                     f"Using HuggingFace tp_plan with {len(layer_specs)} layer specifications; "
                     f"gathered column output patterns={gathered_output_patterns}",
                     ranks=[0],
@@ -797,14 +764,14 @@ def lm_head_entries(tp_plan):
                 setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model))
                 setattr(model, "ds_autotp_parsed", True)
                 return
-            print_dist(
+            log_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])
+            log_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:

From 000c92f5fb96328b5e839494bb28d6e2bd36abc9 Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 3 Aug 2026 00:17:15 +0800
Subject: [PATCH 06/10] Freeze AutoTP partition sizes at layer construction

get_shard_size_list() reads the process-wide tp_shard globals num_kv_heads and
tp_grain_size, which a later init_inference call or a second AutoTP model
overwrites. Recomputing the split in the forward gather and in gather_params
therefore let them disagree with the shards the layer was built with.

Resolve it once in _freeze_partition_sizes() and have every consumer read the
cached value. A tp_world_size of 1 short-circuits the helper so its grain
quantization cannot truncate a replicated parameter.

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/module_inject/layers.py             | 83 ++++++++++---------
 .../test_autotp_universal_checkpoint.py       |  2 +
 2 files changed, 47 insertions(+), 38 deletions(-)

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index cb3a0fca7073..0b611b49c0c6 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -235,34 +235,30 @@ 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, total_size: int,
-                name: Optional[str]) -> torch.Tensor:
-        """Gather the shards of a column parallel output whose full width is ``total_size``.
+    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor,
+                partition_sizes: Tuple[int, ...]) -> torch.Tensor:
+        """Gather the shards of a column parallel output described by ``partition_sizes``.
 
-        The shard widths follow the same deterministic split used to partition the weight, so
-        they are derived locally instead of being discovered with an extra collective. Uneven
-        shards are zero padded to a common width, which keeps the uniform (and faster)
+        The widths were resolved when the layer was built, so they need neither an extra
+        collective nor a second lookup of the tensor parallel globals. Uneven shards are zero
+        padded to a common width, which keeps the uniform (and faster)
         ``all_gather_into_tensor`` collective usable, and are then trimmed back.
         """
         ctx.group = group
-        if group is None:
-            ctx.partition_sizes = (input.shape[-1], )
-            ctx.tp_index = 0
-            return input
+        ctx.partition_sizes = partition_sizes
+        ctx.tp_index = 0
 
-        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], )
+        tp_world_size = len(partition_sizes)
+        if group is None or tp_world_size == 1:
             return input
 
-        ctx.partition_sizes = tuple(get_shard_size_list(total_size, tp_world_size, name))
-        local_size = ctx.partition_sizes[ctx.tp_index]
+        ctx.tp_index = dist.get_rank(group=group)
+        local_size = partition_sizes[ctx.tp_index]
         assert local_size == input.shape[-1], (
-            f"Rank {ctx.tp_index} produced {input.shape[-1]} output features for {name}, but the "
-            f"partition scheme for a width of {total_size} expects {local_size}.")
+            f"Rank {ctx.tp_index} produced {input.shape[-1]} output features, but the partition "
+            f"scheme {partition_sizes} frozen at construction expects {local_size}.")
 
-        max_partition_size = max(ctx.partition_sizes)
+        max_partition_size = max(partition_sizes)
         if local_size == max_partition_size:
             input_padded = input.contiguous()
         else:
@@ -274,14 +270,14 @@ def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, total_size:
         dist.all_gather_into_tensor(buffer, input_padded, group=group)
 
         shards = buffer.view(tp_world_size, *input_padded.shape)
-        return torch.cat([shards[i].narrow(-1, 0, size) for i, size in enumerate(ctx.partition_sizes)], dim=-1)
+        return torch.cat([shards[i].narrow(-1, 0, size) for i, size in enumerate(partition_sizes)], dim=-1)
 
     @staticmethod
-    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]:
+    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None]:
         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, None, None
+        return None, grad_input, None
 
 
 class TensorParallel_Layer(nn.Module, ABC):
@@ -437,6 +433,23 @@ def is_training_mode(self):
         global DEEPSPEED_AUTOTP_MODE
         return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING
 
+    def _freeze_partition_sizes(self, total_size):
+        """Resolve the tensor parallel split of this layer once, while the layer is built.
+
+        ``get_shard_size_list`` reads the process-wide tp_shard globals (``num_kv_heads``,
+        ``tp_grain_size``), which a later ``init_inference`` call or a second AutoTP model
+        overwrites. The split is part of the checkpoint contract, so it is resolved here and
+        every later consumer -- the forward gather, the parameter gather and the checkpoint
+        metadata -- reads the cached value rather than querying those globals again.
+        """
+        if self.tp_world_size == 1:
+            # Nothing to split, so bypass the shard helper and its grain quantization, which
+            # would otherwise drop the tail of a dimension that is not a multiple of the grain.
+            self._partition_sizes = (total_size, )
+        else:
+            self._partition_sizes = tuple(get_shard_size_list(total_size, self.tp_world_size, self.name))
+        return self._partition_sizes
+
     @torch.no_grad()
     def _all_gather_shards(self, shard, partition_sizes, dim):
         """Reassemble a parameter from its tensor parallel shards along ``dim``.
@@ -676,6 +689,7 @@ def __init__(self, module, mp_group, **kwargs):
         self.weight = module.weight
         self.bias = module.bias
         self._orig_weight_shape = tuple(module.weight.shape)
+        self._freeze_partition_sizes(self._orig_weight_shape[1])
 
         if self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
@@ -705,8 +719,7 @@ def gather_params(self, params_list):
             weight.data = weight.data.contiguous()
             return
 
-        partition_sizes = get_shard_size_list(self._orig_weight_shape[1], self.tp_world_size, self.name)
-        weight.data = self._all_gather_shards(weight, partition_sizes, dim=1).contiguous()
+        weight.data = self._all_gather_shards(weight, self._partition_sizes, dim=1).contiguous()
 
     @torch.no_grad()
     def _tp_partition(self, params_list):
@@ -723,21 +736,18 @@ def uneven_partition(self, params_list):
             if param is None or idx > 0:
                 # don't slipt bias
                 return
-            _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size,
-                                                                    self.name),
-                                                dim=1)[self.tp_index]
+            _partition = params_list[idx].split(self._partition_sizes, dim=1)[self.tp_index]
 
             _partition = self.move(_partition).detach()
             params_list[idx].data = _partition
 
     def _mark_uc_metadata(self):
-        partition_sizes = get_shard_size_list(self._orig_weight_shape[1], self.tp_world_size, self.name)
         self._set_param_uc_meta(self.weight,
                                 partition_type='row',
                                 partition_dim=1,
                                 logical_shape=self._orig_weight_shape,
                                 output_shape=(self._orig_weight_shape[0], ),
-                                partition_sizes=partition_sizes,
+                                partition_sizes=self._partition_sizes,
                                 target_partition_shape=tuple(self.weight.shape),
                                 original_shape=self._orig_weight_shape)
         if self.bias is not None:
@@ -761,6 +771,7 @@ def __init__(self, module, mp_group=None, skip_partition=False, gather_output=Fa
         self.gather_output = gather_output
         self._orig_weight_shape = tuple(module.weight.shape)
         self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
+        self._freeze_partition_sizes(self._orig_weight_shape[0])
         if not skip_partition and self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
         self.support_training = True
@@ -780,7 +791,7 @@ def forward(self, input):
             output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)
 
         if self.gather_output:
-            output = GatherFromTensorParallelRegion.apply(self.mp_group, output, self._orig_weight_shape[0], self.name)
+            output = GatherFromTensorParallelRegion.apply(self.mp_group, output, self._partition_sizes)
 
         return output
 
@@ -797,8 +808,7 @@ def gather_params(self, params_list):
 
             # Column parallelism shards dim 0 of both the weight and the bias, so gathering
             # along dim 0 restores the original shape.
-            partition_sizes = get_shard_size_list(self._orig_weight_shape[0], self.tp_world_size, self.name)
-            params_list[idx].data = self._all_gather_shards(param, partition_sizes, dim=0).contiguous()
+            params_list[idx].data = self._all_gather_shards(param, self._partition_sizes, dim=0).contiguous()
 
     @torch.no_grad()
     def _tp_partition(self, params_list):
@@ -810,9 +820,7 @@ def uneven_partition(self, params_list):
             if param is None:
                 #split bias if provide
                 return
-            _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size,
-                                                                    self.name),
-                                                dim=0)[self.tp_index]
+            _partition = params_list[idx].split(self._partition_sizes, dim=0)[self.tp_index]
 
             _partition = self.move(_partition).detach()
 
@@ -820,13 +828,12 @@ def uneven_partition(self, params_list):
 
     def _mark_uc_metadata(self):
         original_out_dim = self._orig_weight_shape[0]
-        partition_sizes = get_shard_size_list(original_out_dim, self.tp_world_size, self.name)
         self._set_param_uc_meta(self.weight,
                                 partition_type='column',
                                 partition_dim=0,
                                 logical_shape=self._orig_weight_shape,
                                 output_shape=(original_out_dim, ),
-                                partition_sizes=partition_sizes,
+                                partition_sizes=self._partition_sizes,
                                 target_partition_shape=tuple(self.weight.shape),
                                 original_shape=self._orig_weight_shape)
         if self.bias is not None:
@@ -835,7 +842,7 @@ def _mark_uc_metadata(self):
                                     partition_dim=0,
                                     logical_shape=self._orig_bias_shape,
                                     output_shape=self._orig_bias_shape,
-                                    partition_sizes=partition_sizes,
+                                    partition_sizes=self._partition_sizes,
                                     target_partition_shape=tuple(self.bias.shape),
                                     original_shape=self._orig_bias_shape,
                                     is_bias=True)
diff --git a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
index 636f28bff776..0b32827aca29 100644
--- a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
+++ b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
@@ -82,7 +82,9 @@ def test_subparam_layer_marks_standardized_param_metadata():
 
 def test_linear_layer_marks_uneven_column_metadata():
     layer = LinearLayer(torch.nn.Linear(8, 101, bias=True), mp_group=None, name="lm_head")
+    # Stand in for a layer built under tp=2; the split is normally frozen during construction.
     layer.tp_world_size = 2
+    layer._freeze_partition_sizes(101)
     layer.weight.data = layer.weight.data[:51].contiguous()
     layer.bias.data = layer.bias.data[:51].contiguous()
     layer._mark_uc_metadata()

From de6159f8add1c65576e5ee51d6b0f7756395e6aa Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 3 Aug 2026 15:22:37 +0800
Subject: [PATCH 07/10] Add shape consistency guard for merging tied parameters
 in pipeline replicas

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/checkpoint/ds_to_universal.py       | 10 +++++++-
 .../checkpoint/test_autotp_uc_checkpoint.py   | 24 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py
index c639ae9a555f..25625f391ee7 100755
--- a/deepspeed/checkpoint/ds_to_universal.py
+++ b/deepspeed/checkpoint/ds_to_universal.py
@@ -733,7 +733,15 @@ def _group_per_tp_shapes(slice_shapes_by_tp, pp_degree, tp_degree):
     for tp in range(tp_degree):
         tp_dict = {}
         for pp in range(pp_degree):
-            tp_dict.update(slice_shapes_by_tp[pp * tp_degree + tp])
+            stage_shapes = slice_shapes_by_tp[pp * tp_degree + tp]
+            for name in tp_dict.keys() & stage_shapes.keys():
+                # A parameter tied across pipeline stages is stored in every stage's file.
+                # The replicas must agree, otherwise the tie was partitioned inconsistently
+                # and the update below would silently keep the stage that was read last.
+                assert tp_dict[name] == stage_shapes[name], (
+                    f"Pipeline replicas of {name} on tp rank {tp} disagree on shape: "
+                    f"{tp_dict[name]} vs {stage_shapes[name]}.")
+            tp_dict.update(stage_shapes)
         per_tp.append(tp_dict)
     all_names = set()
     for d in per_tp:
diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
index e611e446daf0..4dc26573ef23 100644
--- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
+++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
@@ -8,6 +8,7 @@
 import types
 from types import SimpleNamespace
 
+import pytest
 import torch
 import torch.nn as nn
 
@@ -568,6 +569,29 @@ def test_group_per_tp_shapes_handles_pp_local_params():
     assert result['layer1.weight'] == [(2, 4), (2, 4)]
 
 
+def test_group_per_tp_shapes_rejects_disagreeing_pipeline_replicas():
+    # A tied parameter is stored in every PP stage's file. Replicas that disagree mean the
+    # tie was partitioned inconsistently, which must not be papered over by keeping the
+    # stage that happened to be read last.
+    slice_shapes_by_tp = [
+        {
+            'tied_modules.embed.weight': (3, 4)
+        },  # pp0_tp0
+        {
+            'tied_modules.embed.weight': (2, 4)
+        },  # pp0_tp1
+        {
+            'tied_modules.embed.weight': (1, 4)
+        },  # pp1_tp0 disagrees with pp0_tp0
+        {
+            'tied_modules.embed.weight': (2, 4)
+        },  # pp1_tp1
+    ]
+
+    with pytest.raises(AssertionError, match='disagree on shape'):
+        _group_per_tp_shapes(slice_shapes_by_tp, pp_degree=2, tp_degree=2)
+
+
 class TestRealCheckpointUniversalConversionTPxPP(DistributedTest):
     # Generate a real ZeRO-1 checkpoint with TP=2, PP=2 on CPU (gloo) using the
     # production DeepSpeed writer, then convert it to a universal checkpoint.

From f0ce0b30ce86c77ea9242c970351e361541b41d3 Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 3 Aug 2026 16:32:54 +0800
Subject: [PATCH 08/10] Keep the sub-grain remainder when sharding a dimension

get_shard_size() quantizes a split to tp_grain_size by flooring the dimension
into whole grains, so total_size % tp_grain_size was dropped and the shards no
longer tiled the dimension. A GPT-2 vocabulary of 50257 over two ranks yielded
25152 + 25088 = 50240, silently losing the last 17 rows.

Give that tail to the last rank instead. Every other rank keeps the kernel
alignment tp_grain_size exists for, and the shards reconstruct the dimension
exactly, so the sum check in get_shard_size_list() is now an internal invariant
rather than a configuration error a user cannot act on.

The band where a dimension holds fewer grains than there are ranks still leaves
the high ranks with an empty shard. That is pre-existing behaviour, unrelated to
the dropped remainder, and is left alone here.

With the remainder preserved, a tp_world_size of 1 no longer needs to bypass the
shard helper to avoid truncation, so that special case is removed.

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/module_inject/layers.py         |  7 +---
 deepspeed/module_inject/tp_shard.py       | 19 ++++++----
 tests/unit/module_inject/test_tp_shard.py | 43 +++++++++++++++++++++++
 3 files changed, 56 insertions(+), 13 deletions(-)
 create mode 100644 tests/unit/module_inject/test_tp_shard.py

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 785edcf59e27..a4a58306c26e 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -442,12 +442,7 @@ def _freeze_partition_sizes(self, total_size):
         every later consumer -- the forward gather, the parameter gather and the checkpoint
         metadata -- reads the cached value rather than querying those globals again.
         """
-        if self.tp_world_size == 1:
-            # Nothing to split, so bypass the shard helper and its grain quantization, which
-            # would otherwise drop the tail of a dimension that is not a multiple of the grain.
-            self._partition_sizes = (total_size, )
-        else:
-            self._partition_sizes = tuple(get_shard_size_list(total_size, self.tp_world_size, self.name))
+        self._partition_sizes = tuple(get_shard_size_list(total_size, self.tp_world_size, self.name))
         return self._partition_sizes
 
     @torch.no_grad()
diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index 992ad1cd9da5..cd48d7990ecc 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -61,8 +61,14 @@ def get_shard_size(total_size, mp_size, name=None, rank=None):
         return total_size * my_slices // num_kv_heads
     else:
         if total_size >= tp_grain_size:
-            grain_size = total_size // tp_grain_size
-            return (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size
+            grain_size, remainder = divmod(total_size, tp_grain_size)
+            shard_size = (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size
+            if rank == mp_size - 1:
+                # Quantizing to tp_grain_size would otherwise drop total_size % tp_grain_size
+                # and silently truncate the dimension. Giving that tail to the last rank keeps
+                # every other rank aligned for the compute kernels.
+                shard_size += remainder
+            return shard_size
         else:
             return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
@@ -77,10 +83,9 @@ def get_shard_size_list(total_size, mp_size, name=None):
     for i in range(mp_size):
         shard_sizes.append(get_shard_size(total_size, mp_size, name, i))
     # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
-    # reconstruct the original tensor. tp_grain_size quantization can violate this when the
-    # dimension is not a multiple of the grain size.
+    # reconstruct the original tensor.
     assert sum(shard_sizes) == total_size, (
-        f"AutoTP shard sizes {shard_sizes} for layer '{name}' do not sum to the dimension size {total_size} "
-        f"with tp_size={mp_size} and tp_grain_size={tp_grain_size}. Choose a tp_grain_size that divides "
-        f"{total_size}.")
+        f"AutoTP shard sizes {shard_sizes} for layer '{name}' do not sum to the dimension size "
+        f"{total_size} with tp_size={mp_size}, tp_grain_size={tp_grain_size} and "
+        f"num_kv_heads={num_kv_heads}.")
     return shard_sizes
diff --git a/tests/unit/module_inject/test_tp_shard.py b/tests/unit/module_inject/test_tp_shard.py
new file mode 100644
index 000000000000..91d4ad52b03f
--- /dev/null
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+from deepspeed.module_inject import tp_shard
+from deepspeed.module_inject.tp_shard import get_shard_size_list, set_num_kv_heads, set_tp_grain_size
+
+
+@pytest.fixture(autouse=True)
+def restore_tp_shard_globals():
+    # tp_grain_size and num_kv_heads are process wide, so leaking them would change how
+    # unrelated tests partition their layers.
+    grain_size, kv_heads = tp_shard.tp_grain_size, tp_shard.num_kv_heads
+    yield
+    set_tp_grain_size(grain_size)
+    set_num_kv_heads(kv_heads)
+
+
+@pytest.mark.parametrize("total_size,tp_size", [(50257, 2), (50257, 8), (151936, 8), (32000, 4)])
+def test_grain_quantized_shards_tile_the_dimension(total_size, tp_size):
+    # A vocabulary that is not a multiple of tp_grain_size used to lose its tail to the grain
+    # quantization, so the shards no longer reconstructed the embedding table.
+    set_tp_grain_size(64)
+
+    shard_sizes = get_shard_size_list(total_size, tp_size, "lm_head")
+
+    assert sum(shard_sizes) == total_size
+    # Only the rank that absorbs the sub-grain tail gives up its alignment.
+    assert sum(1 for size in shard_sizes if size % 64) <= 1, shard_sizes
+
+
+def test_uneven_shards_without_grain_quantization():
+    assert get_shard_size_list(101, 2, "lm_head") == [51, 50]
+
+
+def test_kv_head_shards_tile_the_dimension():
+    set_num_kv_heads(6)
+
+    # 6 kv heads over 4 ranks gives 2/2/1/1 heads, so 384 hidden splits as 128/128/64/64.
+    assert get_shard_size_list(384, 4, "layers.0.self_attn.q_proj") == [128, 128, 64, 64]

From 4953f602bccbff409bc684221c26761138d99178 Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 3 Aug 2026 16:56:59 +0800
Subject: [PATCH 09/10] Cleanup the remaining AutoTP lm_head replacement traces

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/module_inject/auto_tp.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index 17c5c6338894..abc3875f855f 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/deepspeed/module_inject/auto_tp.py
@@ -428,17 +428,6 @@ 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,
-                }
-            log_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:
@@ -492,8 +481,6 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 partition_dim=spec.get_partition_dim(),
                 name=name,
             )
-        if spec.gather_output:
-            log_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):

From c17e2e6bae119ffbd1d53975edf572a2183c0f38 Mon Sep 17 00:00:00 2001
From: iLeGend <824040212@qq.com>
Date: Mon, 3 Aug 2026 10:08:25 +0000
Subject: [PATCH 10/10] fix yapf formatting

Signed-off-by: iLeGend <824040212@qq.com>
---
 deepspeed/module_inject/layers.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index a4a58306c26e..62600057e88d 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -235,8 +235,8 @@ 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,
-                partition_sizes: Tuple[int, ...]) -> torch.Tensor:
+    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, partition_sizes: Tuple[int,
+                                                                                                ...]) -> torch.Tensor:
         """Gather the shards of a column parallel output described by ``partition_sizes``.
 
         The widths were resolved when the layer was built, so they need neither an extra