chore: import upstream snapshot with attribution

tests/unit/runtime/test_tp_plan_extraction.py

@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
+
+
+class TestTPPlanExtraction:
+
+    def test_extract_tp_plan_from_mock_model(self):
+
+        class MockHFModel:
+
+            def __init__(self):
+                self._tp_plan = {"layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise"}
+
+        model = MockHFModel()
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "layers.*.self_attn.q_proj" in tp_plan
+        assert tp_plan["layers.*.self_attn.q_proj"] == "colwise"
+
+    def test_extract_tp_plan_from_model_with_config(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"layers.*.self_attn.q_proj": "colwise"}
+
+        class MockHFModel:
+
+            def __init__(self, config):
+                self.config = config
+
+        config = MockHFConfig()
+        model = MockHFModel(config)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "layers.*.self_attn.q_proj" in tp_plan
+
+    def test_no_tp_plan_model(self):
+        model = torch.nn.Linear(10, 10)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is None
+
+    def test_empty_tp_plan(self):
+
+        class MockHFModel:
+
+            def __init__(self):
+                self._tp_plan = {}
+
+        model = MockHFModel()
+        tp_plan = _get_hf_tp_plan(model)
+
+        # Empty _tp_plan is falsy, so falls through to config then None
+        assert tp_plan is None
+
+    def test_none_tp_plan_falls_back_to_config(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"layers.*.self_attn.q_proj": "colwise"}
+
+        class MockHFModel:
+
+            def __init__(self, config):
+                self.config = config
+                self._tp_plan = None
+
+        config = MockHFConfig()
+        model = MockHFModel(config)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "layers.*.self_attn.q_proj" in tp_plan
+
+    def test_none_tp_plan(self):
+
+        class MockHFModel:
+
+            def __init__(self):
+                pass
+
+        model = MockHFModel()
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is None
+
+    def test_priority_config_over_model(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"config_plan": "colwise"}
+
+        class MockHFModel:
+
+            def __init__(self, config):
+                self.config = config
+                self._tp_plan = {"model_plan": "colwise"}
+
+        config = MockHFConfig()
+        model = MockHFModel(config)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "config_plan" in tp_plan
+        assert "model_plan" not in tp_plan