Refactor AutoTP gathered column tie handling and update related tests for fallback behavior

Author: jinyouzhi

deepspeed/module_inject/auto_tp.py

@@ -215,7 +215,8 @@ def __init__(self,
         self.linear_policies = None
         self.conv_linear_layer = False
         self.partition_config = partition_config
-        self._gathered_column_ties_validated = False
+        self._gathered_column_tie_fallbacks_configured = False
+        self._tied_gathered_column_module_names = set()
         TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host)
 
     def in_module_list(module, module_list):
@@ -494,19 +495,31 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
             print_dist(f"AutoTP: replacing '{name}' with LinearLayer(gather_output=True)", ranks=[0])
         return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
 
-    def _validate_gathered_column_ties(self):
-        """Reject gathered output layers that share their Parameter with an embedding."""
-        if self._gathered_column_ties_validated or self.partition_config is None:
+    def _configure_gathered_column_tie_fallbacks(self):
+        """Configure a replicated fallback for gathered output layers tied to embeddings."""
+        if self._gathered_column_tie_fallbacks_configured or self.partition_config is None:
             return
 
-        embeddings = [(name, module) for name, module in self.module.named_modules()
+        named_modules = list(self.module.named_modules())
+        embeddings = [(name, module) for name, module in named_modules
                       if isinstance(module, nn.Embedding) and hasattr(module, "weight")]
+        get_input_embeddings = getattr(self.module, "get_input_embeddings", None)
+        if callable(get_input_embeddings):
+            input_embedding = get_input_embeddings()
+            if input_embedding is not None and hasattr(input_embedding, "weight"):
+                input_embedding_name = next(
+                    (name for name, module in named_modules if module is input_embedding),
+                    None,
+                )
+                is_known_embedding = any(module is input_embedding for _, module in embeddings)
+                if input_embedding_name is not None and not is_known_embedding:
+                    embeddings.append((input_embedding_name, input_embedding))
         if not embeddings:
-            self._gathered_column_ties_validated = True
+            self._gathered_column_tie_fallbacks_configured = True
             return
 
         model_type = self._get_model_type()
-        for module_name, module in self.module.named_modules():
+        for module_name, module in named_modules:
             if not module_name or isinstance(module, nn.Embedding) or not hasattr(module, "weight"):
                 continue
 
@@ -521,13 +534,14 @@ def _validate_gathered_column_ties(self):
             if spec is None or spec.partition_type != PartitionType.COLUMN or not spec.gather_output:
                 continue
 
-            raise NotImplementedError(
-                f"AutoTP cannot apply gathered column parallelism to '{module_name}.weight' because it is tied "
-                f"to embedding '{tied_embedding_name}.weight' (both reference the same Parameter). Tied output "
-                "weights require a coupled vocabulary-parallel embedding, which is not supported yet. Untie "
-                "the weights before enabling gathered column parallelism.")
+            self._tied_gathered_column_module_names.update((module_name, tied_embedding_name))
+            print_dist(
+                f"AutoTP: '{module_name}.weight' is tied to '{tied_embedding_name}.weight'; leaving both modules "
+                "replicated because coupled vocabulary-parallel embedding is not supported yet.",
+                ranks=[0],
+            )
 
-        self._gathered_column_ties_validated = True
+        self._gathered_column_tie_fallbacks_configured = True
 
     def _get_model_type(self) -> Optional[str]:
         """Extract model type from module config or class name."""
@@ -623,7 +637,7 @@ def _replace_autoep_shared_experts(self, autoep_layer, autoep_name):
 
     def _replace_module(self, r_module, prev_name='', prev_class_name=''):
         if prev_name == '' and prev_class_name == '':
-            self._validate_gathered_column_ties()
+            self._configure_gathered_column_tie_fallbacks()
 
         for name, child in r_module.named_children():
             if getattr(child, "_is_autoep_layer", False):
@@ -650,7 +664,9 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
             # instead of linear_policies. This keeps all pattern logic centralized here.
             if self.partition_config is not None:
                 full_name = class_name + '.' + name if class_name else name
-                if isinstance(child, nn.Embedding):
+                if full_name in self._tied_gathered_column_module_names:
+                    continue
+                elif isinstance(child, nn.Embedding):
                     # Check if embedding matches any pattern
                     param_name = full_name + ".weight"
                     model_type = self._get_model_type()

docs/_tutorials/autotp-training.md

@@ -125,10 +125,12 @@ shards so every tensor-parallel rank receives the complete output.
 
 Gathered column parallelism currently supports untied output layers. If an
 output layer such as `lm_head` shares the same runtime `Parameter` object with
-an embedding, DeepSpeed raises an error before replacing either layer. Preserving
-that tie requires a coupled vocabulary-parallel embedding, which is not yet
-implemented. This check uses actual `Parameter` identity rather than model
-configuration metadata such as `tie_word_embeddings`.
+an embedding, DeepSpeed leaves both modules replicated and applies tensor
+parallelism to the remaining matched layers. This preserves the tie without
+silently cloning the weight, but does not reduce the embedding or output-layer
+memory footprint. A coupled vocabulary-parallel embedding is required to shard
+the tied weight and is not yet implemented. This fallback uses actual `Parameter`
+identity rather than model configuration metadata such as `tie_word_embeddings`.
 
 Additional HuggingFace types such as `local_colwise` and `local_rowwise` are
 not yet handled and fall back to AutoTP preset-based partitioning.

tests/unit/model_parallelism/test_tp_plan_real_models.py

@@ -93,8 +93,8 @@ def test_qwen2_tp_plan_with_zero2(self):
 
             assert not torch.isnan(outputs.loss)
 
-    def test_qwen2_tied_lm_head_is_rejected(self):
-        """Test that the gathered LM-head plan rejects an actual Qwen2 Parameter tie."""
+    def test_qwen2_tied_lm_head_falls_back_to_replicated(self):
+        """Test that an actual Qwen2 Parameter tie remains replicated."""
         skip_on_device()
 
         try:
@@ -124,10 +124,23 @@ def test_qwen2_tied_lm_head_is_rejected(self):
             },
         }
 
