Enhance AutoTP logging for lm_head diagnostics and improve HuggingFace tp_plan extraction logic

Author: jinyouzhi

deepspeed/module_inject/auto_tp.py

@@ -16,6 +16,7 @@
 from .fusedqkv_utils import require_tp_fused_qkvw
 from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
 from deepspeed.utils import groups
+from deepspeed.utils.logging import print_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
 from .autotp_config import TPLayerSpec, AutoTPConfig, PartitionType
 
@@ -425,6 +426,17 @@ def _replace_with_config(self, child, name):
         model_type = self._get_model_type()
         spec = self.partition_config.find_matching_spec(param_name, model_type)
 
+        if any(part in ("lm_head", "embed_out") for part in name.split('.')):
+            spec_details = None
+            if spec is not None:
+                spec_details = {
+                    "patterns": spec.patterns,
+                    "partition_type": spec.partition_type.value,
+                    "gather_output": spec.gather_output,
+                }
+            print_dist(
+                f"AutoTP lm_head spec match: parameter={param_name!r}; matched_spec={spec_details!r}", ranks=[0])
+
         if spec is None:
             # No matching spec found
             if self.partition_config.strict_mode:
@@ -478,6 +490,8 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 partition_dim=spec.get_partition_dim(),
                 name=name,
             )
+        if spec.gather_output:
+            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):

deepspeed/runtime/engine.py

@@ -135,7 +135,7 @@
 from ..git_version_info import version
 
 from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler
-from deepspeed.utils.logging import print_json_dist, print_configuration, set_log_level_from_string
+from deepspeed.utils.logging import print_dist, print_json_dist, print_configuration, set_log_level_from_string
 
 from deepspeed.accelerator import get_accelerator
 
@@ -703,6 +703,43 @@ def _apply_autotp_partitioning(self, model, tp_config):
         if hasattr(tp_config, "get_partition_config_object"):
             partition_config = tp_config.get_partition_config_object()
 
