From 82fd78fd959b44ef9c84152f1682273bea61ac64 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Mon, 10 Aug 2026 09:00:44 +0800
Subject: [PATCH 01/20] Replace tp_shard process-wide globals with per-model
 AutoTPMeta (#8231)

tp_shard kept num_kv_heads / num_attention_heads / n_embd / tp_grain_size as
process-wide mutable globals set during AutoTP replacement. A second AutoTP
model loaded into the same process overwrote them, so the first model's later
sharding / gather / checkpoint conversion silently read the wrong values.

Move that state onto a frozen AutoTPMeta dataclass computed once from the model
config and threaded through every sharding helper and TP layer. Each model now
carries its own kv-head / grain state, so multiple AutoTP models (teacher /
student, online distillation, RL actor + reference) can coexist in one process.

Ulysses sequence parallelism, which repurposed the same global, gets its own
private kv-head state so it no longer depends on whichever AutoTP model was
loaded last.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/module_inject/auto_tp.py            |  64 +++--
 .../module_inject/auto_tp_model_utils.py      |  19 +-
 deepspeed/module_inject/fusedqkv_utils.py     |  71 +++---
 deepspeed/module_inject/layers.py             |  52 ++--
 deepspeed/module_inject/replace_module.py     |  56 ++--
 deepspeed/module_inject/tp_shard.py           | 114 +++++----
 deepspeed/runtime/engine.py                   |  48 +---
 deepspeed/sequence/layer.py                   |  38 ++-
 .../checkpoint/test_autotp_uc_checkpoint.py   |   8 +-
 .../test_autotp_custom_patterns.py            | 241 +++++++++---------
 .../model_parallelism/test_autotp_training.py |   7 +-
 .../test_tp_plan_real_models.py               |   3 +-
 .../module_inject/test_fused_repartition.py   |  76 +++---
 .../test_tp_partition_config_path.py          |   2 +
 tests/unit/module_inject/test_tp_shard.py     |  81 +++---
 .../test_autotp_universal_checkpoint.py       |  42 ++-
 .../unit/sequence_parallelism/test_ulysses.py |   4 +-
 .../unit/v1/moe/test_autoep_autotp_runtime.py |   1 +
 18 files changed, 449 insertions(+), 478 deletions(-)

diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index e47c5e3ddcc5..95800a354c29 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/deepspeed/module_inject/auto_tp.py
@@ -15,7 +15,7 @@
 from .layers import *
 from deepspeed.accelerator import get_accelerator
 from .fusedqkv_utils import require_tp_fused_qkvw
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 from deepspeed.utils import groups
 from deepspeed.utils.logging import log_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
@@ -205,7 +205,9 @@ def __init__(self,
                  linear_layer_setting,
                  orig_layer_impl,
                  keep_module_on_host=False,
-                 partition_config: Optional[AutoTPConfig] = None):
+                 partition_config: Optional[AutoTPConfig] = None,
+                 model_config=None,
+                 tp_grain_size: int = 1):
         self.module = module
         self.all_reduce_linears = all_reduce_linears
         self.prefix = prefix
@@ -213,6 +215,9 @@ def __init__(self,
 
         self.mp_size = None
         self.mp_group = None
+        # Per-model TP metadata threaded through every layer / helper so each AutoTP instance
+        # shards by its own model's kv-head / grain values.
+        self.tp_meta = AutoTPMeta.from_model_config(model_config, tp_grain_size)
         self.linear_layer_setting = linear_layer_setting
         self.orig_layer_impl = orig_layer_impl
         self.linear_policies = None
@@ -378,14 +383,14 @@ def _replace(self, child, name, conv_linear_layer):
         # For Yuan model
         if 'Yuan' in str(self.module):
             if 'v_proj' in name:
-                return Yuan_LinearLayer(child, self.mp_group)
+                return Yuan_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
 
             elif 'o_proj' in name:
-                return Yuan_LinearAllreduce(child, self.mp_group)
+                return Yuan_LinearAllreduce(child, self.mp_group, tp_meta=self.tp_meta)
 
         # For MLP including chunk layer.
         if 'gate_up_proj' in name or ('dense_h_to_4h' in name and 'GLM' in str(self.module)):
-            return GateUpPack_LinearLayer(child, self.mp_group)
+            return GateUpPack_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
             # For Arctic model, bypass to all_reduce replacement for w2 weights
         arctic_w2_all_reduce_linear = False
         if 'Arctic' in str(self.module) and 'w2' in name:
@@ -398,21 +403,21 @@ def _replace(self, child, name, conv_linear_layer):
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
-                return Conv_LinearALlreduce(child, self.mp_group, name=name)
+                return Conv_LinearALlreduce(child, self.mp_group, name=name, tp_meta=self.tp_meta)
             elif name == "lm_head" or name == 'embed_out':
-                return LmHeadLinearAllreduce(child, self.mp_group)
+                return LmHeadLinearAllreduce(child, self.mp_group, tp_meta=self.tp_meta)
 
-            return LinearAllreduce(child, self.mp_group, name=name)
+            return LinearAllreduce(child, self.mp_group, name=name, tp_meta=self.tp_meta)
         else:
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
-                conv_LinearLayer(child, self.mp_group)
+                conv_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
             elif require_tp_fused_qkvw(name, self.mp_size):
                 #Check and handle fused qkv for TP
-                return fused_LinearLayer(child, self.mp_group, fused_module=self.module)
+                return fused_LinearLayer(child, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta)
 
-            return LinearLayer(child, self.mp_group, name=name)
+            return LinearLayer(child, self.mp_group, name=name, tp_meta=self.tp_meta)
 
     def _replace_with_config(self, child, name):
         """
@@ -450,10 +455,10 @@ def _replace_with_config(self, child, name):
     def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create row-parallel layer (AllReduce after forward)."""
         if self.conv_linear_layer:
-            return Conv_LinearALlreduce(module, self.mp_group, name=name)
+            return Conv_LinearALlreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
         # Check for lm_head / embed_out
         if name == "lm_head" or name == 'embed_out':
-            return LmHeadLinearAllreduce(module, self.mp_group)
+            return LmHeadLinearAllreduce(module, self.mp_group, tp_meta=self.tp_meta)
 
         if spec.shape is not None:
             return SubParamLinearAllreduce(
@@ -462,17 +467,22 @@ def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
+                tp_meta=self.tp_meta,
             )
-        return LinearAllreduce(module, self.mp_group, name=name)
+        return LinearAllreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
 
     def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create column-parallel layer (AllReduce in backward)."""
         if self.conv_linear_layer:
-            return conv_LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+            return conv_LinearLayer(module,
+                                    self.mp_group,
+                                    name=name,
+                                    gather_output=spec.gather_output,
+                                    tp_meta=self.tp_meta)
         # Only use fused-QKV heuristics when no partition_config is provided.
         elif self.partition_config is None and require_tp_fused_qkvw(name, self.mp_size):
             # Check and handle fused qkv for TP
-            return fused_LinearLayer(module, self.mp_group, fused_module=self.module)
+            return fused_LinearLayer(module, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta)
         if spec.shape is not None:
             if spec.gather_output:
                 raise NotImplementedError("AutoTP gather_output does not yet support shaped sub-parameter layers.")
@@ -482,8 +492,9 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
+                tp_meta=self.tp_meta,
             )
-        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output, tp_meta=self.tp_meta)
 
     def _configure_gathered_column_tie_fallbacks(self):
         """Configure a replicated fallback for gathered output layers tied to embeddings."""
@@ -561,7 +572,7 @@ def _slice_embedding(self, child, name, conv_linear_layer):
         mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
 
         original_shape = tuple(child.weight.shape)
-        partition_sizes = get_shard_size_list(original_shape[1], self.mp_size, name)
+        partition_sizes = get_shard_size_list(original_shape[1], self.mp_size, self.tp_meta, name)
         if hasattr(child.weight, 'ds_tensor'):
             data = child.weight.ds_tensor.data.split(partition_sizes, dim=1)
         else:
@@ -645,7 +656,7 @@ def update_mp_params(self, child, name=None):
                 param_val = getattr(child, param)
                 # get_shard_size selects its partitioning strategy from the module name, so the
                 # attributes must be sharded under the same name as the weights they describe.
-                setattr(child, param, get_shard_size(param_val, self.mp_size, name, rank=tp_index))
+                setattr(child, param, get_shard_size(param_val, self.mp_size, self.tp_meta, name, rank=tp_index))
         setattr(child, "replaced", True)
 
     def update_linear_policies(self):
@@ -754,21 +765,6 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
                 self._replace_module(child, name, class_name)
         return r_module
 
-    @staticmethod
-    def get_model_num_kv_heads(config):
-        num_kv_heads = None
-        # multi_query_group_num is for chatglm2 & chatglm3
-        kv_head_names = [
-            'multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads',
-            'attention_heads'
-        ]
-        for name in kv_head_names:
-            if hasattr(config, name):
-                num_kv_heads = getattr(config, name)
-                if num_kv_heads is not None:
-                    break
-        return num_kv_heads
-
     def _replace_last_linear_module(self, r_module):
         if hasattr(r_module, "lm_head"):
             name = "lm_head"
diff --git a/deepspeed/module_inject/auto_tp_model_utils.py b/deepspeed/module_inject/auto_tp_model_utils.py
index 8b4a3a6c852e..fbe0c9ac9635 100644
--- a/deepspeed/module_inject/auto_tp_model_utils.py
+++ b/deepspeed/module_inject/auto_tp_model_utils.py
@@ -8,7 +8,7 @@
 from deepspeed import comm as dist
 import torch
 from typing import Optional
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 
 
 class _HeadCountProxy:
@@ -27,16 +27,23 @@ def __setattr__(self, name, value):
             setattr(self._module, name, value)
 
 
-def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None):
+def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None, meta: Optional[AutoTPMeta] = None):
     tp_world_size = dist.get_world_size(group=mp_group)
     return get_shard_size_list(
         num_heads,
         tp_world_size,
+        meta or AutoTPMeta(),
         num_kv_heads=num_kv_heads,
     )
 
 
-def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=None, num_kv_heads=None):
+def install_head_sharded_helper(module,
+                                name,
+                                wrapper,
+                                mp_group=None,
+                                num_heads=None,
+                                num_kv_heads=None,
+                                meta: Optional[AutoTPMeta] = None):
     """Give ``module`` a head-slicing wrapper around one of its own methods.
 
     The wrapper is bound to this instance instead of installed on its class. A class-wide patch
@@ -48,7 +55,7 @@ def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=
     if original_name not in module.__dict__:
         # Wrapping an already wrapped instance would make it delegate to itself.
         setattr(module, original_name, getattr(module, name))
-    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads) if num_heads is not None else None
+    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads, meta) if num_heads is not None else None
     setattr(
         module, name,
         functools.partial(wrapper, module, mp_group=mp_group, head_shard_sizes=shard_sizes, total_num_heads=num_heads))
