wehub-resource-sync
123 lines
4.8 KiB
Python
123 lines
4.8 KiB
Python
# 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
|