chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Test that partition_config receives correct full hierarchical module paths.
+
+The bug: AutoTP._replace_module built ``full_name`` from ``prev_name`` (the
+immediate parent only) instead of ``class_name`` (the accumulated hierarchical
+path).  Patterns like ``model.layers.0.self_attn.q_proj`` never matched
+because the name was just ``0.self_attn.q_proj``.
+"""
+
+import pytest
+import torch.nn as nn
+
+from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec
+
+
+class SubAttn(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(32, 32, bias=False)
+        self.k_proj = nn.Linear(32, 32, bias=False)
+        self.v_proj = nn.Linear(32, 32, bias=False)
+        self.o_proj = nn.Linear(32, 32, bias=False)
+
+
+class DecoderLayer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.self_attn = SubAttn()
+        self.mlp = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32))
+
+
+class DummyModel(nn.Module):
+
+    def __init__(self, num_layers=2):
+        super().__init__()
+        self.embed = nn.Embedding(100, 32)
+        self.layers = nn.ModuleList([DecoderLayer() for _ in range(num_layers)])
+        self.head = nn.Linear(32, 100, bias=False)
+
+
+def _build_config():
+    """Partition config that matches q_proj and o_proj via regex."""
+    return AutoTPConfig(layer_specs=[
+        TPLayerSpec(patterns=[r".*\.self_attn\.q_proj"], partition_type=PartitionType.COLUMN),
+        TPLayerSpec(patterns=[r".*\.self_attn\.o_proj"], partition_type=PartitionType.ROW),
+    ])
+
+
+def _capture_matched_names(model, config):
+    """Run _replace_module and capture full_name values that match a spec."""
+    matched_names = []
+    original = AutoTP._replace_with_config
+
+    def capture(self, child, full_name):
+        # Only capture if a spec actually matches
+        param_name = full_name + ".weight"
+        model_type = self._get_model_type() if hasattr(self, '_get_model_type') else None
+        spec = config.find_matching_spec(param_name, model_type)
+        if spec is not None:
+            matched_names.append(full_name)
+        return None
+
+    AutoTP._replace_with_config = capture
+    try:
+        autotp = AutoTP(
+            module=model,
+            all_reduce_linears=[],
+            prefix="model",
+            state_dict=None,
+            linear_layer_setting=None,
+            orig_layer_impl=None,
+            partition_config=config,
+        )
+        autotp._replace_module(model)
+    finally:
+        AutoTP._replace_with_config = original
+    return matched_names
+
+
+def test_partition_config_receives_full_path():
+    """Verify that pattern matching sees the full hierarchical path."""
+    model = DummyModel(num_layers=2)
+    config = _build_config()
+    matched_names = _capture_matched_names(model, config)
+
+    for layer_idx in range(2):
+        assert f"layers.{layer_idx}.self_attn.q_proj" in matched_names, \
+            f"Expected 'layers.{layer_idx}.self_attn.q_proj', got: {matched_names}"
+        assert f"layers.{layer_idx}.self_attn.o_proj" in matched_names, \
+            f"Expected 'layers.{layer_idx}.self_attn.o_proj', got: {matched_names}"
+
+
+def test_no_truncated_paths():
+    """Ensure paths are never truncated to just the immediate parent prefix."""
+    model = DummyModel(num_layers=3)
+    config = _build_config()
+    matched_names = _capture_matched_names(model, config)
+
+    for name in matched_names:
+        assert name.startswith("layers."), \
+            f"Path should start with 'layers.', got: {name}"
+        assert ".self_attn." in name, \
+            f"Path should contain '.self_attn.', got: {name}"
+        assert name.count(".") >= 3, \
+            f"Path should have at least 3 dots (layers.N.self_attn.X_proj), got: {name}"
+
+
+def test_nested_depth_correct():
+    """Verify correct count and paths at 3 layers deep."""
+    model = DummyModel(num_layers=3)
+    config = _build_config()
+    matched_names = _capture_matched_names(model, config)
+
+    expected_count = 3 * 2  # 3 layers × (q_proj + o_proj)
+    assert len(matched_names) == expected_count, \
+        f"Expected {expected_count} matches, got {len(matched_names)}: {matched_names}"
+
+    for layer_idx in range(3):
+        assert f"layers.{layer_idx}.self_attn.q_proj" in matched_names
+        assert f"layers.{layer_idx}.self_attn.o_proj" in matched_names
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.module_inject.tp_plan_converter import TPPlanConverter
+from deepspeed.module_inject.autotp_config import PartitionType
+
+
+class TestTPPlanConverter:
+
+    def test_wildcard_to_regex_basic(self):
+        assert TPPlanConverter._wildcard_to_regex("layers.*.q_proj") == r".*layers\..*\.q_proj"
+        assert TPPlanConverter._wildcard_to_regex("self_attn.*.weight") == r".*self_attn\..*\.weight"
+        assert TPPlanConverter._wildcard_to_regex("layers.*.self_attn.q_proj") == r".*layers\..*\.self_attn\.q_proj"
+
+    def test_wildcard_to_regex_special_chars(self):
+        assert TPPlanConverter._wildcard_to_regex("layers.0.q_proj") == r".*layers\.0\.q_proj"
+        assert TPPlanConverter._wildcard_to_regex("mlp.gate_proj") == r".*mlp\.gate_proj"
+
+    def test_colwise_rowwise_conversion(self):
+        hf_plan = {"layers.*.q_proj": "colwise", "layers.*.o_proj": "rowwise"}
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert len(specs) == 2
+
+        q_spec = [s for s in specs if "q_proj" in s.patterns[0]][0]
+        o_spec = [s for s in specs if "o_proj" in s.patterns[0]][0]
+
+        assert q_spec.partition_type == PartitionType.COLUMN
+        assert o_spec.partition_type == PartitionType.ROW
+
+    def test_pattern_weight_suffix(self):
+        hf_plan = {"layers.*.q_proj": "colwise"}
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert len(specs) == 1
+        assert specs[0].patterns[0].endswith(r"\.weight$")
+
+    def test_pattern_weight_suffix_already_present(self):
+        hf_plan = {"layers.*.q_proj.weight": "colwise"}
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert len(specs) == 1
+        assert specs[0].patterns[0].endswith(r"\.weight$")
+
+    def test_empty_plan(self):
+        hf_plan = {}
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert len(specs) == 0
+
+    def test_multiple_patterns(self):
+        hf_plan = {
+            "layers.*.self_attn.q_proj": "colwise",
+            "layers.*.self_attn.k_proj": "colwise",
+            "layers.*.self_attn.v_proj": "colwise",
+            "layers.*.self_attn.o_proj": "rowwise",
+            "layers.*.mlp.gate_proj": "colwise",
+            "layers.*.mlp.up_proj": "colwise",
+            "layers.*.mlp.down_proj": "rowwise",
+        }
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert len(specs) == 7
+
+        colwise_count = sum(1 for s in specs if s.partition_type == PartitionType.COLUMN)
+        rowwise_count = sum(1 for s in specs if s.partition_type == PartitionType.ROW)
+
+        assert colwise_count == 5
+        assert rowwise_count == 2
+
+    def test_pattern_matches_param_name(self):
+        import re
+
+        hf_plan = {"layers.*.self_attn.q_proj": "colwise", "layers.*.mlp.down_proj": "rowwise"}
+        specs = TPPlanConverter.convert(hf_plan)
+
+        q_pattern = [s for s in specs if "q_proj" in s.patterns[0]][0]
+        down_pattern = [s for s in specs if "down_proj" in s.patterns[0]][0]
+
+        assert re.match(q_pattern.patterns[0], "model.layers.0.self_attn.q_proj.weight")
+        assert re.match(q_pattern.patterns[0], "model.layers.10.self_attn.q_proj.weight")
+        assert not re.match(q_pattern.patterns[0], "model.layers.0.self_attn.k_proj.weight")
+
+        assert re.match(down_pattern.patterns[0], "model.layers.5.mlp.down_proj.weight")
+
+    def test_unsupported_style_returns_none(self):
+        """Unsupported styles cause convert() to return None for fallback."""
+        hf_plan = {"layers.*.q_proj": "colwise_rep", "layers.*.o_proj": "rowwise"}
+        result = TPPlanConverter.convert(hf_plan)
+        assert result is None
+
+    def test_alternate_prefixes(self):
+        """Test tp_plan with non-layers prefix"""
+        hf_plan = {
+            "model.layers.*.self_attn.q_proj": "colwise",
+            "transformer.layers.*.self_attn.o_proj": "rowwise",
+        }
+
+        layer_specs = TPPlanConverter.convert(hf_plan)
+        assert len(layer_specs) == 2
+        assert any("model\\.layers" in s.patterns[0] for s in layer_specs)
+        assert any("transformer\\.layers" in s.patterns[0] for s in layer_specs)
+
+    def test_alternate_projection_names(self):
+        """Test tp_plan with qkv and Wq/Wk/Wv style names"""
+        hf_plan = {
+            "layers.*.attn.qkv": "colwise",
+            "layers.*.attn.out_proj": "rowwise",
+            "layers.*.attn.Wq": "colwise",
+            "layers.*.attn.Wk": "colwise",
+            "layers.*.attn.Wv": "colwise",
+        }
+
+        layer_specs = TPPlanConverter.convert(hf_plan)
+        assert len(layer_specs) == 5
+        colwise_count = sum(1 for s in layer_specs if s.partition_type == PartitionType.COLUMN)
+        rowwise_count = sum(1 for s in layer_specs if s.partition_type == PartitionType.ROW)
+
+        assert colwise_count == 4  # qkv + Wq/Wk/Wv
+        assert rowwise_count == 1  # out_proj