@@ -64,7 +71,9 @@ def _head_shard(num_heads, mp_group=None, head_shard_sizes=None, total_num_heads
     tp_world_size = dist.get_world_size(group=mp_group)
     tp_index = dist.get_rank(group=mp_group)
     full_num_heads = total_num_heads if total_num_heads is not None else num_heads
-    shard_sizes = head_shard_sizes or get_shard_size_list(full_num_heads, tp_world_size)
+    # The fallback split is only reached when no per-head shard sizes were recorded, i.e. when
+    # the model is not GQA-aware, so an even split (default meta, no kv-heads) is correct here.
+    shard_sizes = head_shard_sizes or get_shard_size_list(full_num_heads, tp_world_size, AutoTPMeta())
     if len(shard_sizes) != tp_world_size or sum(shard_sizes) != full_num_heads:
         raise ValueError(f"Head shard sizes {shard_sizes} do not partition {full_num_heads} heads across "
                          f"{tp_world_size} tensor-parallel ranks.")
diff --git a/deepspeed/module_inject/fusedqkv_utils.py b/deepspeed/module_inject/fusedqkv_utils.py
index 2d54e231819b..a43af5224e17 100644
--- a/deepspeed/module_inject/fusedqkv_utils.py
+++ b/deepspeed/module_inject/fusedqkv_utils.py
@@ -4,7 +4,7 @@
 # DeepSpeed Team
 import torch
 from deepspeed.utils.logging import warning_once
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, get_num_kv_heads, get_n_embd, get_num_attention_heads
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 
 
 def split_by_qkvlist_and_refuse(qkv_list, split_size, split_dim=0, cat_dim=0):
@@ -51,7 +51,7 @@ def get_fused_qkv_type(module):
     return FUSED_QKV_TYPE_DICT[max(module_name_matches, key=len)]
 
 
-def fused_qkv_subparam_sizes(module, weight_shape):
+def fused_qkv_subparam_sizes(module, weight_shape, meta: AutoTPMeta):
     """Sizes of the sub-parameters a fused qkv weight is cut into, or None.
 
     ``prepare_tp_fused_qkvw`` splits most layouts into q/k/v (or a single block) and
@@ -68,15 +68,15 @@ def fused_qkv_subparam_sizes(module, weight_shape):
     if fused_type == 'bloomtype':
         return (total_size, )
     if fused_type in ('glmtype', 'qwentype'):
-        if get_num_kv_heads() == 2:
-            hidden_dim = get_n_embd()
-            kv_dim = (total_size - hidden_dim) // get_num_kv_heads()
+        if meta.num_kv_heads == 2:
+            hidden_dim = meta.n_embd
+            kv_dim = (total_size - hidden_dim) // meta.num_kv_heads
             return (hidden_dim, kv_dim, kv_dim)
         third = total_size // 3
         return (third, third, third)
     if fused_type == 'phi3type':
-        head_dim = weight_shape[1] // get_num_attention_heads()
-        kv_dim = get_num_kv_heads() * head_dim
+        head_dim = weight_shape[1] // meta.num_attention_heads
+        kv_dim = meta.num_kv_heads * head_dim
         return (total_size - 2 * kv_dim, kv_dim, kv_dim)
     # codegentype interleaves blocks across ranks and bigcodetype replicates the kv block,
     # so neither is a per-sub-parameter split.
@@ -102,79 +102,80 @@ def set_fused_qkv_shard_state(module, shard_widths, tp_index):
     module.attn.split_size = query_width
 
 
-def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index):
+def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index, meta: AutoTPMeta):
 
     if src is None:
         return
 
-    def _codegen_type_transpose(input, mp_size, codegen_mp_num=4):
+    def _codegen_type_transpose(input, mp_size, codegen_mp_num=4, meta=meta):
         # codegen_mp_num defined in https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py
-        assert get_num_kv_heads() % (
+        assert meta.num_kv_heads % (
             mp_size * codegen_mp_num) == 0, "codgen autoTP requires num_kv_heads % (mp_size*codegen_mp_num) == 0"
         #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias)
 
         shape = input.shape
-        dst_shape = get_shard_size(shape[0], mp_size, rank=gpu_index)
+        dst_shape = get_shard_size(shape[0], mp_size, meta, rank=gpu_index)
         num_mp_blocks = input.reshape(codegen_mp_num, shape[0] // codegen_mp_num, shape[1])
 
         #num_mp_blocks : [codegen_mp_num, 3*hidden_dim/codegen_mp_num, :]
         src_split = list(torch.split(num_mp_blocks, num_mp_blocks.shape[1] // 3, dim=1))
         src_split = [x.reshape(codegen_mp_num * mp_size, -1, shape[1]) for x in src_split]
 
-        split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size(shape[0] // 3, mp_size, rank=gpu_index),
-                                                     0, 1)
+        split_fusedqkv = split_by_qkvlist_and_refuse(src_split,
+                                                     get_shard_size(shape[0] // 3, mp_size, meta, rank=gpu_index), 0,
+                                                     1)
         tp_fuseqkv_weight = torch.cat(split_fusedqkv, dim=0).reshape(shape[0], -1)
 
         return tp_fuseqkv_weight[gpu_index * dst_shape:(gpu_index + 1) * dst_shape]
 
-    def _glm_type_transpose(input, mp_size):
+    def _glm_type_transpose(input, mp_size, meta=meta):
         #input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias)
 
         # For chatglm2 & chatglm3(kv_heads=2), need to special handle.
-        if get_num_kv_heads() == 2:
+        if meta.num_kv_heads == 2:
             shape = input.shape
-            hidden_dim = get_n_embd()
-            kv_dim = (shape[0] - hidden_dim) // get_num_kv_heads()
+            hidden_dim = meta.n_embd
+            kv_dim = (shape[0] - hidden_dim) // meta.num_kv_heads
             q = input[:hidden_dim]
             k = input[hidden_dim:hidden_dim + kv_dim]
             v = input[hidden_dim + kv_dim:]
-            q_split = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0)
-            k_split = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0)
-            v_split = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0)
+            q_split = q.split(get_shard_size_list(q.shape[0], mp_size, meta), dim=0)
+            k_split = k.split(get_shard_size_list(k.shape[0], mp_size, meta), dim=0)
+            v_split = v.split(get_shard_size_list(v.shape[0], mp_size, meta), dim=0)
             return torch.cat((q_split[gpu_index], k_split[gpu_index], v_split[gpu_index]), dim=0)
         else:
             shape = input.shape
             src_split = torch.split(input, shape[0] // 3, dim=0)
 
-            split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size))
+            split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size, meta))
             return split_fusedqkv[gpu_index]
 
-    def _bloom_type_transpose(input, mp_size):
+    def _bloom_type_transpose(input, mp_size, meta=meta):
         shape = input.shape
 
-        split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size), dim=0)
+        split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size, meta), dim=0)
         return split_fusedqkv[gpu_index]
 
-    def _bigcode_type_transpose(input, mp_size):
-        n_embd = get_n_embd()
+    def _bigcode_type_transpose(input, mp_size, meta=meta):
+        n_embd = meta.n_embd
         q = input[:n_embd]
         kv = input[n_embd:]
         shape = q.shape
-        split_q = q.split(get_shard_size_list(shape[0], mp_size), dim=0)
+        split_q = q.split(get_shard_size_list(shape[0], mp_size, meta), dim=0)
         return torch.cat((split_q[gpu_index], kv), dim=0)
 
-    def _phi3_type_transpose(input, mp_size):
-        num_kv_heads = get_num_kv_heads()
-        num_heads = get_num_attention_heads()
+    def _phi3_type_transpose(input, mp_size, meta=meta):
+        num_kv_heads = meta.num_kv_heads
+        num_heads = meta.num_attention_heads
         hidden_size = input.shape[1]
         head_dim = hidden_size // num_heads
         q_pos = input.shape[0] - 2 * num_kv_heads * head_dim
         q = input[:q_pos]
         k = input[q_pos:q_pos + num_kv_heads * head_dim]
         v = input[q_pos + num_kv_heads * head_dim:]
-        split_q = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0)
-        split_k = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0)
-        split_v = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0)
+        split_q = q.split(get_shard_size_list(q.shape[0], mp_size, meta), dim=0)
+        split_k = k.split(get_shard_size_list(k.shape[0], mp_size, meta), dim=0)
+        split_v = v.split(get_shard_size_list(v.shape[0], mp_size, meta), dim=0)
         return torch.cat((split_q[gpu_index], split_k[gpu_index], split_v[gpu_index]), dim=0)
 
     def _transpose_fused_qkvw(src, mp_size, fused_qkv_type=None, module=None):
@@ -216,15 +217,15 @@ def shard_value_with_share_qk(
         bias,
         rank,
         world_size,
-        shard_value=True  # True -> shard_value; False -> shard_oproj
-):
+        shard_value,  # True -> shard_value; False -> shard_oproj
+        meta: AutoTPMeta):
     if shard_value:
         total_size = weight.shape[0]
         weight_cat_dim = 0
     else:
         total_size = weight.shape[1]
         weight_cat_dim = 1
-    num_heads = get_num_kv_heads()
+    num_heads = meta.num_kv_heads
     head_dim = total_size // num_heads
     assert (num_heads % world_size == 0)
     if world_size > num_heads // 2:
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 22623b42c73c..61a5dd98420c 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -11,7 +11,7 @@
 from torch.nn import functional as F
 from torch.nn.parameter import Parameter
 from deepspeed.accelerator import get_accelerator
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 from deepspeed.utils.logging import log_dist_once
 from deepspeed.runtime.zero.utils import is_zero_param
 from abc import ABC, abstractmethod
@@ -348,6 +348,10 @@ def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any):
         if kwargs.get('name') is not None:
             self.name = kwargs.get('name')  # Set the layer name if provided.
 
+        # Per-model TP metadata threaded from AutoTP; defaults for layers built outside it
+        # (e.g. the from_weights back-compat constructor).
+        self.tp_meta: AutoTPMeta = kwargs.get('tp_meta') or AutoTPMeta()
+
     @classmethod
     def set_keep_module_on_host(cls, value: bool):
         """
@@ -473,13 +477,11 @@ def is_training_mode(self):
     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.
+        The split depends on this model's kv-head/grain metadata (``self.tp_meta``), 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 re-deriving it.
         """
-        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.tp_meta, self.name))
         return self._partition_sizes
 
     @torch.no_grad()
@@ -964,7 +966,8 @@ def _freeze_partition_sizes(self, total_size):
         self._subparam_sizes = subparam_sizes
         self._subparam_shard_widths = None
         if subparam_sizes is not None:
-            self._subparam_shard_widths = _subparam_shard_widths(subparam_sizes, self.tp_world_size, shard_name)
+            self._subparam_shard_widths = _subparam_shard_widths(subparam_sizes, self.tp_world_size, self.tp_meta,
+                                                                 shard_name)
 
     def _subparam_shape_spec(self, logical_shape):
         shape_spec = list(logical_shape)
@@ -1084,8 +1087,8 @@ def __init__(self, module, mp_group, skip_partition=False, **kwargs):
         self.fused_module = FusedModuleWrapper(kwargs.get('fused_module'))
         # prepare_tp_fused_qkvw takes its own shard sizes without a layer name, so the widths
         # describing its split must be resolved the same way.
-        self._subparam_layout_spec = (fused_qkv_subparam_sizes(kwargs.get('fused_module'),
-                                                               tuple(module.weight.shape)), None)
+        self._subparam_layout_spec = (fused_qkv_subparam_sizes(kwargs.get('fused_module'), tuple(module.weight.shape),
+                                                               kwargs.get('tp_meta') or AutoTPMeta()), None)
         super().__init__(module, mp_group, skip_partition, **kwargs)
 
     def _freeze_partition_sizes(self, total_size):
@@ -1101,7 +1104,8 @@ def _tp_partition_unsupported_layout(self, params_list):
             if param is None:
                 return
 
-            _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index)
+            _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index,
+                                               self.tp_meta)
 
             _partition = self.move(_partition).detach()
 
@@ -1118,13 +1122,15 @@ def _tp_partition(self, params_list):
             weight = params_list[0]
         elif len(params_list) == 2:
             weight, bias = params_list[0], params_list[1]
-        _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name),
+        _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.tp_meta,
+                                                           self.name),
                                        dim=1)[self.tp_index]
         _partition = self.move(_partition).detach()
         weight.data = _partition
 
         if bias is not None:
-            _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name),
+            _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.tp_meta,
+                                                             self.name),
                                          dim=0)[self.tp_index]
             _partition = self.move(_partition).detach()
 
@@ -1141,7 +1147,7 @@ class Yuan_LinearAllreduce(LinearAllreduce):
     @torch.no_grad()
     def _tp_partition(self, params_list):
         weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
-                                                 self.tp_world_size, False)
+                                                 self.tp_world_size, False, self.tp_meta)
         params_list[0].data = weight
         if bias is not None:
             params_list[1].data = bias
@@ -1175,7 +1181,7 @@ class Yuan_LinearLayer(LinearLayer):
     @torch.no_grad()
     def _tp_partition(self, params_list):
         weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
-                                                 self.tp_world_size, True)
+                                                 self.tp_world_size, True, self.tp_meta)
         params_list[0].data = self.move(weight).detach()
         if bias is not None:
             params_list[1].data = self.move(bias).detach()
@@ -1220,7 +1226,7 @@ def _tp_partition(self, params_list):
                 return
             param.data = param.data.transpose(-1, -2).contiguous()
 
-            _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name),
+            _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.tp_meta, self.name),
                                      dim=1)[self.tp_index]
 
             _partition = self.move(_partition).detach()
@@ -1517,7 +1523,7 @@ def _bias_subparam_shape_spec(output_shape, bias_partition_dim, subparam_sizes):
     return tuple(shape_spec)
 
 
-def _subparam_shard_widths(subparam_sizes, tp_world_size, name=None):
+def _subparam_shard_widths(subparam_sizes, tp_world_size, meta: AutoTPMeta, name=None):
     """Per-rank width of each sub-parameter, as one list per sub-parameter.
 
     Sub-parameters follow the same deterministic split as ordinary layers, so a fused
@@ -1525,7 +1531,7 @@ def _subparam_shard_widths(subparam_sizes, tp_world_size, name=None):
     """
     widths = []
     for size in subparam_sizes:
-        per_rank = get_shard_size_list(size, tp_world_size, name)
+        per_rank = get_shard_size_list(size, tp_world_size, meta, name)
         if min(per_rank) == 0:
             # Those ranks contribute zeros to the row-parallel all-reduce, so the result stays
             # correct and this matches how separate q/k/v projections already behave. Serving
@@ -1664,9 +1670,10 @@ def __init__(self, module, mp_group, shape, partition_dim=0, **kwargs):
          self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                     self.partition_dim, self.name)
         # Resolve the per-rank widths once, for the same reason _freeze_partition_sizes does:
-        # get_shard_size_list reads process-wide globals that a later model can overwrite.
+        # the split depends on this model's tp_meta.
         self._subparam_shard_widths = _subparam_shard_widths(
-            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.name)
+            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.tp_meta,
+            self.name)
         self._bias_shape_spec = _bias_subparam_shape_spec(self._output_shape, self._bias_partition_dim,
                                                           self._subparam_sizes)
         if self.bias is not None and self.bias.numel() != _shape_prod(self._output_shape):
@@ -1796,9 +1803,10 @@ def __init__(self, module, mp_group, shape, partition_dim=1, **kwargs):
          self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                     self.partition_dim, self.name)
         # Resolve the per-rank widths once, for the same reason _freeze_partition_sizes does:
-        # get_shard_size_list reads process-wide globals that a later model can overwrite.
+        # the split depends on this model's tp_meta.
         self._subparam_shard_widths = _subparam_shard_widths(
-            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.name)
+            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.tp_meta,
+            self.name)
 
         if self._should_materialize_tp_partition():
             self._tp_partition([self.weight, self.bias])
diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py
index 263369fc0484..4af88b696062 100644
--- a/deepspeed/module_inject/replace_module.py
+++ b/deepspeed/module_inject/replace_module.py
@@ -17,7 +17,6 @@
 from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d
 from deepspeed.module_inject.layers import is_autotp_training_mode
 from deepspeed import comm as dist
-from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads, set_tp_grain_size
 
 from .load_checkpoint import load_model_with_checkpoint
 import time
@@ -278,7 +277,19 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
         if hasattr(config, 'get_partition_config_object'):
             partition_config = config.get_partition_config_object()
 
-        # 1. Create AutoTP object
+        # 1. Resolve which sub-config describes the layers being replaced (e.g. Mllama vision
+        # vs text); AutoTP reads its TP metadata from it.
+        if hasattr(model_config, "vision_config"):
+            if "MllamaVisionEncoderLayer" in str(module):
+                meta_config = model_config.vision_config
+            elif hasattr(model_config, "text_config"):
+                meta_config = model_config.text_config
+            else:
+                meta_config = model_config
+        else:
+            meta_config = model_config
+
+        # 2. Create AutoTP object; it derives its tp_meta from meta_config.
         _autotp = AutoTP(module,
                          all_reduce_linears,
                          prefix,
@@ -286,45 +297,14 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
                          linear_layer_setting,
                          orig_layer_impl,
                          config.keep_module_on_host,
-                         partition_config=partition_config)
+                         partition_config=partition_config,
+                         model_config=meta_config,
+                         tp_grain_size=config.tensor_parallel.tp_grain_size)
 
-        # 2. Set the tensor parallelism config
+        # 3. Set the tensor parallelism config
         _autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group)
 
-        # 3. Try to get num_key_heads from model_config.num_key_value_heads
-        if hasattr(model_config, "vision_config"):
-            if "MllamaVisionEncoderLayer" in str(module):
-                num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config)
-            elif hasattr(model_config, "text_config"):
-                num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config)
-            else:
-                num_kv_heads = _autotp.get_model_num_kv_heads(model_config)
-        else:
-            num_kv_heads = _autotp.get_model_num_kv_heads(model_config)
-
-        # 4. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
-        set_num_kv_heads(num_kv_heads)
-
-        # 4.1 Get n_embd
-        n_embd = None
-        multi_query_n_embd_names = ['n_embd', 'hidden_size']
-        for name in multi_query_n_embd_names:
-            if hasattr(model_config, name):
-                n_embd = getattr(model_config, name)
-            if n_embd != None:
-                break
-
-        # 4.2 set n_embd
-        set_n_embd(n_embd)
-
-        # 4.3 set attention_heads
-        if hasattr(model_config, 'num_attention_heads'):
-            set_num_attention_heads(getattr(model_config, 'num_attention_heads'))
-
-        # 4.4 set tp_grain_size
-        set_tp_grain_size(config.tensor_parallel.tp_grain_size)
-
-        # 5. Set linear policies
+        # 4. Set linear policies
         _autotp.update_linear_policies()
 
         # 6. Replace modules
diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index f976ff22bc33..6d092f21c5dd 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -3,69 +3,80 @@
 
 # DeepSpeed Team
 
-from deepspeed import comm as dist
-
-# Defaults for optional TP globals. These can be overridden by setters.
-num_kv_heads = None
-num_attention_heads = None
-n_embd = None
-tp_grain_size = 1
-
-
-def set_num_kv_heads(num):
-    global num_kv_heads
-    num_kv_heads = num
-
-
-def set_num_attention_heads(num):
-    global num_attention_heads
-    num_attention_heads = num
-
-
-def set_n_embd(num):
-    global n_embd
-    n_embd = num
-
-
-def set_tp_grain_size(num):
-    global tp_grain_size
-    tp_grain_size = num
-
-
-def get_num_kv_heads():
-    global num_kv_heads
-    if 'num_kv_heads' in globals():
-        return num_kv_heads
-    return None
+from dataclasses import dataclass
+from typing import Optional
 
+from deepspeed import comm as dist
 
-def get_num_attention_heads():
-    global num_attention_heads
-    return num_attention_heads
 
+@dataclass(frozen=True)
+class AutoTPMeta:
+    """Per-model tensor-parallel metadata AutoTP derives from the model config.
 
-def get_shard_size(total_size, mp_size, name=None, rank=None, mp_group=None, num_kv_heads=None):
+    Each model carries its own kv-head / grain values so its layers, fused-QKV repacking and
+    checkpoint conversion shard consistently, and more than one AutoTP model can live in the
+    same process.
+    """
+    num_kv_heads: Optional[int] = None
+    num_attention_heads: Optional[int] = None
+    n_embd: Optional[int] = None
+    tp_grain_size: int = 1
+
+    @classmethod
+    def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta":
+        """The single source of truth for reading kv-head / hidden / attention-head counts.
+
+        ``model_config`` may be ``None`` for callers that build a synthetic AutoTP with no real
+        model (some unit tests); they get an empty meta and the resulting even-grain split.
+        """
+        if model_config is None:
+            return cls(tp_grain_size=tp_grain_size)
+        # multi_query_group_num is for chatglm2 & chatglm3
+        num_kv_heads = None
+        for name in ('multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads',
+                     'attention_heads'):
+            if hasattr(model_config, name):
+                num_kv_heads = getattr(model_config, name)
+                if num_kv_heads is not None:
+                    break
+        n_embd = None
+        for name in ('n_embd', 'hidden_size'):
+            if hasattr(model_config, name):
+                n_embd = getattr(model_config, name)
+            if n_embd is not None:
+                break
+        num_attention_heads = getattr(model_config, 'num_attention_heads', None)
+        return cls(num_kv_heads=num_kv_heads,
+                   num_attention_heads=num_attention_heads,
+                   n_embd=n_embd,
+                   tp_grain_size=tp_grain_size)
+
+
+def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None, num_kv_heads=None):
     """Size of one shard of ``total_size`` split across a tensor-parallel group of ``mp_size``.
 
+    ``meta`` carries this model's ``num_kv_heads`` / ``tp_grain_size`` so the split is stable
+    for the lifetime of the model instead of depending on whichever AutoTP model was loaded
+    last.
+
     ``rank`` is the rank *within the tensor-parallel group*, i.e. in ``[0, mp_size)``, matching
     ``dist.get_rank(group=mp_group)`` and the index used by ``get_shard_size_list``. It is not a
     global rank.
 
-    ``num_kv_heads`` overrides the module-level global of the same name when set. Passing it
-    explicitly lets callers split fused sub-parameters (Q/K/V) against their respective head
-    counts without reimplementing the KV-head-aligned partition logic at the call site. When
-    ``None`` (default), the function falls back to the global ``num_kv_heads`` set via
-    ``set_num_kv_heads`` for backward compatibility.
+    ``num_kv_heads`` overrides ``meta.num_kv_heads`` when set. Passing it explicitly lets callers
+    split fused sub-parameters (Q/K/V) against their respective head counts without
+    reimplementing the KV-head-aligned partition logic at the call site.
     """
     if num_kv_heads is None:
-        num_kv_heads = globals()["num_kv_heads"]
+        num_kv_heads = meta.num_kv_heads
+    tp_grain_size = meta.tp_grain_size
     last_linear = ["lm_head", "embed_out"]
     # MoE MLP layer use near even division will get better perf.
     moe_mlp_layer = ["gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"]
     not_moe_mlp_layer = True
     if name != None and any(s in str(name) for s in moe_mlp_layer):
         not_moe_mlp_layer = False
-    # When we have num_kv_heads defined, uneven division is possible, otherwise enforce near even division
+    # When num_kv_heads is defined, uneven division is possible, otherwise enforce near even division
     if rank is None:
         if mp_group is not None:
             rank = dist.get_rank(group=mp_group)
@@ -93,21 +104,16 @@ def get_shard_size(total_size, mp_size, name=None, rank=None, mp_group=None, num
             return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
 
-def get_n_embd():
-    global n_embd
-    return n_embd
-
-
-def get_shard_size_list(total_size, mp_size, name=None, num_kv_heads=None):
+def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None, num_kv_heads=None):
     shard_sizes = []
     if num_kv_heads is None:
-        num_kv_heads = globals()["num_kv_heads"]
+        num_kv_heads = meta.num_kv_heads
     for i in range(mp_size):
-        shard_sizes.append(get_shard_size(total_size, mp_size, name, i, num_kv_heads=num_kv_heads))
+        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i, num_kv_heads=num_kv_heads))
     # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
     # 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 "
-        f"{total_size} with tp_size={mp_size}, tp_grain_size={tp_grain_size} and "
+        f"{total_size} with tp_size={mp_size}, tp_grain_size={meta.tp_grain_size} and "
         f"num_kv_heads={num_kv_heads}.")
     return shard_sizes
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 18085d75005f..8d975caad2fa 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -726,20 +726,14 @@ def _apply_autotp_partitioning(self, model, tp_config):
             partition_config = tp_config.get_partition_config_object()
 
         model_config = getattr(model, "config", None)
-        # The direct Hugging Face tp_plan path bypasses replace_transformer_layer, which
-        # normally initializes the shard-size globals that AutoTP layers consult. Without
-        # them attention projections are split by grain size and can be cut mid-head, so
-        # the model's later reshape onto head_dim fails.
-        from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads
-        from deepspeed.module_inject.tp_shard import set_tp_grain_size
-
-        # 1. Try to get num_key_heads from model_config.num_key_value_heads
-        if hasattr(model_config, "text_config"):
-            num_kv_heads = AutoTP.get_model_num_kv_heads(model_config.text_config)
-        else:
-            num_kv_heads = AutoTP.get_model_num_kv_heads(model_config)
+        # AutoTP derives its per-model sharding metadata from the model config; the warning
+        # below only needs the kv-head count, which that same metadata already carries.
+        from deepspeed.module_inject.tp_shard import AutoTPMeta
+
+        head_config = getattr(model_config, "text_config", model_config)
+        num_kv_heads = AutoTPMeta.from_model_config(head_config).num_kv_heads
 
-        # 2. Ranks beyond the KV head count get no attention shard. This still computes the
+        # Ranks beyond the KV head count get no attention shard. This still computes the
         # correct result because the row-parallel all-reduce sums their empty contribution,
         # but attention work concentrates on the first num_kv_heads ranks.
         if num_kv_heads is not None and tp_size > num_kv_heads:
@@ -750,28 +744,6 @@ def _apply_autotp_partitioning(self, model, tp_config):
                 ranks=[0],
                 level=logging.WARNING)
 
-        # 3. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
-        set_num_kv_heads(num_kv_heads)
-
-        # 3.1 Get n_embd
-        n_embd = None
-        multi_query_n_embd_names = ['n_embd', 'hidden_size']
-        for name in multi_query_n_embd_names:
-            if hasattr(model_config, name):
-                n_embd = getattr(model_config, name)
-            if n_embd != None:
-                break
-
-        # 3.2 set n_embd
-        set_n_embd(n_embd)
-
-        # 3.3 set attention_heads
-        if hasattr(model_config, 'num_attention_heads'):
-            set_num_attention_heads(getattr(model_config, 'num_attention_heads'))
-
-        # 3.4 set tp_grain_size
-        set_tp_grain_size(tp_config.tensor_parallel.tp_grain_size)
-
         from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
         hf_tp_plan = _get_hf_tp_plan(model)
 
@@ -783,7 +755,9 @@ def _apply_autotp_partitioning(self, model, tp_config):
                             linear_layer_setting=(torch.nn.Linear, torch.nn.Embedding),
                             orig_layer_impl=None,
                             keep_module_on_host=tp_config.keep_module_on_host,
-                            partition_config=partition_config)
+                            partition_config=partition_config,
+                            model_config=model_config,
+                            tp_grain_size=tp_config.tensor_parallel.tp_grain_size)
             autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group)
             autotp.update_linear_policies()
             autotp._replace_module(model)
@@ -822,6 +796,8 @@ def _apply_autotp_partitioning(self, model, tp_config):
                     orig_layer_impl=None,
                     keep_module_on_host=tp_config.keep_module_on_host,
                     partition_config=tp_plan_config,
+                    model_config=model_config,
+                    tp_grain_size=tp_config.tensor_parallel.tp_grain_size,
                 )
                 autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group)
                 autotp.update_linear_policies()
diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py
index 1ae18c14ced8..2088f9686ac0 100644
--- a/deepspeed/sequence/layer.py
+++ b/deepspeed/sequence/layer.py
@@ -12,9 +12,27 @@
 
 import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
-from deepspeed.module_inject.tp_shard import get_shard_size_list, set_num_kv_heads, get_num_kv_heads
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 from deepspeed.utils import groups
 
+# Ulysses sequence parallelism keeps its own kv-head count, memoized on the first uneven
+# all-to-all. The state lives here, independent of AutoTP.
+_ulysses_num_kv_heads = None
+
+
+def set_ulysses_num_kv_heads(num):
+    global _ulysses_num_kv_heads
+    _ulysses_num_kv_heads = num
+
+
+def get_ulysses_num_kv_heads():
+    return _ulysses_num_kv_heads
+
+
+def _ulysses_meta():
+    return AutoTPMeta(num_kv_heads=_ulysses_num_kv_heads)
+
+
 try:
     from torchembed._triton import fused_rope_forward as _torchembed_rope_forward
     _torchembed_available = True
@@ -134,7 +152,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
     assert batch_dim_idx in [0, 1], "batch_dim_idx must be either 0 or 1"
 
     if not (scatter_idx < 2):
-        input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size)
+        input_splits = get_shard_size_list(inp_shape[scatter_idx], seq_world_size, _ulysses_meta())
         input = input.transpose(0, scatter_idx).contiguous()
         local_heads = input_splits[groups._get_sequence_parallel_rank()]
         output_splits = [local_heads] * seq_world_size
@@ -168,7 +186,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         elif batch_dim_idx == 1:  #s,b,h
             input = input.transpose(1, 2).contiguous()  #s,h,b
         seq_len, h, batch_size = input.shape
-        num_local_heads_list = get_shard_size_list(get_num_kv_heads(), seq_world_size)
+        num_local_heads_list = get_shard_size_list(get_ulysses_num_kv_heads(), seq_world_size, _ulysses_meta())
         local_heads = num_local_heads_list[groups._get_sequence_parallel_rank()]
         h_dim = h // local_heads
         local_seq_len = seq_len // seq_world_size
@@ -179,7 +197,7 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         coeff = local_seq_len_with_heads // local_heads  #per head: dim size of local_seq_len*hdim
 
         #uneven seq_world_size coeff,  total_heads/local_heads.
-        heads_scale_coeff = get_num_kv_heads() / local_heads
+        heads_scale_coeff = get_ulysses_num_kv_heads() / local_heads
 
         output_splits = [num_local_heads * coeff for num_local_heads in num_local_heads_list]
         output_buff_d1_size = int(heads_scale_coeff * local_seq_len_with_heads)
@@ -196,9 +214,9 @@ def uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group):
         #total_num_large_heads=sum([2,2,2])=7
         #total_num_small_heads=sum([1])=1
 
-        chunk_num_heads_small = get_num_kv_heads() // seq_world_size  # even heads compatible
+        chunk_num_heads_small = get_ulysses_num_kv_heads() // seq_world_size  # even heads compatible
         chunk_num_heads_large = chunk_num_heads_small + 1
-        num_chunk_heads_large = get_num_kv_heads() % seq_world_size
+        num_chunk_heads_large = get_ulysses_num_kv_heads() % seq_world_size
         num_chunk_heads_small = seq_world_size - num_chunk_heads_large
         total_num_large_heads = num_chunk_heads_large * chunk_num_heads_large
         total_num_small_heads = num_chunk_heads_small * chunk_num_heads_small
@@ -243,14 +261,14 @@ def single_all_to_all(input, scatter_idx, gather_idx, batch_dim_idx, group, asyn
     # we only need num_heads once
     num_heads = input.shape[2]
 
-    if get_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):
+    if get_ulysses_num_kv_heads() is not None or (num_heads % seq_world_size != 0 and not scatter_idx < 2):
         # Assuming here that the number of heads for q is consistent with kv
         # If not, additional logic is required for cases like GQA
-        if get_num_kv_heads() is None:
+        if get_ulysses_num_kv_heads() is None:
             assert num_heads > seq_world_size, f"Number of heads ({num_heads}) must be larger than sequence parallel size ({seq_world_size})"
             # set heads at first call by num_total_heads.
-            # then use ``get_num_kv_heads() is not None`` to re-entry uneven path.
-            set_num_kv_heads(num_heads)
+            # then use ``get_ulysses_num_kv_heads() is not None`` to re-entry uneven path.
+            set_ulysses_num_kv_heads(num_heads)
         assert async_op == False, "uneven head sp does not support async op"
         return uneven_heads_all2all(input, scatter_idx, gather_idx, batch_dim_idx, group)
 
diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
index efc85905cdba..af615fa2cc27 100644
--- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
+++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py
@@ -1105,7 +1105,7 @@ def test_save_convert_load_uneven_lm_head(self, tmpdir):
 
         tp_group = groups.get_tensor_model_parallel_group()
         tp_rank = groups.get_tensor_model_parallel_rank()
-        shard_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        shard_sizes = list(engine.module.lm_head._partition_sizes)
         assert shard_sizes == [51, 50], shard_sizes
         assert engine.module.lm_head.weight.shape[0] == shard_sizes[tp_rank]
 
@@ -1194,7 +1194,7 @@ def test_save_convert_load_uneven_row_parallel(self, tmpdir):
         # Column and row parallelism must shard the same dimension identically.
         assert attn.q_proj.weight.shape[0] == attn.o_proj.weight.shape[1]
         head_dim = hidden_dim // num_heads
-        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.name)
+        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.tp_meta, attn.q_proj.name)
         assert head_shards == [2, 2, 1, 1], head_shards
         dim_shards = [h * head_dim for h in head_shards]
         assert attn.q_proj.weight.shape[0] == dim_shards[tp_rank], (attn.q_proj.weight.shape, dim_shards)
@@ -1383,7 +1383,7 @@ def test_strict_correctness(self, tmpdir):
 
         tp_group = groups.get_tensor_model_parallel_group()
         tp_rank = groups.get_tensor_model_parallel_rank()
-        shard_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        shard_sizes = list(engine.module.lm_head._partition_sizes)
         assert sum(shard_sizes) == vocab_size, shard_sizes
         assert shard_sizes == [4, 3, 3], shard_sizes
         assert engine.module.lm_head.weight.shape[0] == shard_sizes[tp_rank]
@@ -1506,7 +1506,7 @@ def test_strict_correctness(self, tmpdir):
         assert attn.q_proj.weight.shape[0] == attn.o_proj.weight.shape[1], (attn.q_proj.weight.shape,
                                                                             attn.o_proj.weight.shape)
         head_dim = hidden_dim // num_heads
-        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.name)
+        head_shards = get_shard_size_list(num_heads, self.world_size, attn.q_proj.tp_meta, attn.q_proj.name)
         assert head_shards == [2, 1, 1], head_shards
         dim_shards = [h * head_dim for h in head_shards]
         assert attn.q_proj.weight.shape[0] == dim_shards[tp_rank], (attn.q_proj.weight.shape, dim_shards)
diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
index e47a7f7d9f1d..a381f57f9720 100644
--- a/tests/unit/model_parallelism/test_autotp_custom_patterns.py
+++ b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
@@ -18,7 +18,7 @@
 from deepspeed.module_inject.layers import collect_autotp_universal_checkpoint_info
 from deepspeed.checkpoint.constants import PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, TP_REPLICATED_PARAMETER_PATTERNS
 from deepspeed.module_inject.autotp_config import AutoTPConfig
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, set_num_kv_heads
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 from deepspeed.module_inject.auto_tp import AutoTP
 from deepspeed.module_inject.auto_tp_model_utils import (build_bloom_alibi_tensor, build_mpt_alibi_tensor,
                                                          get_alibi_mask, install_head_sharded_helper)
@@ -169,7 +169,8 @@ def apply_autotp_with_partition_config(model, tp_size, partition_config):
                     linear_layer_setting=None,
                     orig_layer_impl=None,
                     keep_module_on_host=False,
-                    partition_config=autotp_config)
+                    partition_config=autotp_config,
+                    model_config=getattr(model, "config", None))
     autotp.set_tensor_parallel_config(tp_size, groups.get_tensor_model_parallel_group())
     autotp.update_linear_policies()
     autotp._replace_module(model)
@@ -499,11 +500,8 @@ def test_update_mp_params_uses_group_local_rank(monkeypatch):
     child.num_heads = 12
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(3)
-    try:
-        autotp.update_mp_params(child)
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta(num_kv_heads=3)
+    autotp.update_mp_params(child)
 
     # Three KV groups split as [2, 1], so the second TP rank owns four query heads.
     assert child.num_heads == 4
@@ -518,15 +516,13 @@ def test_update_mp_params_shards_attributes_like_their_weights(monkeypatch):
     child.hidden_size = 12
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(3)
-    try:
-        autotp.update_mp_params(child, "model.layers.0.mlp")
-    finally:
-        set_num_kv_heads(None)
+    meta = AutoTPMeta(num_kv_heads=3)
+    autotp.tp_meta = meta
+    autotp.update_mp_params(child, "model.layers.0.mlp")
 
     # MLP layers are excluded from the KV-head split, so the attribute has to follow the same
     # near-even split that the MLP weights use rather than the [2, 1] KV-group split.
-    assert child.hidden_size == get_shard_size(12, 2, "model.layers.0.mlp", rank=1)
+    assert child.hidden_size == get_shard_size(12, 2, meta, "model.layers.0.mlp", rank=1)
 
 
 def test_sliced_embedding_publishes_row_partition_metadata(monkeypatch):
@@ -534,6 +530,7 @@ def test_sliced_embedding_publishes_row_partition_metadata(monkeypatch):
     autotp = object.__new__(AutoTP)
     autotp.mp_group = tp_group
     autotp.mp_size = 2
+    autotp.tp_meta = AutoTPMeta()
     embedding = nn.Embedding(5, 4)
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
@@ -556,66 +553,74 @@ def test_mpt_alibi_covers_every_head_of_an_uneven_split(self):
         init_tp_engine(tp_size=2)
 
         num_heads = 5
-        set_num_kv_heads(num_heads)
-        try:
 
-            class MptTransformer(nn.Module):
+        class MptTransformer(nn.Module):
 
-                def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
-                    return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
+            def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
+                return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
 
-            transformer = MptTransformer()
-            install_head_sharded_helper(transformer, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
+        transformer = MptTransformer()
+        install_head_sharded_helper(transformer,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    num_heads=num_heads,
+                                    num_kv_heads=num_heads)
 
-            alibi = transformer.build_mpt_alibi_tensor(num_heads, 3)
+        alibi = transformer.build_mpt_alibi_tensor(num_heads, 3)
 
-            # AutoTP splits 5 heads over 2 ranks as [3, 2]; an even split would give every rank
-            # 2 heads and drop the last one entirely.
-            expected_heads = get_shard_size_list(num_heads, dist.get_world_size())
-            offset = sum(expected_heads[:dist.get_rank()])
-            assert alibi.shape[0] == expected_heads[dist.get_rank()]
-            torch.testing.assert_close(alibi[:, 0, 0].cpu(),
-                                       torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
-        finally:
-            set_num_kv_heads(None)
+        # AutoTP splits 5 heads over 2 ranks as [3, 2]; an even split would give every rank
+        # 2 heads and drop the last one entirely.
+        expected_heads = get_shard_size_list(num_heads, dist.get_world_size(), AutoTPMeta(num_kv_heads=num_heads))
+        offset = sum(expected_heads[:dist.get_rank()])
+        assert alibi.shape[0] == expected_heads[dist.get_rank()]
+        torch.testing.assert_close(alibi[:, 0, 0].cpu(),
+                                   torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
 
     def test_head_sharded_helper_leaves_the_class_untouched(self):
         skip_on_device()
         init_tp_engine(tp_size=2)
 
         num_heads = 5
-        set_num_kv_heads(num_heads)
-        try:
-
-            class MptTransformer(nn.Module):
 
-                def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
-                    return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
-
-            class MptSubclass(MptTransformer):
-                pass
-
-            first = MptTransformer()
-            install_head_sharded_helper(first, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
-            expected = first.build_mpt_alibi_tensor(num_heads, 3)
-
-            # Injecting a second model of the same architecture must not make either of them
-            # delegate to the other's wrapper.
-            second = MptTransformer()
-            install_head_sharded_helper(second, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
-            torch.testing.assert_close(second.build_mpt_alibi_tensor(num_heads, 3), expected)
-            torch.testing.assert_close(first.build_mpt_alibi_tensor(num_heads, 3), expected)
+        class MptTransformer(nn.Module):
 
-            # A model of the same class that was never injected keeps its own method, and a
-            # subclass of it inherits that method rather than an installed wrapper.
-            plain = MptTransformer()
-            assert plain.build_mpt_alibi_tensor(num_heads, 3).shape[0] == num_heads
+            def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, device=None):
+                return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
 
-            derived = MptSubclass()
-            install_head_sharded_helper(derived, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor)
-            torch.testing.assert_close(derived.build_mpt_alibi_tensor(num_heads, 3), expected)
-        finally:
-            set_num_kv_heads(None)
+        class MptSubclass(MptTransformer):
+            pass
+
+        first = MptTransformer()
+        install_head_sharded_helper(first,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    num_heads=num_heads,
+                                    num_kv_heads=num_heads)
+        expected = first.build_mpt_alibi_tensor(num_heads, 3)
+
+        # Injecting a second model of the same architecture must not make either of them
+        # delegate to the other's wrapper.
+        second = MptTransformer()
+        install_head_sharded_helper(second,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    num_heads=num_heads,
+                                    num_kv_heads=num_heads)
+        torch.testing.assert_close(second.build_mpt_alibi_tensor(num_heads, 3), expected)
+        torch.testing.assert_close(first.build_mpt_alibi_tensor(num_heads, 3), expected)
+
+        # A model of the same class that was never injected keeps its own method, and a
+        # subclass of it inherits that method rather than an installed wrapper.
+        plain = MptTransformer()
+        assert plain.build_mpt_alibi_tensor(num_heads, 3).shape[0] == num_heads
+
+        derived = MptSubclass()
+        install_head_sharded_helper(derived,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    num_heads=num_heads,
+                                    num_kv_heads=num_heads)
+        torch.testing.assert_close(derived.build_mpt_alibi_tensor(num_heads, 3), expected)
 
     def test_head_sharded_helper_freezes_the_models_split(self):
         skip_on_device()
@@ -627,29 +632,24 @@ def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, devic
                 return torch.arange(heads, dtype=torch.float32).view(heads, 1, 1).expand(heads, 1, sequence_length)
 
         num_heads = 6
-        set_num_kv_heads(3)
-        try:
-            transformer = MptTransformer()
-            install_head_sharded_helper(transformer,
-                                        'build_mpt_alibi_tensor',
-                                        build_mpt_alibi_tensor,
-                                        num_heads=num_heads,
-                                        num_kv_heads=3)
-
-            # Initializing another model can replace this process-wide setting. The first
-            # model's helper must keep the [4, 2] split frozen with its weights.
-            set_num_kv_heads(2)
-            expected_sizes = [4, 2]
-            # AutoTP replaces the model's public head count with this rank's local count.
-            local_num_heads = expected_sizes[dist.get_rank()]
-            alibi = transformer.build_mpt_alibi_tensor(local_num_heads, 3)
-
-            offset = sum(expected_sizes[:dist.get_rank()])
-            assert alibi.shape[0] == expected_sizes[dist.get_rank()]
-            torch.testing.assert_close(alibi[:, 0, 0].cpu(),
-                                       torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
-        finally:
-            set_num_kv_heads(None)
+        transformer = MptTransformer()
+        install_head_sharded_helper(transformer,
+                                    'build_mpt_alibi_tensor',
+                                    build_mpt_alibi_tensor,
+                                    num_heads=num_heads,
+                                    num_kv_heads=3)
+
+        # The helper freezes the [4, 2] split from this model's own num_kv_heads at install
+        # time.
+        expected_sizes = [4, 2]
+        # AutoTP replaces the model's public head count with this rank's local count.
+        local_num_heads = expected_sizes[dist.get_rank()]
+        alibi = transformer.build_mpt_alibi_tensor(local_num_heads, 3)
+
+        offset = sum(expected_sizes[:dist.get_rank()])
+        assert alibi.shape[0] == expected_sizes[dist.get_rank()]
+        torch.testing.assert_close(alibi[:, 0, 0].cpu(),
+                                   torch.arange(offset, offset + alibi.shape[0], dtype=torch.float32))
 
     def test_bloom_alibi_uses_original_total_after_injection(self):
         skip_on_device()
@@ -711,24 +711,22 @@ def test_gate_up_gather_restores_sub_param_order(self):
 
         hidden_dim = 4
         gate_up_dim = 6
-        set_num_kv_heads(3)
-        try:
-            torch.manual_seed(17)
-            linear = nn.Linear(hidden_dim,
-                               gate_up_dim * 2,
-                               bias=False,
-                               dtype=preferred_dtype(),
-                               device=get_accelerator().current_device_name())
-            full_weight = deepcopy(linear.weight.data)
-            layer = GateUpPack_LinearLayer(deepcopy(linear), groups.get_tensor_model_parallel_group())
-
-            # The gate and the up halves are each cut in two, so a rank-order concatenation of
-            # the shards would interleave them instead of restoring the original weight.
-            gathered = nn.Parameter(layer.weight.data.clone())
-            layer.gather_params([gathered, None])
-            torch.testing.assert_close(gathered.data, full_weight)
-        finally:
-            set_num_kv_heads(None)
+        torch.manual_seed(17)
+        linear = nn.Linear(hidden_dim,
+                           gate_up_dim * 2,
+                           bias=False,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device_name())
+        full_weight = deepcopy(linear.weight.data)
+        layer = GateUpPack_LinearLayer(deepcopy(linear),
+                                       groups.get_tensor_model_parallel_group(),
+                                       tp_meta=AutoTPMeta(num_kv_heads=3))
+
+        # The gate and the up halves are each cut in two, so a rank-order concatenation of
+        # the shards would interleave them instead of restoring the original weight.
+        gathered = nn.Parameter(layer.weight.data.clone())
+        layer.gather_params([gathered, None])
+        torch.testing.assert_close(gathered.data, full_weight)
 
     def test_gate_up_fused_weight_partition(self):
         skip_on_device()
@@ -882,29 +880,26 @@ def test_gather_uses_the_layers_own_shard_widths(self):
 
         hidden_dim = 8
         head_size = 12
-        set_num_kv_heads(3)
-        try:
-            torch.manual_seed(7)
-            linear = nn.Linear(hidden_dim,
-                               head_size * 3,
-                               bias=True,
-                               dtype=preferred_dtype(),
-                               device=get_accelerator().current_device_name())
-            full_weight = deepcopy(linear.weight.data)
-            full_bias = deepcopy(linear.bias.data)
-
-            layer = SubParamLinearLayer(deepcopy(linear),
-                                        groups.get_tensor_model_parallel_group(),
-                                        shape=((head_size, head_size, head_size), -1),
-                                        partition_dim=0,
-                                        name="self_attn.qkv_proj")
-            assert layer._subparam_shard_widths == [[8, 4], [8, 4], [8, 4]]
-
-            layer.gather_params([layer.weight, layer.bias])
-            torch.testing.assert_close(layer.weight.data, full_weight)
-            torch.testing.assert_close(layer.bias.data, full_bias)
-        finally:
-            set_num_kv_heads(None)
+        torch.manual_seed(7)
+        linear = nn.Linear(hidden_dim,
+                           head_size * 3,
+                           bias=True,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device_name())
+        full_weight = deepcopy(linear.weight.data)
+        full_bias = deepcopy(linear.bias.data)
+
+        layer = SubParamLinearLayer(deepcopy(linear),
+                                    groups.get_tensor_model_parallel_group(),
+                                    shape=((head_size, head_size, head_size), -1),
+                                    partition_dim=0,
+                                    name="self_attn.qkv_proj",
+                                    tp_meta=AutoTPMeta(num_kv_heads=3))
+        assert layer._subparam_shard_widths == [[8, 4], [8, 4], [8, 4]]
+
+        layer.gather_params([layer.weight, layer.bias])
+        torch.testing.assert_close(layer.weight.data, full_weight)
+        torch.testing.assert_close(layer.bias.data, full_bias)
 
     def test_gqa_uneven_qkv_fused_forward(self):
         skip_on_device()
diff --git a/tests/unit/model_parallelism/test_autotp_training.py b/tests/unit/model_parallelism/test_autotp_training.py
index f694ef18efcf..6cf2e35e7bed 100644
--- a/tests/unit/model_parallelism/test_autotp_training.py
+++ b/tests/unit/model_parallelism/test_autotp_training.py
@@ -17,7 +17,6 @@
 from contextlib import contextmanager
 from torch import nn
 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
@@ -502,7 +501,7 @@ def run_tp_layer_fwd_bwd(tp_size,
         loss.backward()
 
         expected_out = torch_out
-        output_partition_sizes = get_shard_size_list(torch_out.shape[-1], tp_size, linear.name)
+        output_partition_sizes = list(linear._partition_sizes)
         tp_rank = groups.get_tensor_model_parallel_rank()
         if not gather_output:
             shard_offset = sum(output_partition_sizes[:tp_rank])
@@ -689,7 +688,7 @@ def test_uneven_linear_gather_params(self):
         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)
+        output_partition_sizes = list(tp_layer._partition_sizes)
         expected_tp_params = output_partition_sizes[tp_rank] * (hidden_dim + 1)
 
         assert expected_tp_params == sum(p.numel() for p in tp_layer.parameters())
@@ -778,7 +777,7 @@ def test_consolidated_checkpoint(self):
         engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
 
         tp_rank = groups.get_tensor_model_parallel_rank()
-        output_partition_sizes = get_shard_size_list(vocab_size, self.world_size, "lm_head")
+        output_partition_sizes = list(engine.module.lm_head._partition_sizes)
         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)
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 b21a57c5d9c7..aa20f6933e47 100644
--- a/tests/unit/model_parallelism/test_tp_plan_real_models.py
+++ b/tests/unit/model_parallelism/test_tp_plan_real_models.py
@@ -9,7 +9,6 @@
 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
@@ -129,7 +128,7 @@ def test_qwen2_tp_plan_with_uneven_vocab(self):
         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")
+        output_partition_sizes = list(model.lm_head._partition_sizes)
         assert engine.autotp_size() == 2
         assert isinstance(model.lm_head, LinearLayer)
         assert model.lm_head.gather_output
diff --git a/tests/unit/module_inject/test_fused_repartition.py b/tests/unit/module_inject/test_fused_repartition.py
index f971e4ebf65f..972b28703076 100644
--- a/tests/unit/module_inject/test_fused_repartition.py
+++ b/tests/unit/module_inject/test_fused_repartition.py
@@ -2,39 +2,35 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
-"""Fused AutoTP layers must keep the shard widths that were frozen when they were built.
+"""Fused AutoTP layers keep the shard widths frozen when they were built.
 
-``tp_shard`` keeps the grain size and the kv head count in process-wide globals that every
-AutoTP model overwrites while it is being replaced.  A layer that re-derives its split at
-partition time therefore cuts the weight differently once a second model has been loaded,
-which silently disagrees with the gather and with the universal checkpoint metadata.
+Each AutoTP model carries its own :class:`AutoTPMeta` (kv-head count, grain size, ...), so a
+layer's split stays aligned with its gather and its universal checkpoint metadata regardless
+of what other models do.  These tests build layers with an explicit ``tp_meta`` and check the
+split is stable.
 """
 
 import pytest
 import torch
 
-from deepspeed.module_inject import tp_shard
 from deepspeed.module_inject.layers import GateUpPack_LinearLayer, fused_LinearLayer
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 
 
-@pytest.fixture
-def clean_tp_shard_globals():
-    yield
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(None)
-
-
-def _build_gate_up_layer(out_features, tp_world_size, tp_index):
-    layer = GateUpPack_LinearLayer(torch.nn.Linear(3, out_features, bias=False), mp_group=None, name="dense_h_to_4h")
+def _build_gate_up_layer(out_features, tp_world_size, tp_index, meta):
+    layer = GateUpPack_LinearLayer(torch.nn.Linear(3, out_features, bias=False),
+                                   mp_group=None,
+                                   name="dense_h_to_4h",
+                                   tp_meta=meta)
     layer.tp_world_size = tp_world_size
     layer.tp_index = tp_index
     layer._freeze_partition_sizes(out_features)
     return layer
 
 
-def test_gate_up_partition_ignores_later_grain_size_changes(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=0)
+def test_gate_up_partition_ignores_later_grain_size_changes():
+    meta = AutoTPMeta(tp_grain_size=1)
+    layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=0, meta=meta)
     assert layer._subparam_shard_widths == [[3, 2], [3, 2]]
 
     full_weight = torch.arange(30, dtype=torch.float32).view(10, 3)
