Fix AutoTP metadata updates for unsharded modules (#8299)

## Summary

- defer module metadata updates until recursive AutoTP replacement has
completed
- update logical dimensions only when a subtree contains parameters
carrying AutoTP TP metadata
- preserve metadata for replicated vision modules and high-dimensional
weights while retaining existing sharded Linear behavior
- add regression coverage for PatchEmbed-style Conv3d wrappers,
unsharded attention subtrees, and genuinely sharded parameters

Fixes #8285.

## Problem

`update_mp_params()` previously ran before Linear replacement and scaled
logical attributes for every traversed module. In multimodal models this
changed replicated vision metadata, such as Qwen2-VL PatchEmbed
`embed_dim`, even though the corresponding Conv3d weight was not
partitioned. The metadata then disagreed with the full weight shape and
broke vision forward.

## Validation

Current master baseline `e70b90352e93e07862c45ebb946d75b813f025a2`
reproduces on 2-GPU NCCL with `patch_embed.embed_dim=60` while
`proj.out_channels=120`.

- CPU targeted: `7 passed, 20 deselected`
- 2-process Gloo metadata check: unsharded `num_heads=16`, sharded
`hidden_size=6` on both ranks
- 2-GPU custom-pattern suite: `27 passed`
- 4-GPU AutoTP/ZeRO adjacent suite: `33 passed, 2 skipped, 24
deselected`
- synthetic training on 1/2/3/5 RTX 3090 GPUs: forward, backward, and
optimizer step all complete with zero cross-rank output, PatchEmbed
gradient, and PatchEmbed parameter difference
- 5-GPU synthetic, 3 runs with 3 warmup and 20 measured steps each:
`5.008 ms` mean step latency, `409,216 tokens/s`, `28,170,752` peak
allocated bytes, `52,428,800` peak reserved bytes
- Qwen2-VL-2B AutoTP=2 vision forward: PatchEmbed `1280 -> 1280`, vision
heads `16 -> 16`, language q_proj `[1536,1536] -> [768,1536]`, zero
cross-rank output difference
- changed-file pre-commit hooks, `DS_BUILD_OPS=0 python setup.py check`,
`compileall`, and `git diff --check` pass

---------

Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

Author: zupengwang

deepspeed/module_inject/auto_tp.py

@@ -625,12 +625,17 @@ def hook(grad):
     def update_mp_params(self, child, name=None):
         if getattr(child, "replaced", False) == True:
             return
+        if not any(hasattr(param, DS_AUTOTP_UC_META) for param in child.parameters()):
+            setattr(child, "replaced", True)
+            return
         tp_index = dist.get_rank(group=self.mp_group) if self.mp_group is not None else 0
-        # Fused-expert containers (Mixtral/Llama4/Qwen-MoE style) hold their weights as 3D
-        # parameters that AutoTP does not shard, so their dimension attributes must stay whole.
-        # Halving e.g. Llama4TextExperts.hidden_size while its weights keep the full size breaks
-        # the experts' batched matmul.
-        if any(param.dim() >= 3 for param in child.parameters(recurse=False)):
+        # AutoTP does not shard high-dimensional weights, so their dimension attributes must stay whole.
+        # Some wrappers, such as Qwen2-VL PatchEmbed, keep that weight in a direct child module.
+        has_unsharded_high_dimensional_param = any(param.dim() >= 3 for param in child.parameters(recurse=False))
+        if not has_unsharded_high_dimensional_param:
+            has_unsharded_high_dimensional_param = any(param.dim() >= 3 for module in child.children()
+                                                       for param in module.parameters(recurse=False))
+        if has_unsharded_high_dimensional_param:
             setattr(child, "replaced", True)
             return
         param_list = [
@@ -683,8 +688,8 @@ def _replace_autoep_shared_experts(self, autoep_layer, autoep_name):
                 key = next(lp for lp in self.linear_policies if isinstance(child, lp))
                 setattr(autoep_layer, child_name, self.linear_policies[key](child, full_name, self.conv_linear_layer))
             else:
-                self.update_mp_params(child, full_name)
                 self._replace_module(child, full_name, "")
+                self.update_mp_params(child, full_name)
 
     def _replace_module(self, r_module, prev_name='', prev_class_name=''):
         if prev_name == '' and prev_class_name == '':
@@ -732,8 +737,8 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
                     if new_child is not None:
                         setattr(r_module, name, new_child)
                 else:
-                    self.update_mp_params(child, full_name)
                     self._replace_module(child, name, class_name)
+                    self.update_mp_params(child, full_name)
             # Traditional path: use linear_policies for type-based routing
             elif child.__class__ in self.linear_policies:
                 setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name,
@@ -750,8 +755,10 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
                 setattr(r_module, name, self.linear_policies[key](child, prev_name + '.' + name,
                                                                   self.conv_linear_layer))
             else:
-                self.update_mp_params(child, name)
                 self._replace_module(child, name, class_name)
+                # Descendants have now been replaced and carry universal-checkpoint TP metadata.
+                # Keep logical dimensions whole when no parameter in this subtree was actually sharded.
+                self.update_mp_params(child, name)
         return r_module
 
     @staticmethod

tests/unit/model_parallelism/test_autotp_custom_patterns.py

@@ -16,7 +16,8 @@
 from deepspeed.module_inject.layers import (GateUpPack_LinearLayer, LinearAllreduce, LinearLayer,
                                             SubParamLinearAllreduce, SubParamLinearLayer, fused_LinearLayer)
 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.checkpoint.constants import (DS_AUTOTP_UC_META, 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.auto_tp import AutoTP
@@ -496,6 +497,8 @@ def test_update_mp_params_uses_group_local_rank(monkeypatch):
     autotp.mp_group = tp_group
     autotp.mp_size = 2
     child = nn.Module()
+    child.proj = nn.Linear(1, 1, bias=False)
+    setattr(child.proj.weight, DS_AUTOTP_UC_META, {})
     child.num_heads = 12
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
@@ -515,6 +518,8 @@ def test_update_mp_params_shards_attributes_like_their_weights(monkeypatch):
     autotp.mp_group = tp_group
     autotp.mp_size = 2
     child = nn.Module()
+    child.proj = nn.Linear(1, 1, bias=False)
+    setattr(child.proj.weight, DS_AUTOTP_UC_META, {})
     child.hidden_size = 12
 
     monkeypatch.setattr(dist, "get_rank", lambda group=None: 1 if group is tp_group else 0)
@@ -529,6 +534,70 @@ def test_update_mp_params_shards_attributes_like_their_weights(monkeypatch):
     assert child.hidden_size == get_shard_size(12, 2, "model.layers.0.mlp", rank=1)
 
 
+def test_update_mp_params_preserves_unsharded_high_dimensional_modules(monkeypatch):
+    tp_group = object()
+    autotp = object.__new__(AutoTP)
+    autotp.mp_group = tp_group
+    autotp.mp_size = 2
+    child = nn.Module()
+    child.embed_dim = 12
+    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)
+
+    assert child.embed_dim == 12
+
+
+def test_replace_module_preserves_metadata_for_unsharded_subtree(monkeypatch):
+    tp_group = object()
+    autotp = object.__new__(AutoTP)
+    autotp.mp_group = tp_group
+    autotp.mp_size = 2
+    autotp.partition_config = None
+    autotp.state_dict = None
+    autotp.linear_policies = {}
+    autotp.prefix = ""
+    visual = nn.Module()
+    visual.attn = nn.Module()
+    visual.attn.num_heads = 16
+    visual.attn.qkv = nn.Linear(12, 36)
+    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)
+
+    assert visual.attn.num_heads == 16
+
+
+def test_update_mp_params_follows_actual_tp_parameter_metadata(monkeypatch):
+    tp_group = object()
+    autotp = object.__new__(AutoTP)
+    autotp.mp_group = tp_group
+    autotp.mp_size = 2
+    child = nn.Module()
+    child.hidden_size = 12
+    child.proj = nn.Linear(12, 12)
+    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)
+
+    assert child.hidden_size == 6
+
+
 def test_sliced_embedding_publishes_row_partition_metadata(monkeypatch):
     tp_group = object()
     autotp = object.__new__(AutoTP)