-        with pytest.raises(NotImplementedError, match="coupled vocabulary-parallel embedding"):
-            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
 
+        assert engine.autotp_size() == 2
+        assert isinstance(model.model.layers[0].self_attn.q_proj, LinearLayer)
+        assert isinstance(model.model.embed_tokens, torch.nn.Embedding)
+        assert isinstance(model.lm_head, torch.nn.Linear)
         assert model.lm_head.weight is model.model.embed_tokens.weight
+        assert model.lm_head.weight.shape == (1000, 128)
+
+        input_ids = torch.randint(0, 1000, (1, 16)).to(get_accelerator().current_device_name())
+        dist.broadcast(
+            input_ids,
+            src=groups.get_tensor_model_parallel_src_rank(),
+            group=groups.get_tensor_model_parallel_group(),
+        )
+        outputs = engine(input_ids)
+        assert outputs.logits.shape == (1, 16, 1000)
 
     def test_custom_model_with_custom_tp_plan(self):
         """Test custom model + custom tp_plan"""

tests/unit/module_inject/test_tp_partition_config_path.py

@@ -165,12 +165,11 @@ def test_gathered_lm_head_uses_column_parallel_layer_when_untied():
     assert model.lm_head.gather_output
 
 
-def test_gathered_lm_head_rejects_runtime_parameter_tie_before_replacement():
+def test_gathered_lm_head_falls_back_for_runtime_parameter_tie():
     model = OutputModel(tied=True)
     assert model.lm_head.weight is model.embed_tokens.weight
 
-    with pytest.raises(NotImplementedError, match="coupled vocabulary-parallel embedding"):
-        _build_gathered_lm_head_autotp(model)._replace_module(model)
+    _build_gathered_lm_head_autotp(model)._replace_module(model)
 
     assert isinstance(model.embed_tokens, nn.Embedding)
     assert isinstance(model.lm_head, nn.Linear)