@@ -42,8 +38,11 @@ def test_gate_up_partition_ignores_later_grain_size_changes(clean_tp_shard_globa
     first = torch.nn.Parameter(full_weight.clone())
     layer._tp_partition([first, None])
 
-    # A second AutoTP model would install its own grain size before this layer is gathered.
-    tp_shard.set_tp_grain_size(4)
+    # A second AutoTP model would carry a different grain size; the first layer's split must
+    # not move with it, because the layer resolved its split from its own tp_meta.
+    other_meta = AutoTPMeta(tp_grain_size=4)
+    _ = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=0, meta=other_meta)
+
     second = torch.nn.Parameter(full_weight.clone())
     layer._tp_partition([second, None])
 
@@ -51,13 +50,13 @@ def test_gate_up_partition_ignores_later_grain_size_changes(clean_tp_shard_globa
     assert torch.equal(first.data, second.data)
 
 
-def test_gate_up_partition_covers_the_whole_weight(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
+def test_gate_up_partition_covers_the_whole_weight():
+    meta = AutoTPMeta(tp_grain_size=1)
     full_weight = torch.arange(30, dtype=torch.float32).view(10, 3)
 
     shards = []
     for tp_index in range(2):
-        layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=tp_index)
+        layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=tp_index, meta=meta)
         param = torch.nn.Parameter(full_weight.clone())
         layer._tp_partition([param, None])
         shards.append(param.data)