+        model_config = getattr(model, "config", None)
+        base_tp_plan = getattr(model_config, "base_model_tp_plan", None) if model_config is not None else None
+        class_tp_plan = getattr(type(model), "_tp_plan", None)
+        runtime_tp_plan = getattr(model, "__dict__", {}).get("_tp_plan")
+        from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
+        hf_tp_plan = _get_hf_tp_plan(model)
+
+        def lm_head_entries(tp_plan):
+            if not isinstance(tp_plan, dict):
+                return {}
+            return {
+                pattern: style
+                for pattern, style in tp_plan.items()
+                if any(part in ("lm_head", "embed_out") for part in pattern.split('.'))
+            }
+
+        lm_head_modules = [
+            name for name, _ in model.named_modules()
+            if name and any(part in ("lm_head", "embed_out") for part in name.split('.'))
+        ]
+        selected_route = "custom partition_config" if partition_config is not None else "HuggingFace tp_plan or AutoTP"
+        model_class = f"{type(model).__module__}.{type(model).__qualname__}"
+        print_dist(
+            f"AutoTP tp_plan diagnostics: model_class={model_class}; route={selected_route}; "
+            f"base_model_tp_plan={base_tp_plan!r}; type(model)._tp_plan={class_tp_plan!r}; "
+            f"instance_tp_plan={runtime_tp_plan!r}; "
+            f"effective_tp_plan={hf_tp_plan!r}",
+            ranks=[0],
+        )
+        print_dist(
+            f"AutoTP lm_head diagnostics: modules={lm_head_modules!r}; "
+            f"base_entries={lm_head_entries(base_tp_plan)!r}; class_entries={lm_head_entries(class_tp_plan)!r}; "
+            f"runtime_entries={lm_head_entries(runtime_tp_plan)!r}; "
+            f"effective_entries={lm_head_entries(hf_tp_plan)!r}",
+            ranks=[0],
+        )
+
         if partition_config is not None:
             autotp = AutoTP(module=model,
                             all_reduce_linears=(),
@@ -723,19 +760,22 @@ def _apply_autotp_partitioning(self, model, tp_config):
             setattr(model, "ds_autotp_parsed", True)
             return
 
-        model_config = getattr(model, "config", None)
         from deepspeed.module_inject import replace_transformer_layer
-
-        from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
-
-        hf_tp_plan = _get_hf_tp_plan(model)
         if hf_tp_plan:
             from deepspeed.module_inject.tp_plan_converter import TPPlanConverter
             from deepspeed.module_inject.autotp_config import AutoTPConfig
 
             layer_specs = TPPlanConverter.convert(hf_tp_plan)
             if layer_specs is not None:
-                logger.info(f"Using HuggingFace tp_plan with {len(layer_specs)} layer specifications")
+                gathered_output_patterns = [
+                    pattern for pattern, style in hf_tp_plan.items()
+                    if style.lower() in ("colwise_rep", "colwise_gather_output")
+                ]
+                print_dist(
+                    f"Using HuggingFace tp_plan with {len(layer_specs)} layer specifications; "
+                    f"gathered column output patterns={gathered_output_patterns}",
+                    ranks=[0],
+                )
                 tp_plan_config = AutoTPConfig(tp_size=tp_size, layer_specs=layer_specs)
                 autotp = AutoTP(
                     module=model,
@@ -752,6 +792,14 @@ def _apply_autotp_partitioning(self, model, tp_config):
                 autotp._replace_module(model)
                 setattr(model, "ds_autotp_parsed", True)
                 return
+            print_dist(
+                f"AutoTP: effective HuggingFace tp_plan could not be converted; falling back to heuristic AutoTP. "
+                f"styles={sorted(set(hf_tp_plan.values()))!r}",
+                ranks=[0],
+            )
+        else:
+            print_dist("AutoTP: no effective HuggingFace tp_plan was found; falling back to heuristic AutoTP.",
+                       ranks=[0])
 
         parser_dict = AutoTP.tp_parser(model)
         for client_module, injection_policy in parser_dict:

deepspeed/runtime/tensor_parallel/config.py

@@ -149,28 +149,26 @@ def get_tensor_parallel_config(ds_config):
 def _get_hf_tp_plan(model):
     """Extract tp_plan from HuggingFace model.
 
-    Merge unique runtime entries into base_model_tp_plan. HuggingFace often adds
-    duplicate runtime entries with a 'model.' prefix, but model-level entries such
-    as lm_head may only exist in the runtime plan.
+    Merge plans from the model config, model class, and runtime instance.
+    HuggingFace may replace the instance plan with an expanded base-model plan,
+    while model-level entries such as lm_head remain only on the model class.
     """
     config = getattr(model, 'config', None)
     base_plan = getattr(config, 'base_model_tp_plan', None) if config else None
-    runtime_plan = getattr(model, '_tp_plan', None)
-
-    if base_plan:
-        if not runtime_plan:
-            return base_plan
-
-        merged_plan = dict(base_plan)
-        base_patterns = set(base_plan)
-        for pattern, style in runtime_plan.items():
-            unprefixed_pattern = pattern[len('model.'):] if pattern.startswith('model.') else pattern
-            if pattern in merged_plan or unprefixed_pattern in base_patterns:
+    class_plan = getattr(type(model), '_tp_plan', None)
+    instance_dict = getattr(model, '__dict__', {})
+    runtime_plan = instance_dict.get('_tp_plan')
+
+    merged_plan = {}
+    canonical_patterns = set()
+    for plan in (base_plan, class_plan, runtime_plan):
+        if not isinstance(plan, dict):
+            continue
+        for pattern, style in plan.items():
+            canonical_pattern = pattern[len('model.'):] if pattern.startswith('model.') else pattern
+            if pattern.startswith('model.') and canonical_pattern in canonical_patterns:
                 continue
             merged_plan[pattern] = style
-        return merged_plan
-
-    if runtime_plan:
-        return runtime_plan
+            canonical_patterns.add(canonical_pattern)
 
-    return None
+    return merged_plan or None

tests/unit/model_parallelism/test_tp_plan_real_models.py

@@ -8,6 +8,8 @@
 import deepspeed.comm as dist
 import deepspeed
 from deepspeed.accelerator import get_accelerator
+from deepspeed.module_inject.layers import LinearLayer
+from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
 from deepspeed.utils import groups
 from unit.common import DistributedTest
 
@@ -23,26 +25,30 @@ class TestTPPlanRealHFModels(DistributedTest):
     world_size = 2
 
     def test_qwen2_tp_plan_with_zero2(self):
-        """Test Qwen2 model + tp_plan + ZeRO2"""
+        """Test an untied Qwen2 LM head with gathered column output and ZeRO2."""
         skip_on_device()
 
         try:
-            from transformers import AutoModelForCausalLM, AutoConfig
+            from transformers import AutoModelForCausalLM, Qwen2Config
         except ImportError:
             pytest.skip("transformers not installed")
 
-        # Create small Qwen2 config
-        config = AutoConfig.from_pretrained(
-            "Qwen/Qwen2-7B",
+        # Construct locally so the test does not download model configuration or weights.
+        config = Qwen2Config(
             vocab_size=1000,
             hidden_size=128,
             intermediate_size=256,
             num_hidden_layers=2,
             num_attention_heads=4,
             num_key_value_heads=4,
+            tie_word_embeddings=False,
         )
 
         model = AutoModelForCausalLM.from_config(config)
+        assert model.lm_head.weight is not model.model.embed_tokens.weight
+
+        tp_plan = _get_hf_tp_plan(model)
+        assert tp_plan["lm_head"] in ("colwise_rep", "colwise_gather_output")
 
         ds_config = {
             "train_micro_batch_size_per_gpu": 1,
@@ -67,6 +73,10 @@ def test_qwen2_tp_plan_with_zero2(self):
         engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
 
         assert engine.autotp_size() == 2
+        assert isinstance(model.lm_head, LinearLayer)
+        assert model.lm_head.gather_output
+        assert model.lm_head.weight.shape == (500, 128)
+        assert model.model.embed_tokens.weight.shape == (1000, 128)
 
         # Train for a few steps
         for _ in range(3):
@@ -77,11 +87,48 @@ def test_qwen2_tp_plan_with_zero2(self):
                 group=groups.get_tensor_model_parallel_group(),
             )
             outputs = engine(input_ids, labels=input_ids)
+            assert outputs.logits.shape == (1, 16, 1000)
             engine.backward(outputs.loss)
             engine.step()
 
             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."""
+        skip_on_device()
+
+        try:
+            from transformers import AutoModelForCausalLM, Qwen2Config
+        except ImportError:
+            pytest.skip("transformers not installed")
+
+        config = Qwen2Config(
+            vocab_size=1000,
+            hidden_size=128,
+            intermediate_size=256,
+            num_hidden_layers=1,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            tie_word_embeddings=True,
+        )
+        model = AutoModelForCausalLM.from_config(config)
+        assert model.lm_head.weight is model.model.embed_tokens.weight
+
+        ds_config = {
+            "train_micro_batch_size_per_gpu": 1,
+            "tensor_parallel": {
+                "autotp_size": 2
+            },
+            "zero_optimization": {
+                "stage": 0
+            },
+        }
+
+        with pytest.raises(NotImplementedError, match="coupled vocabulary-parallel embedding"):
+            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
+
+        assert model.lm_head.weight is model.model.embed_tokens.weight
+
     def test_custom_model_with_custom_tp_plan(self):
         """Test custom model + custom tp_plan"""
         skip_on_device()

tests/unit/runtime/test_tp_plan_extraction.py

@@ -129,3 +129,35 @@ def __init__(self, config):
             "layers.*.self_attn.q_proj": "colwise",
             "lm_head": "colwise_gather_output",
         }
+
+    def test_class_lm_head_plan_survives_runtime_base_plan_override(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"layers.*.self_attn.q_proj": "colwise"}
+
+        class MockHFModel:
+            _tp_plan = {"lm_head": "colwise_gather_output"}
+
+            def __init__(self, config):
+                self.config = config
+                self._tp_plan = {
+                    "layers.*.self_attn.q_proj": "colwise",
+                    "model.layers.*.self_attn.q_proj": "colwise",
+                }
+
+        model = MockHFModel(MockHFConfig())
+        assert model._tp_plan != type(model)._tp_plan
+
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan == {
+            "layers.*.self_attn.q_proj": "colwise",
+            "lm_head": "colwise_gather_output",
+        }
+
+    def test_extract_class_only_tp_plan(self):
+
+        class MockHFModel:
+            _tp_plan = {"lm_head": "colwise_rep"}
+
+        assert _get_hf_tp_plan(MockHFModel()) == {"lm_head": "colwise_rep"}