wehub-resource-sync
570 lines
20 KiB
Python
570 lines
20 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import torch
|
|
import deepspeed.comm as dist
|
|
import deepspeed
|
|
from copy import deepcopy
|
|
from torch import nn
|
|
|
|
from unit.common import DistributedTest, preferred_dtype
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.utils import groups
|
|
from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, SubParamLinearLayer, fused_LinearLayer)
|
|
from deepspeed.module_inject.autotp_config import AutoTPConfig
|
|
from deepspeed.module_inject.auto_tp import AutoTP
|
|
|
|
|
|
def skip_on_device():
|
|
if get_accelerator().device_name() == 'xpu':
|
|
pytest.skip("XPU requires a higher version for test")
|
|
|
|
|
|
class SequentialLinearModel(torch.nn.Module):
|
|
|
|
def __init__(self, hidden_dim, nlayers=1):
|
|
super(SequentialLinearModel, self).__init__()
|
|
self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for _ in range(nlayers)])
|
|
|
|
def forward(self, x):
|
|
for layer in self.linears:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
class CustomLinearModule(torch.nn.Module):
|
|
|
|
def __init__(self, hidden_dim):
|
|
super(CustomLinearModule, self).__init__()
|
|
self.weight = torch.nn.Parameter(torch.empty(hidden_dim, hidden_dim))
|
|
self.bias = torch.nn.Parameter(torch.empty(hidden_dim))
|
|
torch.nn.init.uniform_(self.weight, -0.02, 0.02)
|
|
torch.nn.init.uniform_(self.bias, -0.02, 0.02)
|
|
|
|
def forward(self, x):
|
|
return torch.matmul(x, self.weight.transpose(-1, -2)) + self.bias
|
|
|
|
|
|
class CustomLinearModel(torch.nn.Module):
|
|
|
|
def __init__(self, hidden_dim):
|
|
super(CustomLinearModel, self).__init__()
|
|
self.custom = CustomLinearModule(hidden_dim)
|
|
|
|
def forward(self, x):
|
|
return self.custom(x)
|
|
|
|
|
|
class QKVLinearModule(torch.nn.Module):
|
|
|
|
def __init__(self, hidden_dim):
|
|
super(QKVLinearModule, self).__init__()
|
|
self.qkv_proj = torch.nn.Linear(hidden_dim, hidden_dim * 3)
|
|
|
|
def forward(self, x):
|
|
return self.qkv_proj(x)
|
|
|
|
|
|
class QKVLinearModel(torch.nn.Module):
|
|
|
|
def __init__(self, hidden_dim):
|
|
super(QKVLinearModel, self).__init__()
|
|
self.self_attn = QKVLinearModule(hidden_dim)
|
|
|
|
def forward(self, x):
|
|
return self.self_attn(x)
|
|
|
|
|
|
class DeepAttention(torch.nn.Module):
|
|
"""Mimics HF attention module with separate projection layers."""
|
|
|
|
def __init__(self, hidden_dim):
|
|
super().__init__()
|
|
self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)
|
|
self.o_proj = torch.nn.Linear(hidden_dim, hidden_dim)
|
|
|
|
def forward(self, x):
|
|
return self.o_proj(self.q_proj(x))
|
|
|
|
|
|
class DeepBlock(torch.nn.Module):
|
|
"""Mimics a single HF transformer block."""
|
|
|
|
def __init__(self, hidden_dim):
|
|
super().__init__()
|
|
self.self_attn = DeepAttention(hidden_dim)
|
|
|
|
def forward(self, x):
|
|
return self.self_attn(x)
|
|
|
|
|
|
class DeepModel(torch.nn.Module):
|
|
"""Mimics HF transformer structure: model.layers.[N].self_attn.{q,o}_proj.
|
|
|
|
This creates a 4-level-deep module hierarchy to test that _replace_module
|
|
correctly propagates the full module path during recursion.
|
|
"""
|
|
|
|
def __init__(self, hidden_dim, nlayers=2):
|
|
super().__init__()
|
|
self.layers = torch.nn.ModuleList([DeepBlock(hidden_dim) for _ in range(nlayers)])
|
|
|
|
def forward(self, x):
|
|
for layer in self.layers:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
def init_tp_engine(tp_size, partition_config=None):
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"tensor_parallel": {
|
|
"autotp_size": tp_size,
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0,
|
|
}
|
|
}
|
|
if partition_config is not None:
|
|
config_dict["tensor_parallel"]["partition_config"] = partition_config
|
|
else:
|
|
config_dict["tensor_parallel"]["partition_config"] = {
|
|
"use_default_specs": False,
|
|
"layer_specs": [{
|
|
"patterns": [".*\\.weight$"],
|
|
"partition_type": "skip",
|
|
}],
|
|
}
|
|
if preferred_dtype() is torch.float16:
|
|
config_dict["fp16"] = {"enabled": True}
|
|
elif preferred_dtype() is torch.bfloat16:
|
|
config_dict["bf16"] = {"enabled": True}
|
|
|
|
model = SequentialLinearModel(hidden_dim=8)
|
|
deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
|
|
|
|
def apply_autotp_with_partition_config(model, tp_size, partition_config):
|
|
groups._init_tp_mesh_device(tensor_model_parallel_size=tp_size)
|
|
autotp_config = AutoTPConfig.from_dict(partition_config)
|
|
autotp = AutoTP(module=model,
|
|
all_reduce_linears=[],
|
|
prefix="",
|
|
state_dict=None,
|
|
linear_layer_setting=None,
|
|
orig_layer_impl=None,
|
|
keep_module_on_host=False,
|
|
partition_config=autotp_config)
|
|
autotp.set_tensor_parallel_config(tp_size, groups.get_tensor_model_parallel_group())
|
|
autotp.update_linear_policies()
|
|
autotp._replace_module(model)
|
|
return model
|
|
|
|
|
|
def gather_subparam_output(output, subparam_sizes, mp_group):
|
|
tp_world_size = dist.get_world_size(group=mp_group)
|
|
local_sizes = [size // tp_world_size for size in subparam_sizes]
|
|
output_chunks = torch.split(output, local_sizes, dim=-1)
|
|
gathered_chunks = []
|
|
for chunk in output_chunks:
|
|
chunk = chunk.contiguous()
|
|
gathered = [torch.empty_like(chunk) for _ in range(tp_world_size)]
|
|
dist.all_gather(gathered, chunk, group=mp_group)
|
|
gathered_chunks.append(torch.cat(gathered, dim=-1))
|
|
return torch.cat(gathered_chunks, dim=-1)
|
|
|
|
|
|
def assert_close_for_preferred_dtype(actual, expected):
|
|
atol = 1e-3
|
|
rtol = 2e-2
|
|
if preferred_dtype() is torch.float32:
|
|
atol = 1e-5
|
|
rtol = 1e-5
|
|
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
|
|
|
|
|
|
class TestAutoTPCustomPatterns(DistributedTest):
|
|
world_size = 2
|
|
reuse_dist_env = False
|
|
|
|
def test_custom_pattern_replacement(self):
|
|
skip_on_device()
|
|
partition_config = {
|
|
"use_default_specs":
|
|
False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [".*linears\\.0\\.weight$"],
|
|
"partition_type": "row",
|
|
},
|
|
{
|
|
"patterns": [".*linears\\.1\\.weight$"],
|
|
"partition_type": "column",
|
|
},
|
|
{
|
|
"patterns": [".*linears\\.2\\.weight$"],
|
|
"partition_type": "skip",
|
|
},
|
|
],
|
|
}
|
|
model = SequentialLinearModel(hidden_dim=16, nlayers=3)
|
|
model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
|
|
|
|
assert isinstance(model.linears[0], LinearAllreduce)
|
|
assert isinstance(model.linears[1], LinearLayer)
|
|
assert isinstance(model.linears[2], nn.Linear)
|
|
|
|
def test_custom_patterns_applied_via_config(self):
|
|
skip_on_device()
|
|
partition_config = {
|
|
"use_default_specs":
|
|
False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [".*linears\\.0\\.weight$"],
|
|
"partition_type": "row",
|
|
},
|
|
{
|
|
"patterns": [".*linears\\.1\\.weight$"],
|
|
"partition_type": "column",
|
|
},
|
|
{
|
|
"patterns": [".*linears\\.2\\.weight$"],
|
|
"partition_type": "skip",
|
|
},
|
|
],
|
|
}
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"tensor_parallel": {
|
|
"autotp_size": 2,
|
|
"partition_config": partition_config,
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0,
|
|
}
|
|
}
|
|
if preferred_dtype() is torch.float16:
|
|
config_dict["fp16"] = {"enabled": True}
|
|
elif preferred_dtype() is torch.bfloat16:
|
|
config_dict["bf16"] = {"enabled": True}
|
|
|
|
model = SequentialLinearModel(hidden_dim=16, nlayers=3)
|
|
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
assert isinstance(engine.module.linears[0], LinearAllreduce)
|
|
assert isinstance(engine.module.linears[1], LinearLayer)
|
|
assert isinstance(engine.module.linears[2], nn.Linear)
|
|
|
|
def test_use_default_specs_false_skips_unmatched_layers(self):
|
|
skip_on_device()
|
|
# Verify unmatched layers remain unsharded when defaults are disabled.
|
|
partition_config = {
|
|
"use_default_specs":
|
|
False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [".*linears\\.0\\.weight$"],
|
|
"partition_type": "row",
|
|
},
|
|
{
|
|
"patterns": [".*linears\\.1\\.weight$"],
|
|
"partition_type": "column",
|
|
},
|
|
],
|
|
}
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"tensor_parallel": {
|
|
"autotp_size": 2,
|
|
"partition_config": partition_config,
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0,
|
|
}
|
|
}
|
|
if preferred_dtype() is torch.float16:
|
|
config_dict["fp16"] = {"enabled": True}
|
|
elif preferred_dtype() is torch.bfloat16:
|
|
config_dict["bf16"] = {"enabled": True}
|
|
|
|
model = SequentialLinearModel(hidden_dim=16, nlayers=3)
|
|
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
assert isinstance(engine.module.linears[0], LinearAllreduce)
|
|
assert isinstance(engine.module.linears[1], LinearLayer)
|
|
assert isinstance(engine.module.linears[2], nn.Linear)
|
|
|
|
def test_custom_module_replacement_with_patterns(self):
|
|
skip_on_device()
|
|
# Verify custom linear-like modules are partitioned via patterns.
|
|
partition_config = {
|
|
"use_default_specs": False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [".*custom\\.weight$"],
|
|
"partition_type": "column",
|
|
},
|
|
],
|
|
}
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"tensor_parallel": {
|
|
"autotp_size": 2,
|
|
"partition_config": partition_config,
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0,
|
|
}
|
|
}
|
|
if preferred_dtype() is torch.float16:
|
|
config_dict["fp16"] = {"enabled": True}
|
|
elif preferred_dtype() is torch.bfloat16:
|
|
config_dict["bf16"] = {"enabled": True}
|
|
|
|
model = CustomLinearModel(hidden_dim=16)
|
|
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
assert isinstance(engine.module.custom, LinearLayer)
|
|
|
|
def test_custom_pattern_disables_fused_qkv_heuristic(self):
|
|
skip_on_device()
|
|
# Use a qkv_proj name that would trigger the fused-QKV heuristic, then
|
|
# verify custom patterns override that path and preserve correctness.
|
|
torch.manual_seed(1234)
|
|
hidden_dim = 16
|
|
qkv_sizes = (hidden_dim, hidden_dim, hidden_dim)
|
|
partition_config = {
|
|
"use_default_specs":
|
|
False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [".*self_attn\\.qkv_proj\\.weight$"],
|
|
"partition_type": "column",
|
|
"shape": [list(qkv_sizes), -1],
|
|
"partition_dim": 0,
|
|
},
|
|
],
|
|
}
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-6
|
|
}
|
|
},
|
|
"tensor_parallel": {
|
|
"autotp_size": 2,
|
|
"partition_config": partition_config,
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0,
|
|
}
|
|
}
|
|
if preferred_dtype() is torch.float16:
|
|
config_dict["fp16"] = {"enabled": True}
|
|
elif preferred_dtype() is torch.bfloat16:
|
|
config_dict["bf16"] = {"enabled": True}
|
|
|
|
model = QKVLinearModel(hidden_dim=hidden_dim)
|
|
baseline = deepcopy(model).to(get_accelerator().current_device(), dtype=preferred_dtype())
|
|
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
qkv_layer = engine.module.self_attn.qkv_proj
|
|
# Custom pattern should force SubParamLinearLayer (shape-based path),
|
|
# and avoid the legacy fused-QKV heuristic despite the qkv_proj name.
|
|
assert isinstance(qkv_layer, SubParamLinearLayer)
|
|
assert not isinstance(qkv_layer, fused_LinearLayer)
|
|
|
|
assert qkv_layer.partition_dim == 0
|
|
assert qkv_layer._subparam_sizes == qkv_sizes
|
|
assert qkv_layer._orig_weight_shape == (hidden_dim * 3, hidden_dim)
|
|
|
|
qkv_layer.gather_params([qkv_layer.weight, qkv_layer.bias])
|
|
torch.testing.assert_close(qkv_layer.weight, baseline.self_attn.qkv_proj.weight)
|
|
if qkv_layer.bias is not None:
|
|
torch.testing.assert_close(qkv_layer.bias, baseline.self_attn.qkv_proj.bias)
|
|
|
|
torch.manual_seed(4321)
|
|
inputs = torch.randn(2, hidden_dim, dtype=preferred_dtype(), device=get_accelerator().current_device())
|
|
full_output = baseline(inputs)
|
|
tp_output = engine.module(inputs)
|
|
assert_close_for_preferred_dtype(tp_output, full_output)
|
|
|
|
def test_first_match_precedence(self):
|
|
skip_on_device()
|
|
partition_config = {
|
|
"use_default_specs":
|
|
False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [".*linears\\.0\\.weight$"],
|
|
"partition_type": "skip",
|
|
},
|
|
{
|
|
"patterns": [".*linears\\.0\\.weight$"],
|
|
"partition_type": "column",
|
|
},
|
|
],
|
|
}
|
|
model = SequentialLinearModel(hidden_dim=16, nlayers=1)
|
|
model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
|
|
|
|
assert isinstance(model.linears[0], nn.Linear)
|
|
|
|
def test_deep_model_full_path_propagation(self):
|
|
"""Verify _replace_module propagates accumulated paths through deep hierarchies.
|
|
|
|
Uses a 4-level-deep model (layers.N.self_attn.{q,o}_proj) with patterns
|
|
that require intermediate path components (layers.N). Without correct
|
|
full_name propagation, the recursive path is truncated and patterns
|
|
that include intermediate levels will silently fail to match.
|
|
"""
|
|
skip_on_device()
|
|
partition_config = {
|
|
"use_default_specs":
|
|
False,
|
|
"layer_specs": [
|
|
{
|
|
"patterns": [r".*layers\.\d+\.self_attn\.q_proj\.weight$"],
|
|
"partition_type": "column",
|
|
},
|
|
{
|
|
"patterns": [r".*layers\.\d+\.self_attn\.o_proj\.weight$"],
|
|
"partition_type": "row",
|
|
},
|
|
],
|
|
}
|
|
model = DeepModel(hidden_dim=16, nlayers=2)
|
|
model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
|
|
|
|
# All 4 projections (2 layers x {q_proj, o_proj}) must be replaced.
|
|
# Before the full_name fix, 0 modules were replaced because the mangled
|
|
# path "self_attn.q_proj.weight" could not match "layers.N.self_attn...".
|
|
for i in range(2):
|
|
assert isinstance(model.layers[i].self_attn.q_proj, LinearLayer), \
|
|
f"layers.{i}.self_attn.q_proj was not replaced (path propagation bug?)"
|
|
assert isinstance(model.layers[i].self_attn.o_proj, LinearAllreduce), \
|
|
f"layers.{i}.self_attn.o_proj was not replaced (path propagation bug?)"
|
|
|
|
|
|
def test_invalid_custom_shape_rejected():
|
|
bad_config = {
|
|
"layer_specs": [{
|
|
"patterns": [".*"],
|
|
"partition_type": "column",
|
|
"shape": [2, [1, 1]],
|
|
}]
|
|
}
|
|
with pytest.raises(ValueError, match="nested tuple only allowed at partition_dim"):
|
|
AutoTPConfig.from_dict(bad_config)
|
|
|
|
|
|
class TestAutoTPFusedWeights(DistributedTest):
|
|
world_size = 2
|
|
reuse_dist_env = False
|
|
|
|
def test_gate_up_fused_weight_partition(self):
|
|
skip_on_device()
|
|
init_tp_engine(tp_size=2)
|
|
|
|
hidden_dim = 8
|
|
torch.manual_seed(42)
|
|
linear = nn.Linear(hidden_dim,
|
|
hidden_dim * 2,
|
|
bias=True,
|
|
dtype=preferred_dtype(),
|
|
device=get_accelerator().current_device())
|
|
full_weight = deepcopy(linear.weight.data)
|
|
full_bias = deepcopy(linear.bias.data)
|
|
|
|
layer = SubParamLinearLayer(deepcopy(linear),
|
|
groups.get_tensor_model_parallel_group(),
|
|
shape=(2, -1),
|
|
partition_dim=0,
|
|
name="mlp.gate_up_proj")
|
|
assert layer._subparam_sizes == (hidden_dim, hidden_dim)
|
|
assert layer.weight.shape == (hidden_dim, hidden_dim)
|
|
|
|
layer.gather_params([layer.weight, layer.bias])
|
|
torch.testing.assert_close(layer.weight.data, full_weight)
|
|
torch.testing.assert_close(layer.bias.data, full_bias)
|
|
|
|
def test_gqa_uneven_qkv_fused_weight_partition(self):
|
|
skip_on_device()
|
|
init_tp_engine(tp_size=2)
|
|
|
|
hidden_dim = 8
|
|
q_size, k_size, v_size = 8, 4, 4
|
|
torch.manual_seed(123)
|
|
linear = nn.Linear(hidden_dim,
|
|
q_size + k_size + v_size,
|
|
bias=True,
|
|
dtype=preferred_dtype(),
|
|
device=get_accelerator().current_device())
|
|
full_weight = deepcopy(linear.weight.data)
|
|
full_bias = deepcopy(linear.bias.data)
|
|
|
|
layer = SubParamLinearLayer(deepcopy(linear),
|
|
groups.get_tensor_model_parallel_group(),
|
|
shape=((q_size, k_size, v_size), -1),
|
|
partition_dim=0,
|
|
name="self_attn.qkv_proj")
|
|
assert layer._subparam_sizes == (q_size, k_size, v_size)
|
|
assert layer.weight.shape == ((q_size + k_size + v_size) // 2, hidden_dim)
|
|
|
|
layer.gather_params([layer.weight, layer.bias])
|
|
torch.testing.assert_close(layer.weight.data, full_weight)
|
|
torch.testing.assert_close(layer.bias.data, full_bias)
|
|
|
|
def test_gqa_uneven_qkv_fused_forward(self):
|
|
skip_on_device()
|
|
groups._init_tp_mesh_device(tensor_model_parallel_size=2)
|
|
|
|
hidden_dim = 8
|
|
q_size, k_size, v_size = 8, 4, 4
|
|
torch.manual_seed(321)
|
|
linear = nn.Linear(hidden_dim,
|
|
q_size + k_size + v_size,
|
|
bias=True,
|
|
dtype=preferred_dtype(),
|
|
device=get_accelerator().current_device())
|
|
layer = SubParamLinearLayer(deepcopy(linear),
|
|
groups.get_tensor_model_parallel_group(),
|
|
shape=((q_size, k_size, v_size), -1),
|
|
partition_dim=0,
|
|
name="self_attn.qkv_proj")
|
|
|
|
torch.manual_seed(42)
|
|
inputs = torch.randn(2, hidden_dim, dtype=preferred_dtype(), device=get_accelerator().current_device())
|
|
full_output = linear(inputs)
|
|
tp_output = layer(inputs)
|
|
|
|
gathered_output = gather_subparam_output(tp_output, (q_size, k_size, v_size),
|
|
groups.get_tensor_model_parallel_group())
|
|
assert_close_for_preferred_dtype(gathered_output, full_output)
|