@@ -81,25 +80,23 @@ def __init__(self, split_size):
         self.attn = _QWenAttention(split_size)
 
 
-def _build_qwen_attn_layer(block, hidden, tp_world_size, tp_index):
+def _build_qwen_attn_layer(block, hidden, tp_world_size, tp_index, meta):
     layer = fused_LinearLayer(torch.nn.Linear(hidden, 3 * hidden, bias=False),
                               mp_group=None,
                               skip_partition=True,
                               name="c_attn",
-                              fused_module=block)
+                              fused_module=block,
+                              tp_meta=meta)
     layer.tp_world_size = tp_world_size
     layer.tp_index = tp_index
     layer._freeze_partition_sizes(3 * hidden)
     return layer
 
 
-def test_qwen_split_size_follows_the_frozen_shard_width(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(4)
-    tp_shard.set_n_embd(12)
-
+def test_qwen_split_size_follows_the_frozen_shard_width():
+    meta = AutoTPMeta(num_kv_heads=4, n_embd=12, tp_grain_size=1)
     block = _QWenBlock(split_size=12)
-    layer = _build_qwen_attn_layer(block, hidden=12, tp_world_size=4, tp_index=3)
+    layer = _build_qwen_attn_layer(block, hidden=12, tp_world_size=4, tp_index=3, meta=meta)
 
     weight = torch.nn.Parameter(torch.zeros(36, 12))
     layer._tp_partition([weight, None])
@@ -109,29 +106,26 @@ def test_qwen_split_size_follows_the_frozen_shard_width(clean_tp_shard_globals):
     assert weight.shape[0] == 3 * block.attn.split_size
 
 
-def test_qwen_rejects_a_tensor_parallel_size_that_empties_a_rank(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(4)
-    tp_shard.set_n_embd(12)
+def test_qwen_rejects_a_tensor_parallel_size_that_empties_a_rank():
+    meta = AutoTPMeta(num_kv_heads=4, n_embd=12, tp_grain_size=1)
 
     with pytest.raises(RuntimeError, match="empty query/key/value shard"):
-        _build_qwen_attn_layer(_QWenBlock(split_size=12), hidden=12, tp_world_size=16, tp_index=15)
+        _build_qwen_attn_layer(_QWenBlock(split_size=12), hidden=12, tp_world_size=16, tp_index=15, meta=meta)
 
 
 class _CodeGenBlock(torch.nn.Module):
     pass
 
 
-def test_interleaved_fused_layout_refuses_to_gather(clean_tp_shard_globals):
-    tp_shard.set_tp_grain_size(1)
-    tp_shard.set_num_kv_heads(8)
-    tp_shard.set_n_embd(4)
+def test_interleaved_fused_layout_refuses_to_gather():
+    meta = AutoTPMeta(num_kv_heads=8, n_embd=4, tp_grain_size=1)
 
     layer = fused_LinearLayer(torch.nn.Linear(4, 24, bias=False),
                               mp_group=None,
                               skip_partition=True,
                               name="qkv_proj",
-                              fused_module=_CodeGenBlock())
+                              fused_module=_CodeGenBlock(),
+                              tp_meta=meta)
     layer.tp_world_size = 2
     layer.tp_index = 0
     layer._freeze_partition_sizes(24)
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 c86dc71afa18..0d465af4dbf4 100644
--- a/tests/unit/module_inject/test_tp_partition_config_path.py
+++ b/tests/unit/module_inject/test_tp_partition_config_path.py
@@ -85,6 +85,7 @@ def capture(self, child, full_name):
             linear_layer_setting=None,
             orig_layer_impl=None,
             partition_config=config,
+            model_config=getattr(model, "config", None),
         )
         autotp._replace_module(model)
     finally:
@@ -151,6 +152,7 @@ def _build_gathered_lm_head_autotp(model, mp_size=1):
         linear_layer_setting=None,
         orig_layer_impl=None,
         partition_config=config,
+        model_config=getattr(model, "config", None),
     )
     autotp.set_tensor_parallel_config(mp_size, None)
     autotp.update_linear_policies()
diff --git a/tests/unit/module_inject/test_tp_shard.py b/tests/unit/module_inject/test_tp_shard.py
index dc72ecde8bfb..4066c456cdaa 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -6,26 +6,16 @@
 import pytest
 
 from deepspeed.module_inject import tp_shard
-from deepspeed.module_inject.tp_shard import get_shard_size, 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)
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 
 
 @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)
+    meta = AutoTPMeta(tp_grain_size=64)
 
-    shard_sizes = get_shard_size_list(total_size, tp_size, "lm_head")
+    shard_sizes = get_shard_size_list(total_size, tp_size, meta, "lm_head")
 
     assert sum(shard_sizes) == total_size
     # Only the rank that absorbs the sub-grain tail gives up its alignment.
@@ -33,59 +23,60 @@ def test_grain_quantized_shards_tile_the_dimension(total_size, tp_size):
 
 
 def test_uneven_shards_without_grain_quantization():
-    assert get_shard_size_list(101, 2, "lm_head") == [51, 50]
+    assert get_shard_size_list(101, 2, AutoTPMeta(), "lm_head") == [51, 50]
 
 
 def test_kv_head_shards_tile_the_dimension():
-    set_num_kv_heads(6)
+    meta = AutoTPMeta(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]
+    assert get_shard_size_list(384, 4, meta, "layers.0.self_attn.q_proj") == [128, 128, 64, 64]
+
+
+def test_two_models_do_not_clobber_each_others_meta():
+    # Each model carries its own AutoTPMeta, so loading a second model does not re-shard the
+    # first one.
+    model_a = AutoTPMeta(num_kv_heads=6, tp_grain_size=64)
+    model_b = AutoTPMeta(num_kv_heads=2, tp_grain_size=1)
+
+    a_qproj = get_shard_size_list(384, 4, model_a, "layers.0.self_attn.q_proj")
+    a_lmhead = get_shard_size_list(1001, 2, model_a, "lm_head")
+
+    # A second model is loaded into the same process.
+    _ = get_shard_size_list(384, 4, model_b, "layers.0.self_attn.q_proj")
+
+    # Model A's partition contract is unchanged.
+    assert get_shard_size_list(384, 4, model_a, "layers.0.self_attn.q_proj") == a_qproj
+    assert get_shard_size_list(1001, 2, model_a, "lm_head") == a_lmhead
 
 
 def test_process_group_resolves_noncontiguous_group_rank(monkeypatch):
-    set_tp_grain_size(64)
+    meta = AutoTPMeta(tp_grain_size=64)
     tp_group = object()
     monkeypatch.setattr(tp_shard.dist, "get_rank", lambda group=None: 1 if group is tp_group else 2)
 
-    shard_sizes = get_shard_size_list(50257, 2, "lm_head")
-    assert get_shard_size(50257, 2, "lm_head", mp_group=tp_group) == shard_sizes[1]
+    shard_sizes = get_shard_size_list(50257, 2, meta, "lm_head")
+    assert get_shard_size(50257, 2, meta, "lm_head", mp_group=tp_group) == shard_sizes[1]
 
 
 def test_shard_size_refuses_to_guess_subgroup_rank(monkeypatch):
     monkeypatch.setattr(tp_shard.dist, "get_world_size", lambda: 4)
 
     with pytest.raises(ValueError, match="group-local rank or process group"):
-        get_shard_size(12, 2)
+        get_shard_size(12, 2, AutoTPMeta())
 
 
-def test_explicit_num_kv_heads_is_used():
-    assert get_shard_size_list(
-        384,
-        4,
-        "self_attn.q_proj",
-        num_kv_heads=6,
-    ) == [128, 128, 64, 64]
+def test_explicit_num_kv_heads_overrides_meta():
+    # Fused Q/K/V sub-parameters are split against their own head counts, which differ from
+    # the model-wide kv-head count carried by the meta.
+    meta = AutoTPMeta(num_kv_heads=2)
 
+    assert get_shard_size_list(384, 4, meta, "self_attn.q_proj", num_kv_heads=6) == [128, 128, 64, 64]
 
-def test_explicit_num_kv_heads_matches_global_value():
-    set_num_kv_heads(6)
-    expected = get_shard_size_list(384, 4, "self_attn.q_proj")
 
-    actual = get_shard_size_list(
-        384,
-        4,
-        "self_attn.q_proj",
-        num_kv_heads=6,
-    )
+def test_explicit_num_kv_heads_matches_meta_value():
+    expected = get_shard_size_list(384, 4, AutoTPMeta(num_kv_heads=6), "self_attn.q_proj")
 
-    assert actual == expected
+    actual = get_shard_size_list(384, 4, AutoTPMeta(), "self_attn.q_proj", num_kv_heads=6)
 
-
-def test_uneven_shards_without_grain_quantization_no_kv_heads_used():
-    assert get_shard_size_list(
-        101,
-        2,
-        "lm_head",
-        num_kv_heads=None,
-    ) == [51, 50]
+    assert actual == expected
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 fe996e267f3c..d6478f6f509a 100644
--- a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
+++ b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
@@ -12,6 +12,7 @@
                                             VOCABULARY_PARAMETER_PATTERNS, DS_AUTOTP_UC_META,
                                             UNIVERSAL_CHECKPOINT_VERSION_VALUE)
 from deepspeed.checkpoint.universal_checkpoint import _narrow_sub_params, _resolve_autotp_partition
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 from deepspeed.module_inject.layers import (_build_param_uc_restore_meta, _get_param_uc_conversion_meta,
                                             _subparam_shard_widths, GateUpPack_LinearLayer, LinearAllreduce,
                                             LinearLayer, SubParamLinearAllreduce, SubParamLinearLayer,
@@ -403,7 +404,7 @@ def test_sub_param_shard_widths_round_trip_with_zero_width_ranks(tp_world_size):
     # More ranks than kv heads leaves some ranks holding none of a sub-parameter. Those empty
     # shards still have to tile the sub-parameter and survive a restore round trip.
     sub_param_sizes = (8, 2, 2)
-    widths = _subparam_shard_widths(sub_param_sizes, tp_world_size)
+    widths = _subparam_shard_widths(sub_param_sizes, tp_world_size, AutoTPMeta())
 
     assert any(width == 0 for per_rank in widths for width in per_rank)
     for size, per_rank in zip(sub_param_sizes, widths):
@@ -474,27 +475,22 @@ def test_sub_param_layer_materializes_zero_width_final_dimension(layer_cls):
 
 
 def test_lm_head_forward_uses_frozen_partition_sizes():
-    # The weight columns were cut when the layer was built. A second AutoTP model overwrites the
-    # process-wide tp_shard globals, so re-deriving the split in forward would slice the input
-    # differently than the weight and the row-parallel all-reduce would hide the mismatch.
-    from deepspeed.module_inject import tp_shard
+    # The weight columns were cut when the layer was built; forward must slice the input with
+    # the same frozen partition sizes rather than re-derive them.
     from deepspeed.module_inject.layers import LmHeadLinearAllreduce
+    from deepspeed.module_inject.tp_shard import AutoTPMeta
 
-    grain_size, kv_heads = tp_shard.tp_grain_size, tp_shard.num_kv_heads
-    try:
-        tp_shard.set_tp_grain_size(1)
-        tp_shard.set_num_kv_heads(None)
-        layer = LmHeadLinearAllreduce(torch.nn.Linear(101, 8, bias=False), mp_group=None)
-        layer.tp_world_size = 2
-        layer.tp_index = 1
-        frozen = layer._freeze_partition_sizes(101)
-        assert frozen == (51, 50)
-
-        # A later model narrows the grain, which would change a recomputed split.
-        tp_shard.set_tp_grain_size(64)
-        layer.weight.data = torch.zeros(8, frozen[1])
-
-        layer(torch.zeros(1, 1, 101))
-    finally:
-        tp_shard.set_tp_grain_size(grain_size)
-        tp_shard.set_num_kv_heads(kv_heads)
+    layer = LmHeadLinearAllreduce(torch.nn.Linear(101, 8, bias=False),
+                                  mp_group=None,
+                                  tp_meta=AutoTPMeta(tp_grain_size=1))
+    layer.tp_world_size = 2
+    layer.tp_index = 1
+    frozen = layer._freeze_partition_sizes(101)
+    assert frozen == (51, 50)
+
+    # A later model carries a different grain; this layer's split is frozen from its own meta.
+    other_meta = AutoTPMeta(tp_grain_size=64)
+    assert other_meta.tp_grain_size != layer.tp_meta.tp_grain_size
+    layer.weight.data = torch.zeros(8, frozen[1])
+
+    layer(torch.zeros(1, 1, 101))
diff --git a/tests/unit/sequence_parallelism/test_ulysses.py b/tests/unit/sequence_parallelism/test_ulysses.py
index abb9899892a8..72e8a81c8048 100644
--- a/tests/unit/sequence_parallelism/test_ulysses.py
+++ b/tests/unit/sequence_parallelism/test_ulysses.py
@@ -16,7 +16,7 @@
 from unit.util import skip_on_arch
 from unit.simple_model import *
 from deepspeed.utils import groups
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 #Use mesh device to create data and sequence parallel group
 
 
@@ -205,7 +205,7 @@ def seq_batch_heads_hash(d0, d1, h, offset_d0=0, offset_d1=0, offset_h=0):
             d0_indices = torch.arange(s2h_tensor.shape[0]).reshape(-1, 1, 1, 1)
             d1_indices = torch.arange(s2h_tensor.shape[1]).reshape(1, -1, 1, 1)
             h_indices = torch.arange(s2h_tensor.shape[2]).reshape(1, 1, -1, 1)
-            shard_list = get_shard_size_list(num_heads, groups._get_sequence_parallel_world_size())
+            shard_list = get_shard_size_list(num_heads, groups._get_sequence_parallel_world_size(), AutoTPMeta())
             head_offset = sum(shard_list[:groups._get_sequence_parallel_rank()])
             s2h_truth = torch.zeros_like(s2h_tensor)
             s2h_truth[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices, 0, 0, head_offset)
diff --git a/tests/unit/v1/moe/test_autoep_autotp_runtime.py b/tests/unit/v1/moe/test_autoep_autotp_runtime.py
index 6d0e79a95c8a..b09f1b9d5547 100644
--- a/tests/unit/v1/moe/test_autoep_autotp_runtime.py
+++ b/tests/unit/v1/moe/test_autoep_autotp_runtime.py
@@ -112,6 +112,7 @@ def __init__(self):
         linear_layer_setting=None,
         orig_layer_impl=None,
         partition_config=None,
+        model_config=getattr(model, "config", None),
     )
 
     calls = []

