[AutoTP] Convert HF embedding_rowwise tp_plan entries to SKIP specs

deepspeed/module_inject/tp_plan_converter.py

@@ -6,8 +6,11 @@
 from typing import List, Dict, Optional
 from .autotp_config import TPLayerSpec, PartitionType
 
-SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce"}
+SUPPORTED_STYLES = {
+    "colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce", "embedding_rowwise"
+}
 # `colwise_rep` was renamed to `colwise_gather_output` in huggingface/transformers#42809.
+# `embedding_rowwise` is injected for every tied-embedding model since huggingface/transformers#47579.
 
 
 class TPPlanConverter:
@@ -17,15 +20,21 @@ class TPPlanConverter:
     def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
         """Convert HF tp_plan to DeepSpeed layer specs.
 
-        Entries whose style is not supported are converted to SKIP specs instead of invalidating
-        the whole plan. Discarding the plan used to send models like Llama4 or Qwen3 down the
-        heuristic path, which shards by name patterns alone and has no notion of the modules the
-        plan deliberately excluded — on Llama4 it wraps the MoE router, whose forward returns a
-        tuple, and breaks the model. A SKIP spec keeps such layers untouched on purpose while the
-        supported entries are still applied.
-
-        Returns None only when no entry is convertible, so the caller can fall back to the
-        existing AutoTP path for models whose plan gives us nothing to work with.
+        A style outside SUPPORTED_STYLES raises ValueError and invalidates the whole plan, so a
+        style newly introduced upstream fails loudly and gets a deliberate decision rather than
+        being dropped on the quiet. Converting only the recognized entries could shard one half
+        of a column/row pair, and silently discarding the plan would send models like Llama4 or
+        Qwen3 down the heuristic path, which shards by name patterns alone and has no notion of
+        the modules the plan deliberately excluded. On Llama4 that path wraps the MoE router,
+        whose forward returns a tuple, and breaks the model.
+
+        Supported styles that must stay whole become SKIP specs, which keeps those layers
+        untouched on purpose while the rest of the plan is still applied: `embedding_rowwise`,
+        because vocabulary-parallel embeddings are not supported yet, and
+        `replicated_with_grad_allreduce`, which additionally sums the gradient across the group.
+
+        Returns None only for an empty plan, so the caller can fall back to the existing AutoTP
+        path for models that give us nothing to work with.
         """
         if not hf_tp_plan:
             return None
@@ -50,6 +59,10 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
                 gather_output = partition_style != "colwise"
             elif partition_style == "rowwise":
                 partition_type = PartitionType.ROW
+            elif partition_style == "embedding_rowwise":
+                # Vocabulary-parallel embeddings are not supported yet, so the embedding stays whole and
+                # the gathered-column tie fallback keeps any LM head tied to it replicated as well.
+                partition_type = PartitionType.SKIP
             else:  # replicated_with_grad_allreduce, the only other supported style
                 # The parameter stays whole; only its gradient needs summing across the group.
                 partition_type = PartitionType.SKIP

tests/unit/module_inject/test_tp_partition_config_path.py

@@ -13,6 +13,7 @@
 
 from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec
 from deepspeed.module_inject.layers import LinearLayer
+from deepspeed.module_inject.tp_plan_converter import TPPlanConverter
 
 
 class SubAttn(nn.Module):
@@ -176,6 +177,28 @@ def test_gathered_lm_head_falls_back_for_runtime_parameter_tie():
     assert model.lm_head.weight is model.embed_tokens.weight
 
 
+def test_tied_embedding_plan_leaves_lm_head_and_embedding_replicated():
+    model = OutputModel(tied=True)
+    specs = TPPlanConverter.convert({"embed_tokens": "embedding_rowwise", "lm_head": "colwise_gather_output"})
+
+    autotp = AutoTP(
+        module=model,
+        all_reduce_linears=[],
+        prefix="",
+        state_dict=None,
+        linear_layer_setting=None,
+        orig_layer_impl=None,
+        partition_config=AutoTPConfig(layer_specs=specs),
+    )
+    autotp.set_tensor_parallel_config(1, None)
+    autotp.update_linear_policies()
+    autotp._replace_module(model)
+
+    assert isinstance(model.embed_tokens, nn.Embedding)
+    assert isinstance(model.lm_head, nn.Linear)
+    assert model.lm_head.weight is model.embed_tokens.weight
+
+
 def test_gathered_lm_head_uses_column_parallel_layer_when_output_dim_is_uneven():
     model = OutputModel(tied=False)
     model.lm_head = nn.Linear(32, 101, bias=False)

tests/unit/module_inject/test_tp_plan_converter.py

@@ -126,6 +126,28 @@ def test_unsupported_style_rejects_whole_plan(self):
         with pytest.raises(ValueError, match="unsupported partition style"):
             TPPlanConverter.convert(hf_plan)
 
+    def test_tied_embedding_style_converts_to_skip(self):
+        """`embedding_rowwise` is injected for every tied-embedding model, so rejecting it would
+        strand those models on the heuristic path. Tied embeddings stay replicated here, and the
+        entry carries no gradient all-reduce because the parameter is never split."""
+        hf_plan = {
+            "layers.*.self_attn.q_proj": "colwise",
+            "layers.*.self_attn.o_proj": "rowwise",
+            "embed_tokens": "embedding_rowwise",
+            "lm_head": "colwise_gather_output",
+        }
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert len(specs) == 4
+
+        embed_spec = [s for s in specs if "embed_tokens" in s.patterns[0]][0]
+        lm_head_spec = [s for s in specs if "lm_head" in s.patterns[0]][0]
+
+        assert embed_spec.partition_type == PartitionType.SKIP
+        assert not embed_spec.grad_allreduce
+        assert lm_head_spec.partition_type == PartitionType.COLUMN
+        assert lm_head_spec.gather_output
+
     def test_alternate_prefixes(self):
         """Test tp_plan with non-layers prefix"""
         hf_plan = {