Fix AutoTP training lm_head routing

Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>

Author: gaoxiaomo

deepspeed/module_inject/auto_tp.py

@@ -394,12 +394,15 @@ def _replace(self, child, name, conv_linear_layer):
         down_proj = False
         if 'down_proj' in name:
             down_proj = True
-        if name in self.all_reduce_linears or arctic_w2_all_reduce_linear or down_proj:
+        legacy_lm_head = name == "lm_head" or name == 'embed_out'
+        training_column_lm_head = is_autotp_training_mode() and legacy_lm_head
+        if (name in self.all_reduce_linears or arctic_w2_all_reduce_linear
+                or down_proj) and not training_column_lm_head:
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
                 return Conv_LinearALlreduce(child, self.mp_group, name=name)
-            elif name == "lm_head" or name == 'embed_out':
+            elif legacy_lm_head:
                 return LmHeadLinearAllreduce(child, self.mp_group)
 
             return LinearAllreduce(child, self.mp_group, name=name)
@@ -451,8 +454,7 @@ 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)
-        # Check for lm_head / embed_out
-        if name == "lm_head" or name == 'embed_out':
+        if (name == "lm_head" or name == 'embed_out') and not is_autotp_training_mode():
             return LmHeadLinearAllreduce(module, self.mp_group)
 
         if spec.shape is not None:

tests/unit/module_inject/test_tp_partition_config_path.py

@@ -8,11 +8,13 @@
 because the name was just ``0.self_attn.q_proj``.
 """
 
+from contextlib import contextmanager
+
 import pytest
 import torch.nn as nn
 
 from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec
-from deepspeed.module_inject.layers import LinearLayer
+from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, set_autotp_mode
 from deepspeed.module_inject.tp_plan_converter import TPPlanConverter
 
 
@@ -158,6 +160,49 @@ def _build_gathered_lm_head_autotp(model, mp_size=1):
     return autotp
 
 
+def _build_legacy_lm_head_autotp(model):
+    autotp = AutoTP(
+        module=model,
+        all_reduce_linears=("lm_head", "embed_out"),
+        prefix="",
+        state_dict=None,
+        linear_layer_setting=(nn.Linear, nn.Embedding),
+        orig_layer_impl=None,
+    )
+    autotp.mp_size = 1
+    autotp.mp_group = None
+    autotp.update_linear_policies()
+    return autotp
+
+
+def _build_row_output_head_autotp(model, head="lm_head"):
+    config = AutoTPConfig(layer_specs=[
+        TPLayerSpec(patterns=[rf".*{head}\.weight$"], partition_type=PartitionType.ROW),
+    ])
+    autotp = AutoTP(
+        module=model,
+        all_reduce_linears=(),
+        prefix="",
+        state_dict=None,
+        linear_layer_setting=(nn.Linear, nn.Embedding),
+        orig_layer_impl=None,
+        partition_config=config,
+    )
+    autotp.mp_size = 1
+    autotp.mp_group = None
+    autotp.update_linear_policies()
+    return autotp
+
+
+@contextmanager
+def _training_mode():
+    set_autotp_mode(training=True)
+    try:
+        yield
+    finally:
+        set_autotp_mode(training=False)
+
+
 def test_gathered_lm_head_uses_column_parallel_layer_when_untied():
     model = OutputModel(tied=False)
     _build_gathered_lm_head_autotp(model)._replace_module(model)
@@ -209,5 +254,43 @@ def test_gathered_lm_head_uses_column_parallel_layer_when_output_dim_is_uneven()
     assert model.lm_head.gather_output
 
 
+@pytest.mark.parametrize("head", ["lm_head", "embed_out"])
+def test_legacy_output_head_defaults_to_column_parallel_during_training(head):
+    model = OutputModel(tied=False)
+    if head == "embed_out":
+        model.embed_out = model.lm_head
+        del model.lm_head
+
+    with _training_mode():
+        _build_legacy_lm_head_autotp(model)._replace_last_linear_module(model)
+
+    output_head = getattr(model, head)
+    assert isinstance(output_head, LinearLayer)
+    assert not output_head.gather_output
+
+
+def test_legacy_lm_head_keeps_inference_allreduce_routing():
+    model = OutputModel(tied=False)
+    _build_legacy_lm_head_autotp(model)._replace_last_linear_module(model)
+
+    assert isinstance(model.lm_head, LmHeadLinearAllreduce)
+
+
+def test_explicit_row_parallel_lm_head_is_not_overridden_by_its_name():
+    model = OutputModel(tied=False)
+    with _training_mode():
+        _build_row_output_head_autotp(model)._replace_module(model)
+
+    assert isinstance(model.lm_head, LinearAllreduce)
+    assert not isinstance(model.lm_head, LmHeadLinearAllreduce)
+
+
+def test_explicit_row_parallel_lm_head_keeps_inference_specialization():
+    model = OutputModel(tied=False)
+    _build_row_output_head_autotp(model)._replace_module(model)
+
+    assert isinstance(model.lm_head, LmHeadLinearAllreduce)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])