From 620a9a7b5e35aaff7a7d62758a75eddb7c263851 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Mon, 10 Aug 2026 11:03:38 +0800
Subject: [PATCH 02/20] Unify head-count extraction behind shared attribute
 lists

AutoTP (``AutoTPMeta.from_model_config``) and the inference engine
(``_get_model_head_count`` / ``_get_model_kv_head_count``) each kept their own
attribute-name lists for kv-head and attention-head counts, so the two paths
recognized different model families (e.g. chatglm only on the AutoTP side, legacy
``n_head_kv`` / ``kv_n_heads`` only on the inference side) and could even disagree on a
plain transformer. Consolidate each count behind one shared list and helper in tp_shard
so coverage and probe order live in a single place:

- ``_KV_HEAD_ATTRS`` / ``_kv_head_count_from`` for the key/value head count
- ``_ATTN_HEAD_ATTRS`` / ``_attention_head_count_from`` for the attention head count

Both ``AutoTPMeta.from_model_config`` and the inference engine consume them, so neither
count is re-extracted on either side.

The union keeps the legacy aliases for older configs/checkpoints, annotated with the
transformers version that superseded each (``n_head_kv`` after 4.33, ``kv_n_heads``
superseded at top-level by ``num_key_value_heads`` in 4.40).

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/inference/engine.py       | 15 +++++----
 deepspeed/module_inject/tp_shard.py | 48 +++++++++++++++++++++++------
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
index 6e5188ba5928..6e343a1e20e5 100755
--- a/deepspeed/inference/engine.py
+++ b/deepspeed/inference/engine.py
@@ -25,6 +25,7 @@
 from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
 from deepspeed.accelerator import get_accelerator
 from ..module_inject.policy import TransformerPolicy
+from deepspeed.module_inject.tp_shard import _attention_head_count_from, _kv_head_count_from
 from ..module_inject.auto_tp import AutoTP
 
 from ..module_inject.replace_policy import generic_policies
@@ -249,20 +250,18 @@ def _get_model_head_count(self, module):
         for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
             if source is None:
                 continue
-            for name in ("num_heads", "n_heads", "n_head", "num_attention_heads"):
-                value = getattr(source, name, None)
-                if value is not None:
-                    return value
+            value = _attention_head_count_from(source)
+            if value is not None:
+                return value
         raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
 
     def _get_model_kv_head_count(self, module, num_heads):
         for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
             if source is None:
                 continue
-            for name in ("num_key_value_heads", "num_kv_heads", "n_head_kv", "kv_n_heads"):
-                value = getattr(source, name, None)
-                if value is not None:
-                    return value
+            value = _kv_head_count_from(source)
+            if value is not None:
+                return value
         return num_heads
 
     def _pre_forward_hook(self, module, *inputs, **kwargs):
diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index 6d092f21c5dd..24f4790f5a00 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -8,6 +8,43 @@
 
 from deepspeed import comm as dist
 
+# Attribute names that carry the key/value head count, in probe order. Modern canonical names
+# first, then legacy aliases kept for older configs / checkpoints, then the query-head names as
+# an implicit "kv == q heads" default for plain (non-GQA) transformers. Shared by AutoTP and the
+# inference engine so the two never diverge on which model families they recognize.
+_KV_HEAD_ATTRS = (
+    'num_key_value_heads',  # llama-class, mistral, qwen2, gemma, phi3, dbrx top-level, ...
+    'num_kv_heads',  # falcon, qwen3_moe
+    'multi_query_group_num',  # chatglm2 / chatglm3 (not in stock transformers)
+    'n_head_kv',  # legacy Falcon custom-code; deprecated after transformers 4.33 (-> num_kv_heads)
+    'kv_n_heads',  # dbrx nested attn_config; top-level config exposes num_key_value_heads (>= transformers 4.40)
+    'num_attention_heads',  # plain transformer: kv == q heads
+    'n_heads',
+    'attention_heads',
+)
+
+
+def _kv_head_count_from(config) -> Optional[int]:
+    """First non-None kv-head count found on ``config`` via :data:`_KV_HEAD_ATTRS`, else ``None``."""
+    for name in _KV_HEAD_ATTRS:
+        value = getattr(config, name, None)
+        if value is not None:
+            return value
+    return None
+
+
+# Attribute names that carry the (query) attention head count, in probe order.
+_ATTN_HEAD_ATTRS = ('num_attention_heads', 'num_heads', 'n_heads', 'n_head', 'attention_heads')
+
+
+def _attention_head_count_from(config) -> Optional[int]:
+    """First non-None attention head count found on ``config`` via :data:`_ATTN_HEAD_ATTRS`, else ``None``."""
+    for name in _ATTN_HEAD_ATTRS:
+        value = getattr(config, name, None)
+        if value is not None:
+            return value
+    return None
+
 
 @dataclass(frozen=True)
 class AutoTPMeta:
@@ -31,21 +68,14 @@ def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta"
         """
         if model_config is None:
             return cls(tp_grain_size=tp_grain_size)
-        # multi_query_group_num is for chatglm2 & chatglm3
-        num_kv_heads = None
-        for name in ('multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads',
-                     'attention_heads'):
-            if hasattr(model_config, name):
-                num_kv_heads = getattr(model_config, name)
-                if num_kv_heads is not None:
-                    break
+        num_kv_heads = _kv_head_count_from(model_config)
         n_embd = None
         for name in ('n_embd', 'hidden_size'):
             if hasattr(model_config, name):
                 n_embd = getattr(model_config, name)
             if n_embd is not None:
                 break
-        num_attention_heads = getattr(model_config, 'num_attention_heads', None)
+        num_attention_heads = _attention_head_count_from(model_config)
         return cls(num_kv_heads=num_kv_heads,
                    num_attention_heads=num_attention_heads,
                    n_embd=n_embd,

From 7b4f51450f8b362d148ff005580f51ddc79c5011 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Mon, 10 Aug 2026 11:49:06 +0800
Subject: [PATCH 03/20] Remove dead meta param from alibi head-sharding helpers

``get_head_shard_sizes`` and ``install_head_sharded_helper`` grew a ``meta``
parameter that no caller ever passed -- it was always ``None``, the
``if meta is not None`` kv-head discovery branch was unreachable, and the
``meta or AutoTPMeta()`` fallback always evaluated to ``AutoTPMeta()``. Drop the
parameter and the dead discovery block; the helpers honestly take
``num_heads`` / ``num_kv_heads``, which every caller already supplies.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/module_inject/auto_tp_model_utils.py | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/deepspeed/module_inject/auto_tp_model_utils.py b/deepspeed/module_inject/auto_tp_model_utils.py
index fbe0c9ac9635..400677248839 100644
--- a/deepspeed/module_inject/auto_tp_model_utils.py
+++ b/deepspeed/module_inject/auto_tp_model_utils.py
@@ -27,23 +27,17 @@ def __setattr__(self, name, value):
             setattr(self._module, name, value)
 
 
-def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None, meta: Optional[AutoTPMeta] = None):
+def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None):
     tp_world_size = dist.get_world_size(group=mp_group)
     return get_shard_size_list(
         num_heads,
         tp_world_size,
-        meta or AutoTPMeta(),
+        AutoTPMeta(),
         num_kv_heads=num_kv_heads,
     )
 
 
-def install_head_sharded_helper(module,
-                                name,
-                                wrapper,
-                                mp_group=None,
-                                num_heads=None,
-                                num_kv_heads=None,
-                                meta: Optional[AutoTPMeta] = None):
+def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=None, num_kv_heads=None):
     """Give ``module`` a head-slicing wrapper around one of its own methods.
 
     The wrapper is bound to this instance instead of installed on its class. A class-wide patch
@@ -55,7 +49,7 @@ def install_head_sharded_helper(module,
     if original_name not in module.__dict__:
         # Wrapping an already wrapped instance would make it delegate to itself.
         setattr(module, original_name, getattr(module, name))
-    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads, meta) if num_heads is not None else None
+    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads) if num_heads is not None else None
     setattr(
         module, name,
         functools.partial(wrapper, module, mp_group=mp_group, head_shard_sizes=shard_sizes, total_num_heads=num_heads))

From 189f6fe0dac1ba1eeadd365a320e84c108c20abd Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Mon, 10 Aug 2026 12:14:01 +0800
Subject: [PATCH 04/20] Thread AutoTPMeta through the alibi head-sharding
 helpers

The alibi helpers (``get_head_shard_sizes``, ``install_head_sharded_helper``) took
``num_heads`` / ``num_kv_heads`` as scalars, so the inference engine extracted them via
``_get_model_head_count`` / ``_get_model_kv_head_count`` -- a second copy of the
head-count probe that ``AutoTPMeta.from_model_config`` already does. With AutoTPMeta
carrying both counts, the helpers now take a single ``meta`` and the inference engine
builds one per model (via the shared ``_attention_head_count_from`` /
``_kv_head_count_from`` probes), deleting the two ``_get_model_*`` methods.

The runtime alibi wrappers and ``_head_shard`` are unchanged: they still consume the
``head_shard_sizes`` + ``total_num_heads`` bound at install time, now derived from meta.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/inference/engine.py                 | 68 +++++++++----------
 .../module_inject/auto_tp_model_utils.py      | 17 +++--
 .../test_autotp_custom_patterns.py            | 20 +++---
 3 files changed, 52 insertions(+), 53 deletions(-)

diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py
index 6e343a1e20e5..f247be89a02c 100755
--- a/deepspeed/inference/engine.py
+++ b/deepspeed/inference/engine.py
@@ -25,7 +25,7 @@
 from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
 from deepspeed.accelerator import get_accelerator
 from ..module_inject.policy import TransformerPolicy
-from deepspeed.module_inject.tp_shard import _attention_head_count_from, _kv_head_count_from
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 from ..module_inject.auto_tp import AutoTP
 
 from ..module_inject.replace_policy import generic_policies
@@ -219,50 +219,44 @@ def build_alibi_tensor(self):
             if hasattr(self.module.transformer, 'build_alibi_tensor'):
                 # The heads must be sliced with the same tensor-parallel group that partitioned
                 # the attention weights, so bind it rather than letting the helper guess.
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
-                shard_sizes = get_head_shard_sizes(num_heads, self.mp_group, num_kv_heads)
-                self.module.transformer.build_alibi_tensor = functools.partial(build_bloom_alibi_tensor,
-                                                                               mp_group=self.mp_group,
-                                                                               head_shard_sizes=shard_sizes,
-                                                                               total_num_heads=num_heads)
+                meta = self._autotp_meta(self.module.transformer)
+                shard_sizes = get_head_shard_sizes(meta, self.mp_group)
+                self.module.transformer.build_alibi_tensor = functools.partial(
+                    build_bloom_alibi_tensor,
+                    mp_group=self.mp_group,
+                    head_shard_sizes=shard_sizes,
+                    total_num_heads=meta.num_attention_heads)
             if hasattr(self.module.transformer, 'build_mpt_alibi_tensor'):
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
+                meta = self._autotp_meta(self.module.transformer)
                 install_head_sharded_helper(self.module.transformer, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor,
-                                            self.mp_group, num_heads, num_kv_heads)
+                                            meta, self.mp_group)
         if hasattr(self.module, 'model'):
             if hasattr(self.module.model, 'get_alibi_mask'):
-                num_heads = self._get_model_head_count(self.module.model)
-                num_kv_heads = self._get_model_kv_head_count(self.module.model, num_heads)
-                install_head_sharded_helper(self.module.model, 'get_alibi_mask', get_alibi_mask, self.mp_group,
-                                            num_heads, num_kv_heads)
+                meta = self._autotp_meta(self.module.model)
+                install_head_sharded_helper(self.module.model, 'get_alibi_mask', get_alibi_mask, meta, self.mp_group)
 
     def build_attn_bias(self):
         if hasattr(self.module, 'transformer'):
             if hasattr(self.module.transformer, '_attn_bias'):
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
-                install_head_sharded_helper(self.module.transformer, '_attn_bias', build_mpt_atten_bias_tensor,
-                                            self.mp_group, num_heads, num_kv_heads)
-
-    def _get_model_head_count(self, module):
-        for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
-            if source is None:
-                continue
-            value = _attention_head_count_from(source)
-            if value is not None:
-                return value
-        raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
-
-    def _get_model_kv_head_count(self, module, num_heads):
-        for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
-            if source is None:
-                continue
-            value = _kv_head_count_from(source)
-            if value is not None:
-                return value
-        return num_heads
+                meta = self._autotp_meta(self.module.transformer)
+                install_head_sharded_helper(self.module.transformer, '_attn_bias', build_mpt_atten_bias_tensor, meta,
+                                            self.mp_group)
+
+    def _autotp_meta(self, module):
+        # from_model_config extracts from a single config object, but the head counts may live on
+        # the module, its config, or the top-level model config. Probe each source via the shared
+        # from_model_config and merge per field -- num_attention_heads is mandatory (alibi needs
+        # it), num_kv_heads is None for non-GQA. Subsumes the former
+        # _get_model_head_count / _get_model_kv_head_count pair.
+        metas = [
+            AutoTPMeta.from_model_config(s)
+            for s in (module, getattr(module, "config", None), getattr(self.module, "config", None)) if s is not None
+        ]
+        num_heads = next((m.num_attention_heads for m in metas if m.num_attention_heads is not None), None)
+        if num_heads is None:
+            raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
+        num_kv_heads = next((m.num_kv_heads for m in metas if m.num_kv_heads is not None), None)
+        return AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_kv_heads)
 
     def _pre_forward_hook(self, module, *inputs, **kwargs):
         if self.use_cuda_events:
diff --git a/deepspeed/module_inject/auto_tp_model_utils.py b/deepspeed/module_inject/auto_tp_model_utils.py
index 400677248839..9457fe243e7f 100644
--- a/deepspeed/module_inject/auto_tp_model_utils.py
+++ b/deepspeed/module_inject/auto_tp_model_utils.py
@@ -27,17 +27,19 @@ def __setattr__(self, name, value):
             setattr(self._module, name, value)
 
 
-def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None):
+def get_head_shard_sizes(meta: AutoTPMeta, mp_group=None):
+    num_heads = meta.num_attention_heads
+    num_kv_heads = meta.num_kv_heads
     tp_world_size = dist.get_world_size(group=mp_group)
     return get_shard_size_list(
         num_heads,
         tp_world_size,
-        AutoTPMeta(),
+        meta,
         num_kv_heads=num_kv_heads,
     )
 
 
-def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=None, num_kv_heads=None):
+def install_head_sharded_helper(module, name, wrapper, meta, mp_group=None):
     """Give ``module`` a head-slicing wrapper around one of its own methods.
 
     The wrapper is bound to this instance instead of installed on its class. A class-wide patch
@@ -49,10 +51,15 @@ def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=
     if original_name not in module.__dict__:
         # Wrapping an already wrapped instance would make it delegate to itself.
         setattr(module, original_name, getattr(module, name))
-    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads) if num_heads is not None else None
+    total_num_heads = meta.num_attention_heads
+    shard_sizes = get_head_shard_sizes(meta, mp_group) if total_num_heads is not None else None
     setattr(
         module, name,
-        functools.partial(wrapper, module, mp_group=mp_group, head_shard_sizes=shard_sizes, total_num_heads=num_heads))
+        functools.partial(wrapper,
+                          module,
+                          mp_group=mp_group,
+                          head_shard_sizes=shard_sizes,
+                          total_num_heads=total_num_heads))
 
 
 def _head_shard(num_heads, mp_group=None, head_shard_sizes=None, total_num_heads=None):
diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
index a381f57f9720..7573a779ecd3 100644
--- a/tests/unit/model_parallelism/test_autotp_custom_patterns.py
+++ b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
@@ -563,8 +563,7 @@ def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, devic
         install_head_sharded_helper(transformer,
                                     'build_mpt_alibi_tensor',
                                     build_mpt_alibi_tensor,
-                                    num_heads=num_heads,
-                                    num_kv_heads=num_heads)
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
 
         alibi = transformer.build_mpt_alibi_tensor(num_heads, 3)
 
@@ -594,8 +593,7 @@ class MptSubclass(MptTransformer):
         install_head_sharded_helper(first,
                                     'build_mpt_alibi_tensor',
                                     build_mpt_alibi_tensor,
-                                    num_heads=num_heads,
-                                    num_kv_heads=num_heads)
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
         expected = first.build_mpt_alibi_tensor(num_heads, 3)
 
         # Injecting a second model of the same architecture must not make either of them
@@ -604,8 +602,7 @@ class MptSubclass(MptTransformer):
         install_head_sharded_helper(second,
                                     'build_mpt_alibi_tensor',
                                     build_mpt_alibi_tensor,
-                                    num_heads=num_heads,
-                                    num_kv_heads=num_heads)
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
         torch.testing.assert_close(second.build_mpt_alibi_tensor(num_heads, 3), expected)
         torch.testing.assert_close(first.build_mpt_alibi_tensor(num_heads, 3), expected)
 
@@ -618,8 +615,7 @@ class MptSubclass(MptTransformer):
         install_head_sharded_helper(derived,
                                     'build_mpt_alibi_tensor',
                                     build_mpt_alibi_tensor,
-                                    num_heads=num_heads,
-                                    num_kv_heads=num_heads)
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_heads))
         torch.testing.assert_close(derived.build_mpt_alibi_tensor(num_heads, 3), expected)
 
     def test_head_sharded_helper_freezes_the_models_split(self):
@@ -636,8 +632,7 @@ def build_mpt_alibi_tensor(self, heads, sequence_length, alibi_bias_max=8, devic
         install_head_sharded_helper(transformer,
                                     'build_mpt_alibi_tensor',
                                     build_mpt_alibi_tensor,
-                                    num_heads=num_heads,
-                                    num_kv_heads=3)
+                                    meta=AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=3))
 
         # The helper freezes the [4, 2] split from this model's own num_kv_heads at install
         # time.
@@ -684,7 +679,10 @@ def get_alibi_mask(self, tensor, sequence_length):
                                     dtype=torch.float32).view(-1, 1, 1).expand(-1, sequence_length, sequence_length)
 
         model = AlibiModel()
-        install_head_sharded_helper(model, 'get_alibi_mask', get_alibi_mask, num_heads=5, num_kv_heads=5)
+        install_head_sharded_helper(model,
+                                    'get_alibi_mask',
+                                    get_alibi_mask,
+                                    meta=AutoTPMeta(num_attention_heads=5, num_kv_heads=5))
 
         shard_sizes = [3, 2]
         model.n_head = shard_sizes[dist.get_rank()]

From 1c5cfe44aec5156814e771bd69cbf52c6ee96c5a Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Mon, 10 Aug 2026 14:57:57 +0800
Subject: [PATCH 05/20] Consolidate the gate-up partition test

test_gate_up_partition_ignores_later_grain_size_changes existed to check that a
layer's frozen shard widths survived a second AutoTP model overwriting the
process-wide grain global. With per-model AutoTPMeta the "second model" leg became
vacuous -- a layer holding meta A is unaffected by merely constructing a layer
with meta B -- so the test no longer tested what its name says. Fold its one
piece of real value (the explicit _subparam_shard_widths == [[3,2],[3,2]]
assertion) into test_gate_up_partition_covers_the_whole_weight, which already
exercises the same layer and partition.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 .../module_inject/test_fused_repartition.py   | 25 +++----------------
 1 file changed, 3 insertions(+), 22 deletions(-)

diff --git a/tests/unit/module_inject/test_fused_repartition.py b/tests/unit/module_inject/test_fused_repartition.py
index 972b28703076..f7b7faf544a5 100644
--- a/tests/unit/module_inject/test_fused_repartition.py
+++ b/tests/unit/module_inject/test_fused_repartition.py
@@ -28,28 +28,6 @@ def _build_gate_up_layer(out_features, tp_world_size, tp_index, meta):
     return layer
 
 
-def test_gate_up_partition_ignores_later_grain_size_changes():
-    meta = AutoTPMeta(tp_grain_size=1)
-    layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=0, meta=meta)
-    assert layer._subparam_shard_widths == [[3, 2], [3, 2]]
-
-    full_weight = torch.arange(30, dtype=torch.float32).view(10, 3)
-
-    first = torch.nn.Parameter(full_weight.clone())
-    layer._tp_partition([first, None])
-
-    # A second AutoTP model would carry a different grain size; the first layer's split must
-    # not move with it, because the layer resolved its split from its own tp_meta.
-    other_meta = AutoTPMeta(tp_grain_size=4)
-    _ = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=0, meta=other_meta)
-
-    second = torch.nn.Parameter(full_weight.clone())
-    layer._tp_partition([second, None])
-
-    assert tuple(second.shape) == (6, 3)
-    assert torch.equal(first.data, second.data)
-
-
 def test_gate_up_partition_covers_the_whole_weight():
     meta = AutoTPMeta(tp_grain_size=1)
     full_weight = torch.arange(30, dtype=torch.float32).view(10, 3)
@@ -57,6 +35,9 @@ def test_gate_up_partition_covers_the_whole_weight():
     shards = []
     for tp_index in range(2):
         layer = _build_gate_up_layer(out_features=10, tp_world_size=2, tp_index=tp_index, meta=meta)
+        # The gate and up halves (5 each) split over 2 ranks as [3, 2]; the layer freezes these
+        # widths from its own tp_meta so the partition is deterministic.
+        assert layer._subparam_shard_widths == [[3, 2], [3, 2]]
         param = torch.nn.Parameter(full_weight.clone())
         layer._tp_partition([param, None])
         shards.append(param.data)

From 24659a0e4248a36d5eab5547dea4b0f877af8c93 Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 11:07:22 +0800
Subject: [PATCH 06/20] Fail loudly when the bigcode fused-QKV split has no
 hidden size

_bigcode_type_transpose slices the fused projection at meta.n_embd to separate
the query block from the replicated kv block. n_embd is Optional, and Python
reads input[:None] as "to the end", so a missing hidden size would hand the
whole weight to q and leave kv empty -- a silently wrong shard rather than an
error. Assert the value is present before slicing.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 deepspeed/module_inject/fusedqkv_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/deepspeed/module_inject/fusedqkv_utils.py b/deepspeed/module_inject/fusedqkv_utils.py
index a43af5224e17..f732aa8074db 100644
--- a/deepspeed/module_inject/fusedqkv_utils.py
+++ b/deepspeed/module_inject/fusedqkv_utils.py
@@ -158,6 +158,9 @@ def _bloom_type_transpose(input, mp_size, meta=meta):
 
     def _bigcode_type_transpose(input, mp_size, meta=meta):
         n_embd = meta.n_embd
+        # A missing hidden size would slice as input[:None], handing the whole fused weight to q
+        # and leaving kv empty rather than failing.
+        assert n_embd is not None
         q = input[:n_embd]
         kv = input[n_embd:]
         shape = q.shape

From 3c0dc8e115bf0857f7cf87a0ef965c3665eca67c Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 11:31:19 +0800
Subject: [PATCH 07/20] Require tp_meta when building a fused QKV layer

fused_LinearLayer defaulted a missing tp_meta to an empty AutoTPMeta before
resolving its sub-parameter layout. That default is not a neutral choice here:
with num_kv_heads left as None, fused_qkv_subparam_sizes skips the chatglm
kv-head branch and falls back to splitting the fused weight into three equal
blocks, so the layer would partition to a structurally wrong layout instead of
failing. Index kwargs directly so an omitted tp_meta raises.

Every AutoTP construction site already passes it; only one test relied on the
default and now states the empty meta explicitly, where the codegen layout it
exercises does not consult it.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 deepspeed/module_inject/layers.py                              | 2 +-
 .../tensor_parallel/test_autotp_universal_checkpoint.py        | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 61a5dd98420c..3a779c968fcf 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -1088,7 +1088,7 @@ def __init__(self, module, mp_group, skip_partition=False, **kwargs):
         # prepare_tp_fused_qkvw takes its own shard sizes without a layer name, so the widths
         # describing its split must be resolved the same way.
         self._subparam_layout_spec = (fused_qkv_subparam_sizes(kwargs.get('fused_module'), tuple(module.weight.shape),
-                                                               kwargs.get('tp_meta') or AutoTPMeta()), None)
+                                                               kwargs['tp_meta']), None)
         super().__init__(module, mp_group, skip_partition, **kwargs)
 
     def _freeze_partition_sizes(self, total_size):
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 d6478f6f509a..17da79d86f11 100644
--- a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
+++ b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
@@ -341,7 +341,8 @@ class CodeGenBlock(torch.nn.Module):
                               mp_group=None,
                               skip_partition=True,
                               fused_module=fused_module,
-                              name="qkv_proj")
+                              name="qkv_proj",
+                              tp_meta=AutoTPMeta())
     model = torch.nn.Module()
     model.qkv_proj = layer
 

From a8915b3bbcf0a21ab2df77f7b14ffee8240569f2 Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 14:36:31 +0800
Subject: [PATCH 08/20] Name the kv-head split parameter for what it is

get_shard_size / get_shard_size_list took a num_kv_heads argument that shadowed
meta.num_kv_heads, so the same name meant "this model's kv-head count" in one
place and "use this instead" in another. Call it eff_num_kv_heads: the head
count the split is actually aligned to, defaulting to the meta value.

get_head_shard_sizes was reading meta.num_kv_heads only to pass it straight back
as that argument, which is what get_shard_size_list already does when it is
omitted. Drop the round trip.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 .../module_inject/auto_tp_model_utils.py      |  9 +-----
 deepspeed/module_inject/tp_shard.py           | 29 ++++++++++---------
 tests/unit/module_inject/test_tp_shard.py     |  4 +--
 3 files changed, 18 insertions(+), 24 deletions(-)

diff --git a/deepspeed/module_inject/auto_tp_model_utils.py b/deepspeed/module_inject/auto_tp_model_utils.py
index 9457fe243e7f..422bc651fbc3 100644
--- a/deepspeed/module_inject/auto_tp_model_utils.py
+++ b/deepspeed/module_inject/auto_tp_model_utils.py
@@ -28,15 +28,8 @@ def __setattr__(self, name, value):
 
 
 def get_head_shard_sizes(meta: AutoTPMeta, mp_group=None):
-    num_heads = meta.num_attention_heads
-    num_kv_heads = meta.num_kv_heads
     tp_world_size = dist.get_world_size(group=mp_group)
-    return get_shard_size_list(
-        num_heads,
-        tp_world_size,
-        meta,
-        num_kv_heads=num_kv_heads,
-    )
+    return get_shard_size_list(meta.num_attention_heads, tp_world_size, meta)
 
 
 def install_head_sharded_helper(module, name, wrapper, meta, mp_group=None):
diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index 24f4790f5a00..863669b849b2 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -82,7 +82,7 @@ def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta"
                    tp_grain_size=tp_grain_size)
 
 
-def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None, num_kv_heads=None):
+def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None, eff_num_kv_heads=None):
     """Size of one shard of ``total_size`` split across a tensor-parallel group of ``mp_size``.
 
     ``meta`` carries this model's ``num_kv_heads`` / ``tp_grain_size`` so the split is stable
@@ -93,12 +93,13 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
     ``dist.get_rank(group=mp_group)`` and the index used by ``get_shard_size_list``. It is not a
     global rank.
 
-    ``num_kv_heads`` overrides ``meta.num_kv_heads`` when set. Passing it explicitly lets callers
-    split fused sub-parameters (Q/K/V) against their respective head counts without
-    reimplementing the KV-head-aligned partition logic at the call site.
+    ``eff_num_kv_heads`` is the head count this split is actually aligned to; it defaults to
+    ``meta.num_kv_heads``. Passing it explicitly lets callers split fused sub-parameters (Q/K/V)
+    against their respective head counts without reimplementing the KV-head-aligned partition
+    logic at the call site.
     """
-    if num_kv_heads is None:
-        num_kv_heads = meta.num_kv_heads
+    if eff_num_kv_heads is None:
+        eff_num_kv_heads = meta.num_kv_heads
     tp_grain_size = meta.tp_grain_size
     last_linear = ["lm_head", "embed_out"]
     # MoE MLP layer use near even division will get better perf.
@@ -116,10 +117,10 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
                 raise ValueError("get_shard_size requires a group-local rank or process group when mp_size "
                                  f"({mp_size}) differs from the distributed world size ({world_size}).")
             rank = dist.get_rank()
-    if num_kv_heads is not None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and \
+    if eff_num_kv_heads is not None and total_size % eff_num_kv_heads == 0 and "mlp" not in str(name) and \
         str(name) not in last_linear and not_moe_mlp_layer:
-        my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0)
-        return total_size * my_slices // num_kv_heads
+        my_slices = (eff_num_kv_heads // mp_size) + (1 if rank < (eff_num_kv_heads % mp_size) else 0)
+        return total_size * my_slices // eff_num_kv_heads
     else:
         if total_size >= tp_grain_size:
             grain_size, remainder = divmod(total_size, tp_grain_size)
@@ -134,16 +135,16 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
             return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
 
-def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None, num_kv_heads=None):
+def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None, eff_num_kv_heads=None):
     shard_sizes = []
-    if num_kv_heads is None:
-        num_kv_heads = meta.num_kv_heads
+    if eff_num_kv_heads is None:
+        eff_num_kv_heads = meta.num_kv_heads
     for i in range(mp_size):
-        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i, num_kv_heads=num_kv_heads))
+        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i, eff_num_kv_heads=eff_num_kv_heads))
     # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
     # 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 "
         f"{total_size} with tp_size={mp_size}, tp_grain_size={meta.tp_grain_size} and "
-        f"num_kv_heads={num_kv_heads}.")
+        f"num_kv_heads={eff_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
index 4066c456cdaa..72809d1c8c4a 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -71,12 +71,12 @@ def test_explicit_num_kv_heads_overrides_meta():
     # the model-wide kv-head count carried by the meta.
     meta = AutoTPMeta(num_kv_heads=2)
 
-    assert get_shard_size_list(384, 4, meta, "self_attn.q_proj", num_kv_heads=6) == [128, 128, 64, 64]
+    assert get_shard_size_list(384, 4, meta, "self_attn.q_proj", eff_num_kv_heads=6) == [128, 128, 64, 64]
 
 
 def test_explicit_num_kv_heads_matches_meta_value():
     expected = get_shard_size_list(384, 4, AutoTPMeta(num_kv_heads=6), "self_attn.q_proj")
 
-    actual = get_shard_size_list(384, 4, AutoTPMeta(), "self_attn.q_proj", num_kv_heads=6)
+    actual = get_shard_size_list(384, 4, AutoTPMeta(), "self_attn.q_proj", eff_num_kv_heads=6)
 
     assert actual == expected

From 7f55639bc76fa7dacedd042f949bce9ccaed8981 Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 14:47:24 +0800
Subject: [PATCH 09/20] Drop the unused kv-head override from the shard-size
 helpers

The override existed because get_shard_size read a process-wide num_kv_heads
that was only set during AutoTP replacement. The inference engine patches the
alibi helpers before that happens, so it probed the head count itself and passed
it in to bypass the uninitialized global. That was the parameter's only caller.

With the count carried by a per-model AutoTPMeta, the alibi path receives its
own model's value like everyone else, and the previous commit removed the
round trip that read meta.num_kv_heads only to pass it straight back. Nothing
outside tp_shard supplies the argument now, so remove it and the two tests
written against it; get_shard_size has a single source for the head count again.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 deepspeed/module_inject/tp_shard.py       | 26 ++++++++---------------
 tests/unit/module_inject/test_tp_shard.py | 16 --------------
 2 files changed, 9 insertions(+), 33 deletions(-)

diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index 863669b849b2..dcffe1b6e5dc 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -82,7 +82,7 @@ def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta"
                    tp_grain_size=tp_grain_size)
 
 
-def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None, eff_num_kv_heads=None):
+def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None):
     """Size of one shard of ``total_size`` split across a tensor-parallel group of ``mp_size``.
 
     ``meta`` carries this model's ``num_kv_heads`` / ``tp_grain_size`` so the split is stable
@@ -92,14 +92,8 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
     ``rank`` is the rank *within the tensor-parallel group*, i.e. in ``[0, mp_size)``, matching
     ``dist.get_rank(group=mp_group)`` and the index used by ``get_shard_size_list``. It is not a
     global rank.
-
-    ``eff_num_kv_heads`` is the head count this split is actually aligned to; it defaults to
-    ``meta.num_kv_heads``. Passing it explicitly lets callers split fused sub-parameters (Q/K/V)
-    against their respective head counts without reimplementing the KV-head-aligned partition
-    logic at the call site.
     """
-    if eff_num_kv_heads is None:
-        eff_num_kv_heads = meta.num_kv_heads
+    num_kv_heads = meta.num_kv_heads
     tp_grain_size = meta.tp_grain_size
     last_linear = ["lm_head", "embed_out"]
     # MoE MLP layer use near even division will get better perf.
@@ -107,7 +101,6 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
     not_moe_mlp_layer = True
     if name != None and any(s in str(name) for s in moe_mlp_layer):
         not_moe_mlp_layer = False
-    # When num_kv_heads is defined, uneven division is possible, otherwise enforce near even division
     if rank is None:
         if mp_group is not None:
             rank = dist.get_rank(group=mp_group)
@@ -117,10 +110,11 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
                 raise ValueError("get_shard_size requires a group-local rank or process group when mp_size "
                                  f"({mp_size}) differs from the distributed world size ({world_size}).")
             rank = dist.get_rank()
-    if eff_num_kv_heads is not None and total_size % eff_num_kv_heads == 0 and "mlp" not in str(name) and \
+    # A known kv-head count allows uneven division, otherwise enforce near even division.
+    if num_kv_heads is not None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and \
         str(name) not in last_linear and not_moe_mlp_layer:
-        my_slices = (eff_num_kv_heads // mp_size) + (1 if rank < (eff_num_kv_heads % mp_size) else 0)
-        return total_size * my_slices // eff_num_kv_heads
+        my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0)
+        return total_size * my_slices // num_kv_heads
     else:
         if total_size >= tp_grain_size:
             grain_size, remainder = divmod(total_size, tp_grain_size)
@@ -135,16 +129,14 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
             return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
 
-def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None, eff_num_kv_heads=None):
+def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None):
     shard_sizes = []
-    if eff_num_kv_heads is None:
-        eff_num_kv_heads = meta.num_kv_heads
     for i in range(mp_size):
-        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i, eff_num_kv_heads=eff_num_kv_heads))
+        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i))
     # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
     # 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 "
         f"{total_size} with tp_size={mp_size}, tp_grain_size={meta.tp_grain_size} and "
-        f"num_kv_heads={eff_num_kv_heads}.")
+        f"num_kv_heads={meta.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
index 72809d1c8c4a..d061756e522c 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -64,19 +64,3 @@ def test_shard_size_refuses_to_guess_subgroup_rank(monkeypatch):
 
     with pytest.raises(ValueError, match="group-local rank or process group"):
         get_shard_size(12, 2, AutoTPMeta())
-
-
-def test_explicit_num_kv_heads_overrides_meta():
-    # Fused Q/K/V sub-parameters are split against their own head counts, which differ from
-    # the model-wide kv-head count carried by the meta.
-    meta = AutoTPMeta(num_kv_heads=2)
-
-    assert get_shard_size_list(384, 4, meta, "self_attn.q_proj", eff_num_kv_heads=6) == [128, 128, 64, 64]
-
-
-def test_explicit_num_kv_heads_matches_meta_value():
-    expected = get_shard_size_list(384, 4, AutoTPMeta(num_kv_heads=6), "self_attn.q_proj")
-
-    actual = get_shard_size_list(384, 4, AutoTPMeta(), "self_attn.q_proj", eff_num_kv_heads=6)
-
-    assert actual == expected

From 993cf87955d55fe4d6a509c129e67f3f1232d9ba Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 15:06:35 +0800
Subject: [PATCH 10/20] Record that the Ulysses kv-head memo is still
 process-wide

Giving Ulysses its own kv-head state stops AutoTP from overwriting it, but the
memo itself is set once and never reset, so the first model to take the uneven
path decides the split for every later call in the process. Note the limitation
and what fixing it would cost, since the gather direction cannot recover the
total head count from the tensor shape and would need it threaded through
_SeqAllToAll and its backward pass.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 deepspeed/sequence/layer.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py
index 2088f9686ac0..2dea696918f5 100644
--- a/deepspeed/sequence/layer.py
+++ b/deepspeed/sequence/layer.py
@@ -17,6 +17,11 @@
 
 # Ulysses sequence parallelism keeps its own kv-head count, memoized on the first uneven
 # all-to-all. The state lives here, independent of AutoTP.
+# TODO: this is process-wide and never reset, so the first model to take the uneven path locks
+# every later Ulysses call into it -- a second model with a different head count would then be
+# split against the first one's value. Moving it onto the attention instance means threading the
+# total head count through _SeqAllToAll and its backward pass, because the gather direction
+# cannot recover it from the tensor shape alone.
 _ulysses_num_kv_heads = None
 
 

From c1c9bb9789a94ac2a5d1bde72b869d15f57fd774 Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 15:14:48 +0800
Subject: [PATCH 11/20] Cover every exclusion in the kv-head split condition

get_shard_size only takes the kv-head-aligned path for attention projections
whose dimension divides by the head count; MLP, the last linear and MoE expert
layers are excluded so they keep a near-even split. Only the attention case was
asserted, so deleting the last_linear or MoE branch left the suite green.

Fold the single assertion into a parametrization that pairs each exclusion with
the split it produces, and add the non-divisible case. Removing any one of the
four rules now fails a specific case rather than none.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 tests/unit/module_inject/test_tp_shard.py | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/tests/unit/module_inject/test_tp_shard.py b/tests/unit/module_inject/test_tp_shard.py
index d061756e522c..8f30e07101ce 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -26,11 +26,27 @@ def test_uneven_shards_without_grain_quantization():
     assert get_shard_size_list(101, 2, AutoTPMeta(), "lm_head") == [51, 50]
 
 
-def test_kv_head_shards_tile_the_dimension():
+# 6 kv heads over 4 ranks gives 2/2/1/1 heads, so an attention projection splits 384 as
+# 128/128/64/64. Every other layer kind keeps the near-even 96/96/96/96 split, either because
+# the dimension has no head structure or because even shards serve it better.
+@pytest.mark.parametrize("name,expected", [
+    ("layers.0.self_attn.q_proj", [128, 128, 64, 64]),
+    ("layers.0.mlp.dense_h_to_4h", [96, 96, 96, 96]),
+    ("lm_head", [96, 96, 96, 96]),
+    ("embed_out", [96, 96, 96, 96]),
+    ("layers.0.experts.w1", [96, 96, 96, 96]),
+])
+def test_kv_head_split_applies_only_to_attention(name, expected):
     meta = AutoTPMeta(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, meta, "layers.0.self_attn.q_proj") == [128, 128, 64, 64]
+    assert get_shard_size_list(384, 4, meta, name) == expected
+
+
+def test_kv_head_split_needs_a_divisible_dimension():
+    # 385 is not a multiple of the 6 kv heads, so there is no head-aligned split to make.
+    meta = AutoTPMeta(num_kv_heads=6)
+
+    assert get_shard_size_list(385, 4, meta, "layers.0.self_attn.q_proj") == [97, 96, 96, 96]
 
 
 def test_two_models_do_not_clobber_each_others_meta():

From e79cdf24e3b49f3e3fa45a38ba5c2c70e717926b Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 15:24:54 +0800
Subject: [PATCH 12/20] Test that a second AutoTP model does not reshard the
 first

The isolation this PR provides had no end-to-end coverage. The unit test that
claimed it built two frozen AutoTPMeta objects and called a pure function with
each, which cannot fail regardless of the implementation, so it is replaced.

Inject two models in one process with different kv-head counts and assert the
first one's split survives. 3 kv heads over 2 ranks splits 192 unevenly as
[128, 64] while 2 heads gives [96, 96], and [96, 96] is exactly what the first
model would produce once a shared kv-head count had been overwritten, so the
two outcomes are distinguishable. Re-deriving from the first model's meta after
the second is built covers the value itself rather than the frozen result.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 .../test_autotp_custom_patterns.py            | 54 +++++++++++++++++++
 tests/unit/module_inject/test_tp_shard.py     | 17 ------
 2 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
index 7573a779ecd3..418e39c1bec1 100644
--- a/tests/unit/model_parallelism/test_autotp_custom_patterns.py
+++ b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
@@ -8,6 +8,7 @@
 import deepspeed.comm as dist
 import deepspeed
 from copy import deepcopy
+from types import SimpleNamespace
 from torch import nn
 
 from unit.common import DistributedTest, preferred_dtype
@@ -925,3 +926,56 @@ def test_gqa_uneven_qkv_fused_forward(self):
         gathered_output = gather_subparam_output(tp_output, (q_size, k_size, v_size),
                                                  groups.get_tensor_model_parallel_group())
         assert_close_for_preferred_dtype(gathered_output, full_output)
+
+
+class AttentionOnlyModel(torch.nn.Module):
+    """Minimal stand-in for a decoder layer, named so AutoTP treats it as attention."""
+
+    class _Attention(torch.nn.Module):
+
+        def __init__(self, hidden_dim, proj_dim):
+            super().__init__()
+            self.q_proj = torch.nn.Linear(hidden_dim, proj_dim, bias=False)
+
+    def __init__(self, hidden_dim, proj_dim, num_kv_heads):
+        super().__init__()
+        self.self_attn = AttentionOnlyModel._Attention(hidden_dim, proj_dim)
+        # AutoTPMeta probes attributes by name, so any config-like object will do.
+        self.config = SimpleNamespace(num_key_value_heads=num_kv_heads,
+                                      num_attention_heads=num_kv_heads,
+                                      hidden_size=hidden_dim)
+
+
+class TestAutoTPMultipleModels(DistributedTest):
+    world_size = 2
+    reuse_dist_env = False
+
+    def test_a_second_model_does_not_reshard_the_first(self):
+        skip_on_device()
+        # 3 kv heads over 2 ranks is uneven ([128, 64] of 192), while 2 kv heads divides evenly
+        # ([96, 96]). The second model's split is what a clobbered kv-head count would give the
+        # first one, so the two are distinguishable.
+        partition_config = {
+            "use_default_specs": False,
+            "layer_specs": [{
+                "patterns": [".*q_proj\\.weight$"],
+                "partition_type": "column",
+            }],
+        }
+
+        teacher = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
+        teacher = apply_autotp_with_partition_config(teacher, tp_size=2, partition_config=partition_config)
+        teacher_layer = teacher.self_attn.q_proj
+        teacher_split = list(teacher_layer._partition_sizes)
+        assert teacher_split == [128, 64]
+
+        # A second model with a different kv-head count is built in the same process.
+        student = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=2)
+        student = apply_autotp_with_partition_config(student, tp_size=2, partition_config=partition_config)
+        assert list(student.self_attn.q_proj._partition_sizes) == [96, 96]
+
+        # The teacher still describes, and re-derives, its own split.
+        assert teacher_layer.tp_meta.num_kv_heads == 3
+        assert list(teacher_layer._partition_sizes) == teacher_split
+        assert get_shard_size_list(192, 2, teacher_layer.tp_meta, teacher_layer.name) == teacher_split
+        assert teacher_layer.weight.shape[0] == teacher_split[dist.get_rank()]
diff --git a/tests/unit/module_inject/test_tp_shard.py b/tests/unit/module_inject/test_tp_shard.py
index 8f30e07101ce..fe301fbbc297 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -49,23 +49,6 @@ def test_kv_head_split_needs_a_divisible_dimension():
     assert get_shard_size_list(385, 4, meta, "layers.0.self_attn.q_proj") == [97, 96, 96, 96]
 
 
-def test_two_models_do_not_clobber_each_others_meta():
-    # Each model carries its own AutoTPMeta, so loading a second model does not re-shard the
-    # first one.
-    model_a = AutoTPMeta(num_kv_heads=6, tp_grain_size=64)
-    model_b = AutoTPMeta(num_kv_heads=2, tp_grain_size=1)
-
-    a_qproj = get_shard_size_list(384, 4, model_a, "layers.0.self_attn.q_proj")
-    a_lmhead = get_shard_size_list(1001, 2, model_a, "lm_head")
-
-    # A second model is loaded into the same process.
-    _ = get_shard_size_list(384, 4, model_b, "layers.0.self_attn.q_proj")
-
-    # Model A's partition contract is unchanged.
-    assert get_shard_size_list(384, 4, model_a, "layers.0.self_attn.q_proj") == a_qproj
-    assert get_shard_size_list(1001, 2, model_a, "lm_head") == a_lmhead
-
-
 def test_process_group_resolves_noncontiguous_group_rank(monkeypatch):
     meta = AutoTPMeta(tp_grain_size=64)
     tp_group = object()

From 88acf893c00ace240ab8c095049e827b2d8f03c2 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Fri, 21 Aug 2026 18:10:13 +0800
Subject: [PATCH 13/20] Read AutoTP metadata through multimodal text_config

AutoTPMeta.from_model_config probed the config it was handed directly,
so a multimodal outer config (head counts only under text_config) lost
num_kv_heads / num_attention_heads / hidden_size and the sharding fell
back to an even-grain split that can cut through KV heads under GQA.
Descend into text_config inside from_model_config so every caller
(runtime engine, inference engine, direct AutoTP use) shares the fix.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/module_inject/tp_shard.py           |  6 +++++
 deepspeed/runtime/engine.py                   |  5 ++--
 .../test_autotp_custom_patterns.py            | 23 +++++++++++++++++++
 tests/unit/module_inject/test_tp_shard.py     | 19 +++++++++++++++
 4 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/deepspeed/module_inject/tp_shard.py b/deepspeed/module_inject/tp_shard.py
index dcffe1b6e5dc..092a5bcb4a78 100644
--- a/deepspeed/module_inject/tp_shard.py
+++ b/deepspeed/module_inject/tp_shard.py
@@ -68,6 +68,12 @@ def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta"
         """
         if model_config is None:
             return cls(tp_grain_size=tp_grain_size)
+        # Multimodal configs (e.g. vision-language models) keep the head counts only under
+        # text_config while the outer config carries the modality heads, so descend to the
+        # text config when it exists.
+        text_config = getattr(model_config, "text_config", None)
+        if text_config is not None:
+            model_config = text_config
         num_kv_heads = _kv_head_count_from(model_config)
         n_embd = None
         for name in ('n_embd', 'hidden_size'):
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 8d975caad2fa..89a05c9d96e5 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -728,10 +728,11 @@ def _apply_autotp_partitioning(self, model, tp_config):
         model_config = getattr(model, "config", None)
         # AutoTP derives its per-model sharding metadata from the model config; the warning
         # below only needs the kv-head count, which that same metadata already carries.
+        # from_model_config descends into text_config itself, so multimodal outer configs
+        # work here too.
         from deepspeed.module_inject.tp_shard import AutoTPMeta
 
-        head_config = getattr(model_config, "text_config", model_config)
-        num_kv_heads = AutoTPMeta.from_model_config(head_config).num_kv_heads
+        num_kv_heads = AutoTPMeta.from_model_config(model_config).num_kv_heads
 
         # Ranks beyond the KV head count get no attention shard. This still computes the
         # correct result because the row-parallel all-reduce sums their empty contribution,
diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
index 418e39c1bec1..d1fbb9de0bc8 100644
--- a/tests/unit/model_parallelism/test_autotp_custom_patterns.py
+++ b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
@@ -979,3 +979,26 @@ def test_a_second_model_does_not_reshard_the_first(self):
         assert list(teacher_layer._partition_sizes) == teacher_split
         assert get_shard_size_list(192, 2, teacher_layer.tp_meta, teacher_layer.name) == teacher_split
         assert teacher_layer.weight.shape[0] == teacher_split[dist.get_rank()]
+
+    def test_multimodal_outer_config_keeps_kv_head_split(self):
+        skip_on_device()
+        # Multimodal configs keep the head counts only under text_config; passing the outer
+        # config to AutoTP would lose num_kv_heads and fall back to an even-grain split that
+        # cuts through KV heads.
+        partition_config = {
+            "use_default_specs": False,
+            "layer_specs": [{
+                "patterns": [".*q_proj\\.weight$"],
+                "partition_type": "column",
+            }],
+        }
+
+        model = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
+        text_config = model.config
+        model.config = SimpleNamespace(text_config=text_config)
+
+        model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
+        layer = model.self_attn.q_proj
+        assert layer.tp_meta.num_kv_heads == 3
+        assert list(layer._partition_sizes) == [128, 64]
+        assert layer.weight.shape[0] == [128, 64][dist.get_rank()]
diff --git a/tests/unit/module_inject/test_tp_shard.py b/tests/unit/module_inject/test_tp_shard.py
index fe301fbbc297..b37de1425942 100644
--- a/tests/unit/module_inject/test_tp_shard.py
+++ b/tests/unit/module_inject/test_tp_shard.py
@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import pytest
+from types import SimpleNamespace
 
 from deepspeed.module_inject import tp_shard
 from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
@@ -26,6 +27,24 @@ def test_uneven_shards_without_grain_quantization():
     assert get_shard_size_list(101, 2, AutoTPMeta(), "lm_head") == [51, 50]
 
 
+def test_meta_descends_into_multimodal_text_config():
+    # Vision-language outer configs keep the head counts only under text_config; reading the
+    # outer config directly would lose them and fall back to an even-grain split.
+    outer = SimpleNamespace(text_config=SimpleNamespace(num_key_value_heads=4, num_attention_heads=8, hidden_size=64))
+
+    meta = AutoTPMeta.from_model_config(outer)
+
+    assert meta == AutoTPMeta(num_kv_heads=4, num_attention_heads=8, n_embd=64)
+
+
+def test_meta_descends_only_when_text_config_is_set():
+    plain = SimpleNamespace(num_key_value_heads=3, num_attention_heads=6, hidden_size=32)
+
+    meta = AutoTPMeta.from_model_config(plain)
+
+    assert meta == AutoTPMeta(num_kv_heads=3, num_attention_heads=6, n_embd=32)
+
+
 # 6 kv heads over 4 ranks gives 2/2/1/1 heads, so an attention projection splits 384 as
 # 128/128/64/64. Every other layer kind keeps the near-even 96/96/96/96 split, either because
 # the dimension has no head structure or because even shards serve it better.

From 123912f9459d4fb1d10229ab7b766f6194357302 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Sat, 22 Aug 2026 14:25:07 +0800
Subject: [PATCH 14/20] Move tests/unit/model_parallelism to
 tests/unit/v1/autotp

Per maintainer feedback the tests-to-preserve tree is tests/unit/v1, which
is also the scope the modal GPU workflow selects, so the AutoTP tests run
on multi-accelerator CI once they live there. Update the xpu workflow's
path accordingly.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 .github/workflows/xpu-max1100.yml                               | 2 +-
 .../autotp}/test_autotp_custom_patterns.py                      | 0
 .../{model_parallelism => v1/autotp}/test_autotp_training.py    | 0
 .../autotp}/test_configurable_parallel_mp.py                    | 0
 .../autotp}/test_configurable_parallel_pp.py                    | 0
 tests/unit/{model_parallelism => v1/autotp}/test_tp_plan_e2e.py | 0
 .../autotp}/test_tp_plan_real_models.py                         | 0
 7 files changed, 1 insertion(+), 1 deletion(-)
 rename tests/unit/{model_parallelism => v1/autotp}/test_autotp_custom_patterns.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_autotp_training.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_configurable_parallel_mp.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_configurable_parallel_pp.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_tp_plan_e2e.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_tp_plan_real_models.py (100%)

diff --git a/.github/workflows/xpu-max1100.yml b/.github/workflows/xpu-max1100.yml
index b9768caf6d26..cbca3a291f6e 100644
--- a/.github/workflows/xpu-max1100.yml
+++ b/.github/workflows/xpu-max1100.yml
@@ -72,7 +72,7 @@ jobs:
         export I_MPI_SHM=off
         pytest --verbose accelerator/*
         pytest --verbose autotuning/*
-        pytest --verbose model_parallelism/*
+        pytest --verbose v1/autotp/*
         pytest --verbose monitor/*
         pytest --verbose utils/*
         pytest --verbose runtime/test_ds_config_model.py
diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/v1/autotp/test_autotp_custom_patterns.py
similarity index 100%
rename from tests/unit/model_parallelism/test_autotp_custom_patterns.py
rename to tests/unit/v1/autotp/test_autotp_custom_patterns.py
diff --git a/tests/unit/model_parallelism/test_autotp_training.py b/tests/unit/v1/autotp/test_autotp_training.py
similarity index 100%
rename from tests/unit/model_parallelism/test_autotp_training.py
rename to tests/unit/v1/autotp/test_autotp_training.py
diff --git a/tests/unit/model_parallelism/test_configurable_parallel_mp.py b/tests/unit/v1/autotp/test_configurable_parallel_mp.py
similarity index 100%
rename from tests/unit/model_parallelism/test_configurable_parallel_mp.py
rename to tests/unit/v1/autotp/test_configurable_parallel_mp.py
diff --git a/tests/unit/model_parallelism/test_configurable_parallel_pp.py b/tests/unit/v1/autotp/test_configurable_parallel_pp.py
similarity index 100%
rename from tests/unit/model_parallelism/test_configurable_parallel_pp.py
rename to tests/unit/v1/autotp/test_configurable_parallel_pp.py
diff --git a/tests/unit/model_parallelism/test_tp_plan_e2e.py b/tests/unit/v1/autotp/test_tp_plan_e2e.py
similarity index 100%
rename from tests/unit/model_parallelism/test_tp_plan_e2e.py
rename to tests/unit/v1/autotp/test_tp_plan_e2e.py
diff --git a/tests/unit/model_parallelism/test_tp_plan_real_models.py b/tests/unit/v1/autotp/test_tp_plan_real_models.py
similarity index 100%
rename from tests/unit/model_parallelism/test_tp_plan_real_models.py
rename to tests/unit/v1/autotp/test_tp_plan_real_models.py

From 606ec5233811e61883358f22f9a7ea7112899d3e Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Sat, 22 Aug 2026 14:38:00 +0800
Subject: [PATCH 15/20] Address review: fix step numbering; make lm-head meta
 test two real layers

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/module_inject/replace_module.py     |  2 +-
 .../test_autotp_universal_checkpoint.py       | 21 +++++++++++++++----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py
index 4af88b696062..35b03587c6bd 100644
--- a/deepspeed/module_inject/replace_module.py
+++ b/deepspeed/module_inject/replace_module.py
@@ -307,7 +307,7 @@ def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
         # 4. Set linear policies
         _autotp.update_linear_policies()
 
-        # 6. Replace modules
+        # 5. Replace modules
         if "lm_head" in all_reduce_linears or "embed_out" in all_reduce_linears:
             return _autotp._replace_last_linear_module(module)
         return _autotp._replace_module(module)
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 17da79d86f11..2a1159359241 100644
--- a/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
+++ b/tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py
@@ -488,10 +488,23 @@ def test_lm_head_forward_uses_frozen_partition_sizes():
     layer.tp_index = 1
     frozen = layer._freeze_partition_sizes(101)
     assert frozen == (51, 50)
-
-    # A later model carries a different grain; this layer's split is frozen from its own meta.
-    other_meta = AutoTPMeta(tp_grain_size=64)
-    assert other_meta.tp_grain_size != layer.tp_meta.tp_grain_size
+    assert layer.tp_meta.tp_grain_size == 1
     layer.weight.data = torch.zeros(8, frozen[1])
 
+    # A second layer of a different model is then built and initialized in the same process:
+    # a coarser grain (vocabulary sharding) and a different kv-head count (GQA attention
+    # sharding). Neither may leak into the first layer, whose meta and frozen split describe
+    # its own model only.
+    second = LmHeadLinearAllreduce(torch.nn.Linear(101, 8, bias=False),
+                                   mp_group=None,
+                                   tp_meta=AutoTPMeta(tp_grain_size=64, num_kv_heads=1))
+    second.tp_world_size = 2
+    second.tp_index = 1
+    assert second._freeze_partition_sizes(101) == (64, 37)
+
+    assert layer.tp_meta.tp_grain_size == 1
+    assert layer.tp_meta.num_kv_heads is None
+    assert layer._freeze_partition_sizes(101) == frozen
+    assert layer._partition_sizes == frozen
+
     layer(torch.zeros(1, 1, 101))

From bb63a3e42703945edd279531d1a78b35c078f979 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Sat, 22 Aug 2026 15:33:55 +0800
Subject: [PATCH 16/20] Import print_dist for the replicated-grad-hook log line

The GPU CI run surfaced a NameError in register_replicated_grad_hooks:
print_dist was called without being imported (log_dist was). Add it to
the existing deepspeed.utils.logging import.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 deepspeed/module_inject/auto_tp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index 95800a354c29..56eda5991d81 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/deepspeed/module_inject/auto_tp.py
@@ -17,7 +17,7 @@
 from .fusedqkv_utils import require_tp_fused_qkvw
 from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 from deepspeed.utils import groups
-from deepspeed.utils.logging import log_dist
+from deepspeed.utils.logging import log_dist, print_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
 from deepspeed.module_inject.layers import _build_param_uc_restore_meta
 from deepspeed.checkpoint.constants import DS_AUTOTP_UC_META

From e9612b3c773f38247885c96d75101aad3c9958fb Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Sat, 22 Aug 2026 19:21:38 +0800
Subject: [PATCH 17/20] Revert move of model_parallelism tests out of
 tests/unit/v1

Moving them into tests/unit/v1 exposed test_tp_plan_real_models to the
modal workflow, which tests against transformers main; that main now
injects 'embedding_rowwise' into tp_plan for tied-embedding models
(#8290), so every full-suite modal run on master would fail until the
upstream drift is handled. Move the directory back for now; it will be
relocated into tests/unit/v1 once #8290 is fixed.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 .github/workflows/xpu-max1100.yml                               | 2 +-
 .../autotp => model_parallelism}/test_autotp_custom_patterns.py | 0
 .../{v1/autotp => model_parallelism}/test_autotp_training.py    | 0
 .../test_configurable_parallel_mp.py                            | 0
 .../test_configurable_parallel_pp.py                            | 0
 tests/unit/{v1/autotp => model_parallelism}/test_tp_plan_e2e.py | 0
 .../autotp => model_parallelism}/test_tp_plan_real_models.py    | 0
 7 files changed, 1 insertion(+), 1 deletion(-)
 rename tests/unit/{v1/autotp => model_parallelism}/test_autotp_custom_patterns.py (100%)
 rename tests/unit/{v1/autotp => model_parallelism}/test_autotp_training.py (100%)
 rename tests/unit/{v1/autotp => model_parallelism}/test_configurable_parallel_mp.py (100%)
 rename tests/unit/{v1/autotp => model_parallelism}/test_configurable_parallel_pp.py (100%)
 rename tests/unit/{v1/autotp => model_parallelism}/test_tp_plan_e2e.py (100%)
 rename tests/unit/{v1/autotp => model_parallelism}/test_tp_plan_real_models.py (100%)

diff --git a/.github/workflows/xpu-max1100.yml b/.github/workflows/xpu-max1100.yml
index cbca3a291f6e..b9768caf6d26 100644
--- a/.github/workflows/xpu-max1100.yml
+++ b/.github/workflows/xpu-max1100.yml
@@ -72,7 +72,7 @@ jobs:
         export I_MPI_SHM=off
         pytest --verbose accelerator/*
         pytest --verbose autotuning/*
-        pytest --verbose v1/autotp/*
+        pytest --verbose model_parallelism/*
         pytest --verbose monitor/*
         pytest --verbose utils/*
         pytest --verbose runtime/test_ds_config_model.py
diff --git a/tests/unit/v1/autotp/test_autotp_custom_patterns.py b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
similarity index 100%
rename from tests/unit/v1/autotp/test_autotp_custom_patterns.py
rename to tests/unit/model_parallelism/test_autotp_custom_patterns.py
diff --git a/tests/unit/v1/autotp/test_autotp_training.py b/tests/unit/model_parallelism/test_autotp_training.py
similarity index 100%
rename from tests/unit/v1/autotp/test_autotp_training.py
rename to tests/unit/model_parallelism/test_autotp_training.py
diff --git a/tests/unit/v1/autotp/test_configurable_parallel_mp.py b/tests/unit/model_parallelism/test_configurable_parallel_mp.py
similarity index 100%
rename from tests/unit/v1/autotp/test_configurable_parallel_mp.py
rename to tests/unit/model_parallelism/test_configurable_parallel_mp.py
diff --git a/tests/unit/v1/autotp/test_configurable_parallel_pp.py b/tests/unit/model_parallelism/test_configurable_parallel_pp.py
similarity index 100%
rename from tests/unit/v1/autotp/test_configurable_parallel_pp.py
rename to tests/unit/model_parallelism/test_configurable_parallel_pp.py
diff --git a/tests/unit/v1/autotp/test_tp_plan_e2e.py b/tests/unit/model_parallelism/test_tp_plan_e2e.py
similarity index 100%
rename from tests/unit/v1/autotp/test_tp_plan_e2e.py
rename to tests/unit/model_parallelism/test_tp_plan_e2e.py
diff --git a/tests/unit/v1/autotp/test_tp_plan_real_models.py b/tests/unit/model_parallelism/test_tp_plan_real_models.py
similarity index 100%
rename from tests/unit/v1/autotp/test_tp_plan_real_models.py
rename to tests/unit/model_parallelism/test_tp_plan_real_models.py

From 0fae83e7112a8b0758d140e847a0382cd57406d9 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Sat, 22 Aug 2026 19:51:18 +0800
Subject: [PATCH 18/20] Keep the multi-model AutoTP regression tests inside
 tests/unit/v1

The full model_parallelism directory stays out of tests/unit/v1 until the
transformers-main 'embedding_rowwise' drift (#8290) is handled, but the two
multi-model regressions this PR adds (#8231 reproduction and the multimodal
text_config path) are safe there and belong to the GPU workflow's scope.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 .../test_autotp_custom_patterns.py            |  77 +------------
 .../v1/autotp/test_autotp_multiple_models.py  | 105 ++++++++++++++++++
 2 files changed, 107 insertions(+), 75 deletions(-)
 create mode 100644 tests/unit/v1/autotp/test_autotp_multiple_models.py

diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
index d1fbb9de0bc8..6d733c152a0a 100644
--- a/tests/unit/model_parallelism/test_autotp_custom_patterns.py
+++ b/tests/unit/model_parallelism/test_autotp_custom_patterns.py
@@ -8,7 +8,6 @@
 import deepspeed.comm as dist
 import deepspeed
 from copy import deepcopy
-from types import SimpleNamespace
 from torch import nn
 
 from unit.common import DistributedTest, preferred_dtype
@@ -928,77 +927,5 @@ def test_gqa_uneven_qkv_fused_forward(self):
         assert_close_for_preferred_dtype(gathered_output, full_output)
 
 
-class AttentionOnlyModel(torch.nn.Module):
-    """Minimal stand-in for a decoder layer, named so AutoTP treats it as attention."""
-
-    class _Attention(torch.nn.Module):
-
-        def __init__(self, hidden_dim, proj_dim):
-            super().__init__()
-            self.q_proj = torch.nn.Linear(hidden_dim, proj_dim, bias=False)
-
-    def __init__(self, hidden_dim, proj_dim, num_kv_heads):
-        super().__init__()
-        self.self_attn = AttentionOnlyModel._Attention(hidden_dim, proj_dim)
-        # AutoTPMeta probes attributes by name, so any config-like object will do.
-        self.config = SimpleNamespace(num_key_value_heads=num_kv_heads,
-                                      num_attention_heads=num_kv_heads,
-                                      hidden_size=hidden_dim)
-
-
-class TestAutoTPMultipleModels(DistributedTest):
-    world_size = 2
-    reuse_dist_env = False
-
-    def test_a_second_model_does_not_reshard_the_first(self):
-        skip_on_device()
-        # 3 kv heads over 2 ranks is uneven ([128, 64] of 192), while 2 kv heads divides evenly
-        # ([96, 96]). The second model's split is what a clobbered kv-head count would give the
-        # first one, so the two are distinguishable.
-        partition_config = {
-            "use_default_specs": False,
-            "layer_specs": [{
-                "patterns": [".*q_proj\\.weight$"],
-                "partition_type": "column",
-            }],
-        }
-
-        teacher = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
-        teacher = apply_autotp_with_partition_config(teacher, tp_size=2, partition_config=partition_config)
-        teacher_layer = teacher.self_attn.q_proj
-        teacher_split = list(teacher_layer._partition_sizes)
-        assert teacher_split == [128, 64]
-
-        # A second model with a different kv-head count is built in the same process.
-        student = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=2)
-        student = apply_autotp_with_partition_config(student, tp_size=2, partition_config=partition_config)
-        assert list(student.self_attn.q_proj._partition_sizes) == [96, 96]
-
-        # The teacher still describes, and re-derives, its own split.
-        assert teacher_layer.tp_meta.num_kv_heads == 3
-        assert list(teacher_layer._partition_sizes) == teacher_split
-        assert get_shard_size_list(192, 2, teacher_layer.tp_meta, teacher_layer.name) == teacher_split
-        assert teacher_layer.weight.shape[0] == teacher_split[dist.get_rank()]
-
-    def test_multimodal_outer_config_keeps_kv_head_split(self):
-        skip_on_device()
-        # Multimodal configs keep the head counts only under text_config; passing the outer
-        # config to AutoTP would lose num_kv_heads and fall back to an even-grain split that
-        # cuts through KV heads.
-        partition_config = {
-            "use_default_specs": False,
-            "layer_specs": [{
-                "patterns": [".*q_proj\\.weight$"],
-                "partition_type": "column",
-            }],
-        }
-
-        model = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
-        text_config = model.config
-        model.config = SimpleNamespace(text_config=text_config)
-
-        model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
-        layer = model.self_attn.q_proj
-        assert layer.tp_meta.num_kv_heads == 3
-        assert list(layer._partition_sizes) == [128, 64]
-        assert layer.weight.shape[0] == [128, 64][dist.get_rank()]
+# Multi-model (teacher + student in one process) and multimodal-config regressions live in
+# tests/unit/v1/autotp/test_autotp_multiple_models.py so they run on the GPU workflow too.
diff --git a/tests/unit/v1/autotp/test_autotp_multiple_models.py b/tests/unit/v1/autotp/test_autotp_multiple_models.py
new file mode 100644
index 000000000000..5f7d2d22d550
--- /dev/null
+++ b/tests/unit/v1/autotp/test_autotp_multiple_models.py
@@ -0,0 +1,105 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Regression tests for multiple AutoTP models living in the same process.
+
+The kv-head / grain sharding state used to be process-wide globals in tp_shard, so a second
+AutoTP model overwrote the first one's metadata (#8231). AutoTPMeta is per-model now; these
+tests hold that line.
+"""
+
+import torch
+from types import SimpleNamespace
+
+from unit.common import DistributedTest
+
+import deepspeed.comm as dist
+from deepspeed.utils import groups
+from deepspeed.module_inject.auto_tp import AutoTP
+from deepspeed.module_inject.autotp_config import AutoTPConfig
+from deepspeed.module_inject.tp_shard import get_shard_size_list
+
+# Only q_proj is partitioned; matches the naming the kv-head split logic looks for.
+PARTITION_CONFIG = {
+    "use_default_specs": False,
+    "layer_specs": [{
+        "patterns": [".*q_proj\\.weight$"],
+        "partition_type": "column",
+    }],
+}
+
+
+def apply_autotp(model, tp_size, partition_config):
+    groups._init_tp_mesh_device(tensor_model_parallel_size=tp_size)
+    autotp = AutoTP(module=model,
+                    all_reduce_linears=[],
+                    prefix="",
+                    state_dict=None,
+                    linear_layer_setting=None,
+                    orig_layer_impl=None,
+                    keep_module_on_host=False,
+                    partition_config=AutoTPConfig.from_dict(partition_config),
+                    model_config=getattr(model, "config", None))
+    autotp.set_tensor_parallel_config(tp_size, groups.get_tensor_model_parallel_group())
+    autotp.update_linear_policies()
+    autotp._replace_module(model)
+    return model
+
+
+class AttentionOnlyModel(torch.nn.Module):
+    """Minimal stand-in for a decoder layer, named so AutoTP treats it as attention."""
+
+    class _Attention(torch.nn.Module):
+
+        def __init__(self, hidden_dim, proj_dim):
+            super().__init__()
+            self.q_proj = torch.nn.Linear(hidden_dim, proj_dim, bias=False)
+
+    def __init__(self, hidden_dim, proj_dim, num_kv_heads):
+        super().__init__()
+        self.self_attn = AttentionOnlyModel._Attention(hidden_dim, proj_dim)
+        # AutoTPMeta probes attributes by name, so any config-like object will do.
+        self.config = SimpleNamespace(num_key_value_heads=num_kv_heads,
+                                      num_attention_heads=num_kv_heads,
+                                      hidden_size=hidden_dim)
+
+
+class TestAutoTPMultipleModels(DistributedTest):
+    world_size = 2
+    reuse_dist_env = False
+
+    def test_a_second_model_does_not_reshard_the_first(self):
+        # 3 kv heads over 2 ranks is uneven ([128, 64] of 192), while 2 kv heads divides evenly
+        # ([96, 96]). The second model's split is what a clobbered kv-head count would give the
+        # first one, so the two are distinguishable.
+        teacher = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
+        teacher = apply_autotp(teacher, tp_size=2, partition_config=PARTITION_CONFIG)
+        teacher_layer = teacher.self_attn.q_proj
+        teacher_split = list(teacher_layer._partition_sizes)
+        assert teacher_split == [128, 64]
+
+        # A second model with a different kv-head count is built in the same process.
+        student = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=2)
+        student = apply_autotp(student, tp_size=2, partition_config=PARTITION_CONFIG)
+        assert list(student.self_attn.q_proj._partition_sizes) == [96, 96]
+
+        # The teacher still describes, and re-derives, its own split.
+        assert teacher_layer.tp_meta.num_kv_heads == 3
+        assert list(teacher_layer._partition_sizes) == teacher_split
+        assert get_shard_size_list(192, 2, teacher_layer.tp_meta, teacher_layer.name) == teacher_split
+        assert teacher_layer.weight.shape[0] == teacher_split[dist.get_rank()]
+
+    def test_multimodal_outer_config_keeps_kv_head_split(self):
+        # Multimodal configs keep the head counts only under text_config; passing the outer
+        # config to AutoTP would lose num_kv_heads and fall back to an even-grain split that
+        # cuts through KV heads.
+        model = AttentionOnlyModel(hidden_dim=64, proj_dim=192, num_kv_heads=3)
+        text_config = model.config
+        model.config = SimpleNamespace(text_config=text_config)
+
+        model = apply_autotp(model, tp_size=2, partition_config=PARTITION_CONFIG)
+        layer = model.self_attn.q_proj
+        assert layer.tp_meta.num_kv_heads == 3
+        assert list(layer._partition_sizes) == [128, 64]
+        assert layer.weight.shape[0] == [128, 64][dist.get_rank()]

From 19e2735099408061716eb7758bc06b83d298f7c8 Mon Sep 17 00:00:00 2001
From: Guokai Ma <guokai.ma@intel.com>
Date: Sun, 23 Aug 2026 22:37:16 +0800
Subject: [PATCH 19/20] Move tests/unit/model_parallelism to
 tests/unit/v1/autotp

The whole directory is multi-rank DistributedTest (world_size 2/4), so it
was never executed in PR CI: cpu runners skip it and the modal GPU
workflow's diff-driven selector only covers tests/unit/v1. Moving it under
tests/unit/v1 brings it into the modal GPU workflow's scope, so these tests
now actually run on multi-accelerator CI.

Also update the xpu-max1100 workflow to point at the new location.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>
---
 .github/workflows/xpu-max1100.yml                               | 2 +-
 .../autotp}/test_autotp_custom_patterns.py                      | 0
 .../{model_parallelism => v1/autotp}/test_autotp_training.py    | 0
 .../autotp}/test_configurable_parallel_mp.py                    | 0
 .../autotp}/test_configurable_parallel_pp.py                    | 0
 tests/unit/{model_parallelism => v1/autotp}/test_tp_plan_e2e.py | 0
 .../autotp}/test_tp_plan_real_models.py                         | 0
 7 files changed, 1 insertion(+), 1 deletion(-)
 rename tests/unit/{model_parallelism => v1/autotp}/test_autotp_custom_patterns.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_autotp_training.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_configurable_parallel_mp.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_configurable_parallel_pp.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_tp_plan_e2e.py (100%)
 rename tests/unit/{model_parallelism => v1/autotp}/test_tp_plan_real_models.py (100%)

diff --git a/.github/workflows/xpu-max1100.yml b/.github/workflows/xpu-max1100.yml
index b9768caf6d26..cbca3a291f6e 100644
--- a/.github/workflows/xpu-max1100.yml
+++ b/.github/workflows/xpu-max1100.yml
@@ -72,7 +72,7 @@ jobs:
         export I_MPI_SHM=off
         pytest --verbose accelerator/*
         pytest --verbose autotuning/*
-        pytest --verbose model_parallelism/*
+        pytest --verbose v1/autotp/*
         pytest --verbose monitor/*
         pytest --verbose utils/*
         pytest --verbose runtime/test_ds_config_model.py
diff --git a/tests/unit/model_parallelism/test_autotp_custom_patterns.py b/tests/unit/v1/autotp/test_autotp_custom_patterns.py
similarity index 100%
rename from tests/unit/model_parallelism/test_autotp_custom_patterns.py
rename to tests/unit/v1/autotp/test_autotp_custom_patterns.py
diff --git a/tests/unit/model_parallelism/test_autotp_training.py b/tests/unit/v1/autotp/test_autotp_training.py
similarity index 100%
rename from tests/unit/model_parallelism/test_autotp_training.py
rename to tests/unit/v1/autotp/test_autotp_training.py
diff --git a/tests/unit/model_parallelism/test_configurable_parallel_mp.py b/tests/unit/v1/autotp/test_configurable_parallel_mp.py
similarity index 100%
rename from tests/unit/model_parallelism/test_configurable_parallel_mp.py
rename to tests/unit/v1/autotp/test_configurable_parallel_mp.py
diff --git a/tests/unit/model_parallelism/test_configurable_parallel_pp.py b/tests/unit/v1/autotp/test_configurable_parallel_pp.py
similarity index 100%
rename from tests/unit/model_parallelism/test_configurable_parallel_pp.py
rename to tests/unit/v1/autotp/test_configurable_parallel_pp.py
diff --git a/tests/unit/model_parallelism/test_tp_plan_e2e.py b/tests/unit/v1/autotp/test_tp_plan_e2e.py
similarity index 100%
rename from tests/unit/model_parallelism/test_tp_plan_e2e.py
rename to tests/unit/v1/autotp/test_tp_plan_e2e.py
diff --git a/tests/unit/model_parallelism/test_tp_plan_real_models.py b/tests/unit/v1/autotp/test_tp_plan_real_models.py
similarity index 100%
rename from tests/unit/model_parallelism/test_tp_plan_real_models.py
rename to tests/unit/v1/autotp/test_tp_plan_real_models.py

From fff342f671f24a0ef3be05f924b6f945c368447e Mon Sep 17 00:00:00 2001
From: "Ma, Guokai" <guokai.ma@intel.com>
Date: Wed, 26 Aug 2026 15:54:23 +0800
Subject: [PATCH 20/20] Port #8299 regression tests to per-model AutoTPMeta

The unsharded-metadata regression tests merged from master drive kv-head
counts through the process-wide set_num_kv_heads() helper this branch
removes, so they fail to resolve after the merge. Set the per-model
tp_meta directly instead; the surrounding try/finally is no longer needed
because there is no global state left to restore.

Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
---
 .../v1/autotp/test_autotp_custom_patterns.py  | 21 ++++++-------------
 1 file changed, 6 insertions(+), 15 deletions(-)

diff --git a/tests/unit/v1/autotp/test_autotp_custom_patterns.py b/tests/unit/v1/autotp/test_autotp_custom_patterns.py
index b9c48fbd127c..63c714b42723 100644
--- a/tests/unit/v1/autotp/test_autotp_custom_patterns.py
+++ b/tests/unit/v1/autotp/test_autotp_custom_patterns.py
@@ -540,11 +540,8 @@ def test_update_mp_params_preserves_unsharded_high_dimensional_modules(monkeypat
     child.proj = nn.Conv3d(3, child.embed_dim, kernel_size=1)
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(3)
-    try:
-        autotp.update_mp_params(child, "model.visual.patch_embed")
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta(num_kv_heads=3)
+    autotp.update_mp_params(child, "model.visual.patch_embed")
 
     assert child.embed_dim == 12
 
@@ -565,11 +562,8 @@ def test_replace_module_preserves_metadata_for_unsharded_subtree(monkeypatch):
     visual.attn.proj = nn.Linear(12, 12)
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(None)
-    try:
-        autotp._replace_module(visual, "model.visual")
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta()
+    autotp._replace_module(visual, "model.visual")
 
     assert visual.attn.num_heads == 16
 
@@ -585,11 +579,8 @@ def test_update_mp_params_follows_actual_tp_parameter_metadata(monkeypatch):
     setattr(child.proj.weight, DS_AUTOTP_UC_META, {})
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
-    set_num_kv_heads(None)
-    try:
-        autotp.update_mp_params(child, "model.layers.0.mlp")
-    finally:
-        set_num_kv_heads(None)
+    autotp.tp_meta = AutoTPMeta()
+    autotp.update_mp_params(child, "model.layers.0.mlp")
 
     assert child.hidden_size == 6