mirrors/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/model_parallelism/test_autotp_custom_patterns.py
+569
View File
@@ -0,0 +1,569 @@
# 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)
tests/unit/model_parallelism/test_autotp_training.py
+822
View File
@@ -0,0 +1,822 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import pytest
import deepspeed.comm as dist
import torch
import math
from copy import deepcopy
from unit.common import DistributedTest, preferred_dtype
import deepspeed
from deepspeed.accelerator import get_accelerator
from unit.simple_model import SimpleModel, random_dataloader
from deepspeed.utils import groups
from contextlib import contextmanager
from torch import nn
from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, set_autotp_mode, is_autotp_training_mode
from unit.checkpoint.common import compare_lr_scheduler_states, compare_optimizer_states
import os
from deepspeed.runtime.utils import is_model_parallel_parameter
def skip_on_device():
if get_accelerator().device_name() == 'xpu':
pytest.skip("XPU requires a higher version for test")
def reset_tp_model_init_state():
deepspeed._TP_MODEL_INIT_ARGS = None
set_autotp_mode(training=False)
def _reset_tp_groups(monkeypatch):
monkeypatch.setattr(groups, "_DATA_PARALLEL_GROUP", None)
monkeypatch.setattr(groups, "_MODEL_PARALLEL_GROUP", None)
monkeypatch.setattr(groups, "_TENSOR_MODEL_PARALLEL_GROUP", None)
def _patch_tp_group_creation(monkeypatch, *, rank=2, initialize_mesh_device=None):
new_group_calls = []
def fake_new_group(ranks):
ranks = tuple(ranks)
new_group_calls.append(ranks)
return ranks
_reset_tp_groups(monkeypatch)
monkeypatch.setattr(groups.dist, "get_world_size", lambda group=None: 4)
monkeypatch.setattr(groups.dist, "get_rank", lambda group=None: rank)
monkeypatch.setattr(groups.dist, "new_group", fake_new_group)
monkeypatch.setattr(groups, "log_dist", lambda *args, **kwargs: None)
if initialize_mesh_device is not None:
monkeypatch.setattr(groups.dist, "initialize_mesh_device", initialize_mesh_device)
return new_group_calls
def test_init_tp_mesh_device_debug_detail_uses_explicit_groups(monkeypatch):
def fail_initialize_mesh_device(*args, **kwargs):
raise AssertionError("DeviceMesh should be skipped when TORCH_DISTRIBUTED_DEBUG=DETAIL")
new_group_calls = _patch_tp_group_creation(monkeypatch, initialize_mesh_device=fail_initialize_mesh_device)
monkeypatch.setenv("TORCH_DISTRIBUTED_DEBUG", "DETAIL")
data_parallel_group, tensor_parallel_group = groups._init_tp_mesh_device(tensor_model_parallel_size=2)
assert new_group_calls == [(0, 2), (1, 3), (0, 1), (2, 3)]
assert data_parallel_group == (0, 2)
assert tensor_parallel_group == (2, 3)
assert groups.get_data_parallel_group() == (0, 2)
assert groups.get_tensor_model_parallel_group() == (2, 3)
def test_init_tp_mesh_device_split_error_falls_back_to_explicit_groups(monkeypatch):
def raise_split_error(*args, **kwargs):
raise RuntimeError(groups._DEVICE_MESH_SPLIT_UNSUPPORTED)
new_group_calls = _patch_tp_group_creation(monkeypatch, initialize_mesh_device=raise_split_error)
monkeypatch.delenv("TORCH_DISTRIBUTED_DEBUG", raising=False)
data_parallel_group, tensor_parallel_group = groups._init_tp_mesh_device(tensor_model_parallel_size=2)
assert new_group_calls == [(0, 2), (1, 3), (0, 1), (2, 3)]
assert data_parallel_group == (0, 2)
assert tensor_parallel_group == (2, 3)
assert groups.get_data_parallel_group() == (0, 2)
assert groups.get_tensor_model_parallel_group() == (2, 3)
class DummyMPU:
def __init__(self, tp_world_size=1):
self.rank = dist.get_rank()
self.world_size = dist.get_world_size()
self.tp_world_size = tp_world_size
self.dp_group = dist.get_world_group()
self.tp_group = dist.get_world_group()
def get_model_parallel_rank(self):
return self.rank % self.tp_world_size
def get_model_parallel_world_size(self):
return self.tp_world_size
def get_data_parallel_rank(self):
return self.rank // self.tp_world_size
def get_data_parallel_world_size(self):
return self.world_size // self.tp_world_size
def get_data_parallel_group(self):
return self.dp_group
def get_model_parallel_group(self):
return self.tp_group
class SequentialLinearModel(torch.nn.Module):
def __init__(self, hidden_dim, empty_grad=False, nlayers=1):
super(SequentialLinearModel, self).__init__()
self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for _ in range(nlayers)])
if empty_grad:
self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
self.empty_grad = empty_grad
def forward(self, x, y):
if len(self.linears) == 1:
x = self.linears[0](x)
else:
for i, l in enumerate(self.linears):
x = self.linears[i](x)
return self.cross_entropy_loss(x, y)
@contextmanager
def should_assert_with_msg(expected_message):
try:
yield
except AssertionError as e:
if dist.get_rank() == 0:
print(expected_message)
print(str(e))
if str(e) == expected_message:
pass
else:
raise e
else:
raise AssertionError(f"Expected AssertionError with message '{expected_message}' "
"but no exception was raised.")
@pytest.mark.parametrize("tp_size", [2, 4])
class TestTpParallelStates(DistributedTest):
world_size = 4
def test(self, tp_size: int):
skip_on_device()
dp_size = 4 / tp_size
hidden_dim = 128
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": tp_size,
"partition_config": {
"use_default_specs": False,
"layer_specs": [{
"patterns": [".*\\.weight$"],
"partition_type": "skip",
}],
}
},
"zero_optimization": {
"stage": 0
}
}
model = SimpleModel(hidden_dim=hidden_dim)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
assert groups.get_tensor_model_parallel_world_size() == tp_size
assert groups.get_data_parallel_world_size() == dp_size
class TestTpModelInitCompatibility(DistributedTest):
world_size = 4
reuse_dist_env = False
def test_tp_model_init_merges_config(self):
skip_on_device()
reset_tp_model_init_state()
model = SimpleModel(hidden_dim=8)
deepspeed.tp_model_init(model, tp_size=1, dtype=preferred_dtype())
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"zero_optimization": {
"stage": 0,
}
}
engine, _, _, _ = deepspeed.initialize(model=model,
model_parameters=model.parameters(),
config=config_dict,
mpu=DummyMPU())
assert engine.autotp_size() == 1
assert is_autotp_training_mode()
def test_tp_model_init_config_autotp_size_mismatch(self):
skip_on_device()
reset_tp_model_init_state()
model = SimpleModel(hidden_dim=8)
deepspeed.tp_model_init(model, tp_size=1, dtype=preferred_dtype())
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": 2,
},
"zero_optimization": {
"stage": 0,
}
}
with pytest.raises(ValueError, match="tensor_parallel.autotp_size"):
deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict, mpu=DummyMPU())
def test_tp_model_init_autocreates_tp_group(self):
skip_on_device()
reset_tp_model_init_state()
# Verify tp_model_init creates TP groups when no mpu is provided.
model = SimpleModel(hidden_dim=8)
tp_size = 2
deepspeed.tp_model_init(model, tp_size=tp_size, dtype=preferred_dtype())
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"partition_config": {
"use_default_specs": False,
"layer_specs": [{
"patterns": [".*\\.weight$"],
"partition_type": "skip",
}],
}
},
"zero_optimization": {
"stage": 0,
}
}
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
assert engine.autotp_size() == tp_size
assert groups.get_tensor_model_parallel_world_size() == tp_size
assert groups.get_data_parallel_world_size() == dist.get_world_size() // tp_size
def test_tp_model_init_tp_group_rejects_mpu(self):
skip_on_device()
reset_tp_model_init_state()
model = SimpleModel(hidden_dim=8)
tp_group = dist.new_group(ranks=[0])
deepspeed.tp_model_init(model, tp_size=1, dtype=preferred_dtype(), tp_group=tp_group)
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"zero_optimization": {
"stage": 0,
}
}
with pytest.raises(ValueError, match="tp_model_init provided tp_group"):
deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict, mpu=DummyMPU())
def test_tp_model_init_dtype_mismatch(self):
skip_on_device()
reset_tp_model_init_state()
model = SimpleModel(hidden_dim=8)
deepspeed.tp_model_init(model, tp_size=1, dtype=torch.float16)
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"bf16": {
"enabled": True,
},
"zero_optimization": {
"stage": 0,
}
}
with pytest.raises(ValueError, match="Conflicting dtype"):
deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict, mpu=DummyMPU())
@pytest.mark.sequential
@pytest.mark.parametrize("tp_size", [2, 4])
@pytest.mark.parametrize("tp_overlap_comm", [True, False])
def test_tp_model_init_row_parallel(self, tp_size: int, tp_overlap_comm: bool):
run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=False, use_tp_model_init=True)
@pytest.mark.sequential
@pytest.mark.parametrize("tp_size", [2, 4])
@pytest.mark.parametrize("tp_overlap_comm", [True, False])
def test_tp_model_init_column_parallel(self, tp_size: int, tp_overlap_comm: bool):
run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=True, use_tp_model_init=True)
@pytest.mark.parametrize("tp_size", [2, 4])
class TestTpDataloaderCorrectness(DistributedTest):
world_size = 4
reuse_dist_env = False
def test(self, tp_size: int):
skip_on_device()
hidden_dim = 128
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-6
}
},
"tensor_parallel": {
"autotp_size": tp_size,
"partition_config": {
"use_default_specs": False,
"layer_specs": [{
"patterns": [".*\\.weight$"],
"partition_type": "skip",
}],
}
},
"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 = SimpleModel(hidden_dim=hidden_dim)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
torch.manual_seed(42)
data_loader = random_dataloader(model=model,
total_samples=3,
hidden_dim=hidden_dim,
device=model.device,
dtype=preferred_dtype())
dist.barrier()
with should_assert_with_msg(
"Data inconsistency within the TP group. Please check the Dataloader implementation to ensure consistency."
):
for batch in data_loader:
# batch[0].requires_grad = requires_grad
batch[0] += dist.get_rank()
model(batch[0], batch[1])
model = SimpleModel(hidden_dim=hidden_dim)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
data_loader = random_dataloader(model=model,
total_samples=3,
hidden_dim=hidden_dim,
device=model.device,
dtype=preferred_dtype())
for batch in data_loader:
dist.broadcast(batch[0],
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group())
dist.broadcast(batch[1],
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group())
model(batch[0], batch[1])
def process_linear_layer(hidden_dim, input):
torch.manual_seed(42)
torch_linear = nn.Linear(hidden_dim,
hidden_dim,
dtype=preferred_dtype(),
device=get_accelerator().current_device())
torch_out = torch_linear(input)
torch_loss = torch_out.sum()
torch_loss.backward()
return torch_linear, torch_out
def run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel, use_tp_model_init=False):
skip_on_device()
hidden_dim = 128
batch_size_per_device = 1
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-6
}
},
"tensor_parallel": {
"autotp_size": tp_size,
"tp_overlap_comm": tp_overlap_comm
},
"zero_optimization": {
"stage": 0,
}
}
partition_type = "column" if column_parallel else "row"
config_dict["tensor_parallel"]["partition_config"] = {
"use_default_specs": False,
"layer_specs": [{
"patterns": [".*\\.weight$"],
"partition_type": partition_type,
}],
}
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=hidden_dim)
if use_tp_model_init:
reset_tp_model_init_state()
deepspeed.tp_model_init(model, tp_size=tp_size, dtype=preferred_dtype())
mpu = DummyMPU(tp_world_size=tp_size)
model, _, _, _ = deepspeed.initialize(model=model,
model_parameters=model.parameters(),
config=config_dict,
mpu=mpu)
else:
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
input = torch.randn(batch_size_per_device,
hidden_dim,
dtype=preferred_dtype(),
requires_grad=True,
device=get_accelerator().current_device())
dist.broadcast(input, groups.get_tensor_model_parallel_src_rank(), group=groups.get_tensor_model_parallel_group())
# Note: correctness checks below use standalone TP wrappers and do not
# rely on the model's AutoTP-partitioned parameters.
torch_linear, torch_out = process_linear_layer(hidden_dim, input)
if column_parallel:
linear = LinearLayer(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
out = linear(input.to(get_accelerator().current_device()))
loss = out.sum()
loss.backward()
cur_device_out = torch.chunk(torch_out, tp_size, dim=-1)[groups.get_tensor_model_parallel_rank()]
torch_grad = torch.chunk(torch_linear.weight.grad, tp_size, dim=0)[groups.get_tensor_model_parallel_rank()]
torch_bias_grad = torch.chunk(torch_linear.bias.grad, tp_size, dim=0)[groups.get_tensor_model_parallel_rank()]
torch.testing.assert_close(linear.bias.grad,
torch_bias_grad.to(get_accelerator().current_device()),
atol=1e-3,
rtol=1e-3)
torch.testing.assert_close(linear.weight.grad,
torch_grad.to(get_accelerator().current_device()),
atol=1e-3,
rtol=1e-3)
torch.testing.assert_close(cur_device_out.to(get_accelerator().current_device()).contiguous(),
out.contiguous(),
atol=1e-2,
rtol=1e-2)
else:
linear = LinearAllreduce(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
input_ = torch.chunk(input, tp_size, dim=-1)[groups.get_tensor_model_parallel_rank()]
out = linear(input_.to(get_accelerator().current_device()))
loss = out.sum()
loss.backward()
torch_grad = torch.chunk(torch_linear.weight.grad, tp_size, dim=1)[groups.get_tensor_model_parallel_rank()]
torch_bias_grad = torch_linear.bias.grad
torch.testing.assert_close(linear.bias.grad,
torch_bias_grad.to(get_accelerator().current_device()),
atol=1e-3,
rtol=1e-3)
torch.testing.assert_close(linear.weight.grad,
torch_grad.to(get_accelerator().current_device()),
atol=1e-3,
rtol=1e-3)
torch.testing.assert_close(out, torch_out.to(get_accelerator().current_device()), atol=1e-2, rtol=1e-2)
@pytest.mark.sequential
@pytest.mark.parametrize("tp_size", [2, 4])
@pytest.mark.parametrize("tp_overlap_comm", [True, False])
class TestTpLayerFwdBwd(DistributedTest):
world_size = 4
reuse_dist_env = False
def testRowParallel(self, tp_size: int, tp_overlap_comm: bool):
run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=False)
def testColumnParallel(self, tp_size: int, tp_overlap_comm: bool):
run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=True)
# @pytest.mark.sequential
class TestParamsGather(DistributedTest):
world_size = 4
reuse_dist_env = False
@pytest.mark.parametrize("layer_type", ["linear", "linearallreduce"])
def test(self, layer_type):
skip_on_device()
tp_size = 4
hidden_dim = 128
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-6
}
},
"tensor_parallel": {
"autotp_size": tp_size,
"partition_config": {
"use_default_specs": False,
"layer_specs": [{
"patterns": [".*\\.weight$"],
"partition_type": "skip",
}],
}
},
"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}
torch.manual_seed(42)
model = SequentialLinearModel(hidden_dim=hidden_dim)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
torch_linear = nn.Linear(hidden_dim, hidden_dim, dtype=preferred_dtype(), device="cpu")
total_params = sum(p.numel() for p in torch_linear.parameters())
tp_layer = None
if layer_type == "linear":
tp_layer = LinearLayer(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
elif layer_type == "linearallreduce":
tp_layer = LinearAllreduce(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
else:
raise ValueError(f"Invalid linear type: {config_dict['linear_type']}")
tp_params = sum(p.numel() for p in tp_layer.parameters())
expected_tp_params = 0
# compute expected TP params:
# - column-parallel (LinearLayer): weight & bias both split => total // tp_size
# - row-parallel (LinearAllreduce): weight split, bias (1d tensors) replicated
if layer_type == "linearallreduce":
weight_params = torch_linear.weight.numel()
bias_params = torch_linear.bias.numel()
expected_tp_params = weight_params // tp_size + bias_params
else:
expected_tp_params = total_params // tp_size
assert expected_tp_params == tp_params, (
f"{layer_type}: expected {expected_tp_params} tp params, got {tp_params}")
for name, param in tp_layer.named_parameters(recurse=False):
if is_model_parallel_parameter(param):
param.gather_params([param])
torch_linear = torch_linear.to(get_accelerator().current_device())
is_same_weights = all(
torch.equal(param1, param2) for param1, param2 in zip(tp_layer.parameters(), torch_linear.parameters()))
assert is_same_weights
params1 = sum(p.numel() for p in tp_layer.parameters())
assert total_params == params1
for name, param in tp_layer.named_parameters(recurse=False):
if is_model_parallel_parameter(param):
param._tp_partition([param])
tp_params2 = sum(p.numel() for p in tp_layer.parameters())
assert expected_tp_params == tp_params2
def dummy_init_engine(config):
# This is a dummy initialization function for the DeepSpeed engine.
# We only need to use the config to initialize the distributed settings for the test.
# Add default partition_config for simple test models if not provided
if "tensor_parallel" in config and "partition_config" not in config["tensor_parallel"]:
config["tensor_parallel"]["partition_config"] = {
"use_default_specs": False,
"layer_specs": [{
"patterns": [".*\\.weight$"],
"partition_type": "skip",
}],
}
model = SequentialLinearModel(hidden_dim=8)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
def prepare_tp_model(hidden_dim, nlayers, linear_indices, allreduce_indices, group, return_global_copy=False):
model = SequentialLinearModel(hidden_dim=hidden_dim, nlayers=nlayers).to(preferred_dtype())
base_model = None
if return_global_copy:
base_model = deepcopy(model)
for i in linear_indices:
layer = LinearLayer(model.linears[i], group)
model.linears[i] = layer
for i in allreduce_indices:
layer = LinearAllreduce(model.linears[i], group)
model.linears[i] = layer
return model, base_model
@pytest.mark.parametrize("zero_stage", [0, 1, 2])
@pytest.mark.parametrize("tp_size", [2, 4])
class TestSave(DistributedTest):
world_size = 4
reuse_dist_env = False
def test_save_original_weight(self, tp_size: int, zero_stage: int):
skip_on_device()
hidden_dim = 64
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-6
}
},
"tensor_parallel": {
"autotp_size": tp_size
},
"zero_optimization": {
"stage": zero_stage,
}
}
if preferred_dtype() is torch.float16:
config_dict["fp16"] = {"enabled": True}
elif preferred_dtype() is torch.bfloat16:
config_dict["bf16"] = {"enabled": True}
dummy_init_engine(config_dict)
torch.manual_seed(42)
model, base_model = prepare_tp_model(hidden_dim,
8, [2, 5], [3, 6],
groups.get_tensor_model_parallel_group(),
return_global_copy=True)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
cur_params_numel = sum(p.numel() for p in model.parameters())
base_params_numel = sum(p.numel() for p in base_model.parameters())
assert cur_params_numel < base_params_numel
tp_state_dict = model._consolidated_16bit_state_dict()
def compare_state_dicts(state_dict1, state_dict2):
if state_dict1.keys() != state_dict2.keys():
print("The state_dicts have different keys!")
return False
for key in state_dict1:
if not torch.allclose(state_dict1[key], state_dict2[key], atol=1e-3):
assert state_dict1[key].device == "cpu"
print(f"Parameters for {key} are different!")
return False
return True
base_state_dict = base_model.state_dict()
if dist.get_rank() == 0:
# we should consider the case when zero3 is used in the future.
assert compare_state_dicts(base_state_dict, tp_state_dict), "State_dict is not the same!"
else:
assert tp_state_dict is None, "noly rank0 should have the state_dict"
def test_ckpt_save(self, tmpdir, tp_size: int, zero_stage: int):
skip_on_device()
hidden_dim = 64
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-3
}
},
"zero_optimization": {
"stage": zero_stage,
},
"tensor_parallel": {
"autotp_size": tp_size
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 0.001,
"warmup_num_steps": 1000
}
}
}
if preferred_dtype() is torch.float16:
config_dict["fp16"] = {"enabled": True}
elif preferred_dtype() is torch.bfloat16:
config_dict["bf16"] = {"enabled": True}
dummy_init_engine(config_dict)
trained_model, _ = prepare_tp_model(hidden_dim, 8, [2, 5], [3, 6], groups.get_tensor_model_parallel_group())
loaded_model, _ = prepare_tp_model(hidden_dim, 8, [2, 5], [3, 6], groups.get_tensor_model_parallel_group())
trained_model, _, _, _ = deepspeed.initialize(model=trained_model,
model_parameters=trained_model.parameters(),
config=config_dict)
torch.manual_seed(42)
data_loader = random_dataloader(model=trained_model,
total_samples=3,
hidden_dim=hidden_dim,
device=trained_model.device,
dtype=preferred_dtype())
ckpt_path = os.path.join(tmpdir, 'tp_saved_checkpoint')
for i, batch in enumerate(data_loader):
batch[0].requires_grad = True
loss = trained_model(batch[0], batch[1])
loss = loss
trained_model.backward(loss)
trained_model.step()
trained_model.save_checkpoint(ckpt_path)
loaded_model, _, _, _ = deepspeed.initialize(model=loaded_model,
model_parameters=loaded_model.parameters(),
config=config_dict)
loaded_model.load_checkpoint(ckpt_path, load_optimizer_states=True, load_lr_scheduler_states=True)
compare_optimizer_states(trained_model, loaded_model, hidden_dim, fp16=(preferred_dtype() == torch.float16))
compare_lr_scheduler_states(trained_model, loaded_model)
@pytest.mark.parametrize("zero_stage", [0, 1, 2])
@pytest.mark.parametrize("tp_size", [2, 4])
class TestTpGradNorm(DistributedTest):
world_size = 4
reuse_dist_env = False
def test(self, tp_size: int, zero_stage: int):
skip_on_device()
hidden_dim = 64
config_dict = {
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-6
}
},
"tensor_parallel": {
"autotp_size": tp_size
},
"zero_optimization": {
"stage": zero_stage,
}
}
if preferred_dtype() is torch.float16:
config_dict["fp16"] = {"enabled": True}
elif preferred_dtype() is torch.bfloat16:
if zero_stage == 0:
pytest.skip(
"This test has an overflow data and needs to implement an overflow skip mechanism in BF16_optimizer"
)
config_dict["bf16"] = {"enabled": True}
torch.manual_seed(42)
dummy_init_engine(config=config_dict)
tp_model, base_model = prepare_tp_model(hidden_dim,
8, [2, 5], [3, 6],
groups.get_tensor_model_parallel_group(),
return_global_copy=True)
base_model, base_optimizer, _, _ = deepspeed.initialize(model=base_model,
model_parameters=base_model.parameters(),
config=config_dict)
data_loader = random_dataloader(model=base_model,
total_samples=20,
hidden_dim=hidden_dim,
device=base_model.device,
dtype=preferred_dtype())
for i, batch in enumerate(data_loader):
batch[0].requires_grad = True
loss = base_model(batch[0], batch[1])
loss = loss
base_model.backward(loss)
base_model.step()
base_norm = base_optimizer._global_grad_norm
base_model.destroy()
tp_model, tp_optimizer, _, _ = deepspeed.initialize(model=tp_model,
model_parameters=tp_model.parameters(),
config=config_dict)
for i, batch in enumerate(data_loader):
batch[0].requires_grad = True
loss = tp_model(batch[0], batch[1])
loss = loss
tp_model.backward(loss)
tp_model.step()
tp_norm = tp_optimizer._global_grad_norm
assert math.isclose(base_norm, tp_norm, abs_tol=1e-3), f"base_norm: {base_norm}, tp_norm: {tp_norm}"
tp_params_numel = sum(p.numel() for p in tp_model.parameters())
base_params_numel = sum(p.numel() for p in base_model.parameters())
assert tp_params_numel < base_params_numel, f"tp_params_numel: {tp_params_numel}, base_params_numel: {base_params_numel}"
tests/unit/model_parallelism/test_configurable_parallel_mp.py
+177
View File
@@ -0,0 +1,177 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import os
import torch
import deepspeed
import pytest
import random
import numpy as np
import deepspeed.comm as dist
from deepspeed.accelerator import get_accelerator
from unit.common import DistributedTest, DistributedFixture
from unit.megatron_model import get_gpt2_model, get_megatron_version
from deepspeed.utils.torch import required_torch_version
pytestmark = pytest.mark.skipif(not required_torch_version(min_version=1.5, max_version=1.13),
reason='Megatron-LM package requires Pytorch version >=1.5 and <=1.13')
# TODO: integrated testing of TP and ZeRO 1/2/3
def get_deepspeed_model(model):
ds_config_dict = {
"train_micro_batch_size_per_gpu": 1,
"optimizer": {
"type": "Lamb",
"params": {
"lr": 0.00015
}
},
}
from megatron import mpu
model, _, _, _ = deepspeed.initialize(model=model,
mpu=mpu,
model_parameters=model.parameters(),
config=ds_config_dict)
return model
class ConfigurableMP(DistributedTest):
@pytest.fixture(autouse=True)
def reset_random(self, seed=1234):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
get_accelerator().manual_seed_all(seed)
@pytest.fixture
def inputs(self, bs=1, seq_len=20):
input_ids = torch.randint(low=0, high=1000, size=(bs, seq_len))
position_ids = torch.randint(low=0, high=2, size=(bs, seq_len))
attention_mask = torch.randint(low=0, high=2, size=(bs, seq_len), dtype=torch.bool)
return [input_ids, position_ids, attention_mask]
class TestConfigurableMP(ConfigurableMP):
@pytest.mark.world_size(1)
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_gpt2_basic(self, tmpdir, inputs):
args_defaults = {
'num_layers': 2,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
model = get_gpt2_model(args_defaults)
model = get_deepspeed_model(model)
model.eval()
device_name = get_accelerator().device_name()
baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
tag = 'mp_1'
state_dict = {}
state_dict['checkpoint_version'] = get_megatron_version()
model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
dist.barrier()
model.load_checkpoint(tmpdir, tag=tag, load_optimizer_states=False, load_lr_scheduler_states=False)
test = model(inputs[0], inputs[1], inputs[2])
assert torch.allclose(baseline, test,
atol=1e-07), f"Baseline output {baseline} is not equal to save-then-load output {test}"
@pytest.mark.world_size(2)
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_gpt2_mp2_no_resize(self, tmpdir, inputs):
args_defaults = {
'num_layers': 2,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
model = get_gpt2_model(args_defaults, mp_size=2)
model = get_deepspeed_model(model)
model.eval()
device_name = get_accelerator().device_name()
baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
tag = 'mp_2'
state_dict = {}
state_dict['checkpoint_version'] = get_megatron_version()
model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
dist.barrier()
model.load_checkpoint(tmpdir, tag=tag, load_optimizer_states=False, load_lr_scheduler_states=False)
device_name = get_accelerator().device_name()
test = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
assert torch.allclose(baseline, test, rtol=1.0,
atol=1e-07), f"Baseline output {baseline} is not equal to save-then-load output {test}"
# This fixture provides the baseline model with mp=2 to TestConfigurableMPResize
class baseline_mp2(DistributedFixture):
world_size = 2
def run(self, inputs, class_tmpdir):
args_defaults = {
'num_layers': 2,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
model = get_gpt2_model(args_defaults, mp_size=self.world_size)
model = get_deepspeed_model(model)
model.eval()
with torch.no_grad():
device_name = get_accelerator().device_name()
baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
if dist.get_rank() == 0:
save_path = os.path.join(class_tmpdir, "output.pt")
torch.save(baseline.cpu(), save_path)
state_dict = {}
state_dict['checkpoint_version'] = get_megatron_version()
model.save_checkpoint(class_tmpdir, client_state=state_dict)
class TestConfigurableResizeMP(ConfigurableMP):
world_size = [1, 4]
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test(self, baseline_mp2, inputs, class_tmpdir):
args_defaults = {
'num_layers': 2,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
world_size = os.environ["WORLD_SIZE"]
model = get_gpt2_model(args_defaults, mp_size=world_size)
model = get_deepspeed_model(model)
model.eval()
with torch.no_grad():
model.load_checkpoint(class_tmpdir, load_optimizer_states=False, load_lr_scheduler_states=False)
device_name = get_accelerator().device_name()
test = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
if dist.get_rank() == 0:
load_path = os.path.join(class_tmpdir, "output.pt")
baseline = torch.load(load_path, weights_only=False)
test = test.cpu()
assert torch.allclose(
baseline, test,
atol=1e-03), f"Baseline output {baseline} is not equal to save-then-load output {test}"
tests/unit/model_parallelism/test_configurable_parallel_pp.py
+267
View File
@@ -0,0 +1,267 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import os
import torch
import deepspeed
import pytest
import random
import numpy as np
import deepspeed.comm as dist
from unit.common import DistributedTest, DistributedFixture
from unit.megatron_model import get_megatron_version
from unit.megatron_model import MockGPT2ModelPipe as GPT2ModelPipe
from deepspeed.utils import RepeatingLoader
from deepspeed.accelerator import get_accelerator
from deepspeed.utils.torch import required_torch_version
pytestmark = pytest.mark.skipif(not required_torch_version(min_version=1.5, max_version=1.13),
reason='Megatron-LM package requires Pytorch version >=1.5 and <=1.13')
def get_deepspeed_model(model):
ds_config_dict = {
"train_micro_batch_size_per_gpu": 1,
"optimizer": {
"type": "Lamb",
"params": {
"lr": 0.00015
}
},
}
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config_dict)
return model.to(get_accelerator().device_name())
def get_topology(mp, pp, world_size):
assert world_size % (pp * mp) == 0
dp = world_size // (pp * mp)
from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
topo = PipeModelDataParallelTopology(num_pp=pp, num_mp=mp, num_dp=dp)
return topo
class ConfigurablePP(DistributedTest):
@pytest.fixture(autouse=True)
def reset_random(self, seed=1234):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
get_accelerator().manual_seed_all(seed)
@pytest.fixture
def inputs(self, bs=1, seq_len=1, hidden_size=128):
hidden_states = torch.randn(bs, seq_len, hidden_size)
attention_mask = torch.randint(low=0, high=2, size=(bs, seq_len), dtype=torch.bool)
return (hidden_states, attention_mask)
class TestConfigurablePP(ConfigurablePP):
mp_size = 2
pp_size = 2
world_size = 4 # mp_size * pp_size
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_pp_basic(self, inputs, tmpdir):
# basic test case, mp_size=2, pp_size=2, verify ckpt saving/loading.
args_defaults = {
'num_layers': 8,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
mp_size = self.mp_size
pp_size = self.pp_size
world_size = self.world_size
topo = get_topology(mp_size, pp_size, world_size)
gpt2_pipe_model = GPT2ModelPipe(num_layers=8,
num_stages=pp_size,
mp_size=mp_size,
args_others=args_defaults,
topo=topo)
model = get_deepspeed_model(gpt2_pipe_model)
tag = 'pp_basic'
state_dict = {}
state_dict['checkpoint_version'] = get_megatron_version()
model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
if model.is_first_stage() or model.is_last_stage():
loader = RepeatingLoader([(inputs[0], 0)])
data_iter = iter(loader)
else:
data_iter = None
baseline = model.eval_batch(data_iter=data_iter, compute_loss=False, reduce_output=None)
dist.barrier()
model.load_checkpoint(tmpdir, tag=tag, load_optimizer_states=False, load_lr_scheduler_states=False)
dist.barrier()
test = model.eval_batch(data_iter=data_iter, compute_loss=False, reduce_output=None)
if test is not None:
assert len(baseline) == len(test)
# Compare outputs of each microbatch
for mb in range(len(baseline)):
for b, t in zip(baseline[mb], test[mb]):
if b.is_floating_point(): # don't compare masks
assert torch.allclose(
b, t,
atol=1e-07), f"Baseline output {baseline} is not equal to save-then-load output {test}"
# Fixture for defining the checkpoint path since all tests in
# TestConfigurableResizePP will use the same tmpdir
@pytest.fixture
def checkpoint_tag(mp_size, pp_size, mp_resize, pp_resize):
return f"{mp_size}-{pp_size}-{mp_resize}-{pp_resize}"
# Base class for creating / saving model output for baseline models. This is
# not meant to be used directly as a fixture to any classes
class _baseline(DistributedFixture):
world_size = None
def run(self, inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size):
assert int(os.environ["WORLD_SIZE"]) == (pp_size *
mp_size), "world size does not match provided pp_size and mp_size"
args_defaults = {
'num_layers': 8,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
topo = get_topology(mp_size, pp_size, mp_size * pp_size)
gpt2_pipe_model = GPT2ModelPipe(num_layers=8,
num_stages=pp_size,
mp_size=mp_size,
args_others=args_defaults,
topo=topo)
model = get_deepspeed_model(gpt2_pipe_model)
with torch.no_grad():
inputs = [x.to(get_accelerator().device_name()) for x in inputs]
if model.is_first_stage() or model.is_last_stage():
loader = RepeatingLoader([(inputs[0], 0)])
data_iter = iter(loader)
else:
data_iter = None
baseline = model.eval_batch(data_iter=data_iter, compute_loss=False, reduce_output=None)
if baseline is not None:
# baseline should be [[hidden, True]]]
assert len(baseline) == 1
assert len(baseline[0]) == 1
assert torch.is_tensor(baseline[0][0])
save_path = os.path.join(class_tmpdir, f"output-{checkpoint_tag}.pt")
torch.save(baseline[0][0].cpu(), save_path)
state_dict = {}
state_dict['checkpoint_version'] = get_megatron_version()
model.save_checkpoint(class_tmpdir, tag=checkpoint_tag, client_state=state_dict)
# This may look odd, but there is a limitation with DistributedFixture that
# doesn't allow us to reuse a fixture with different worldsizes. This could be
# implemented in conftest.py::pytest_fixture_setup and common.py::DistributedFixture
class baseline_ws1(_baseline):
world_size = 1
class baseline_ws2(_baseline):
world_size = 2
class baseline_ws4(_baseline):
world_size = 4
class TestConfigurableResizePP(ConfigurablePP):
def _test(self, inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size, mp_resize, pp_resize):
args_defaults = {
'num_layers': 8,
'hidden_size': 128,
'num_attention_heads': 8,
'max_position_embeddings': 128,
}
topo = get_topology(mp_resize, pp_resize, mp_resize * pp_resize)
gpt2_pipe_model = GPT2ModelPipe(num_layers=8,
num_stages=pp_resize,
mp_size=mp_resize,
args_others=args_defaults,
topo=topo)
model = get_deepspeed_model(gpt2_pipe_model)
with torch.no_grad():
model.load_checkpoint(class_tmpdir,
tag=checkpoint_tag,
load_optimizer_states=False,
load_lr_scheduler_states=False)
inputs = [x.to(get_accelerator().device_name()) for x in inputs]
if model.is_first_stage() or model.is_last_stage():
loader = RepeatingLoader([(inputs[0], 0)])
data_iter = iter(loader)
else:
data_iter = None
test = model.eval_batch(data_iter=data_iter, compute_loss=False, reduce_output=None)
if test is not None:
# test should be [[hidden, True]]]
assert len(test) == 1
assert len(test[0]) == 1
assert torch.is_tensor(test[0][0])
test = test[0][0].cpu()
load_path = os.path.join(class_tmpdir, f"output-{checkpoint_tag}.pt")
baseline = torch.load(load_path, weights_only=False)
assert torch.allclose(
baseline, test,
atol=1e-03), f"Baseline output {baseline} is not equal to save-then-load output {test}"
# These tests are divided by baseline model worldsize and test model worldsize
@pytest.mark.world_size(1)
@pytest.mark.parametrize("mp_size, pp_size, mp_resize, pp_resize", [(1, 2, 1, 1)])
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_world_size_2to1(self, inputs, class_tmpdir, checkpoint_tag, baseline_ws2, mp_size, pp_size, mp_resize,
pp_resize):
self._test(inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size, mp_resize, pp_resize)
@pytest.mark.world_size(1)
@pytest.mark.parametrize("mp_size, pp_size, mp_resize, pp_resize", [(2, 2, 1, 1)])
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_world_size_4to1(self, inputs, class_tmpdir, checkpoint_tag, baseline_ws4, mp_size, pp_size, mp_resize,
pp_resize):
self._test(inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size, mp_resize, pp_resize)
@pytest.mark.world_size(2)
@pytest.mark.parametrize("mp_size, pp_size, mp_resize, pp_resize", [(2, 2, 2, 1)])
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_world_size_4to2(self, inputs, class_tmpdir, checkpoint_tag, baseline_ws4, mp_size, pp_size, mp_resize,
pp_resize):
self._test(inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size, mp_resize, pp_resize)
@pytest.mark.world_size(4)
@pytest.mark.parametrize("mp_size, pp_size, mp_resize, pp_resize", [(1, 1, 2, 2)])
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_world_size_1to4(self, inputs, class_tmpdir, checkpoint_tag, baseline_ws1, mp_size, pp_size, mp_resize,
pp_resize):
self._test(inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size, mp_resize, pp_resize)
@pytest.mark.world_size(4)
@pytest.mark.parametrize("mp_size, pp_size, mp_resize, pp_resize", [(1, 2, 1, 4), (2, 1, 2, 2)])
@pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
def test_world_size_2to4(self, inputs, class_tmpdir, checkpoint_tag, baseline_ws2, mp_size, pp_size, mp_resize,
pp_resize):
self._test(inputs, class_tmpdir, checkpoint_tag, mp_size, pp_size, mp_resize, pp_resize)
tests/unit/model_parallelism/test_tp_plan_e2e.py
+274
View File
@@ -0,0 +1,274 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import torch
import deepspeed.comm as dist
import deepspeed
from deepspeed.accelerator import get_accelerator
from deepspeed.utils import groups
from deepspeed.runtime.utils import is_model_parallel_parameter
from unit.common import DistributedTest, preferred_dtype
def skip_on_device():
return
class TestTPPlanEndToEnd(DistributedTest):
world_size = 2
class SimpleHFModel(torch.nn.Module):
class Block(torch.nn.Module):
def __init__(self, hidden_size):
super().__init__()
self.q_proj = torch.nn.Linear(hidden_size, hidden_size * 2)
self.o_proj = torch.nn.Linear(hidden_size * 2, hidden_size)
def forward(self, x):
return self.o_proj(self.q_proj(x))
def __init__(self, hidden_size=64):
super().__init__()
self.hidden_size = hidden_size
self.config = type(
"Config",
(),
{"base_model_tp_plan": {
"*.q_proj": "colwise",
"*.o_proj": "rowwise",
}},
)()
self.layers = torch.nn.ModuleList([self.Block(hidden_size)])
def forward(self, x):
return self.layers[0](x)
def _setup_baseline_linears(self, model):
torch_q = torch.nn.Linear(model.hidden_size, model.hidden_size * 2)
torch_o = torch.nn.Linear(model.hidden_size * 2, model.hidden_size)
torch_q.load_state_dict(model.layers[0].q_proj.state_dict(), strict=True)
torch_o.load_state_dict(model.layers[0].o_proj.state_dict(), strict=True)
if preferred_dtype() == torch.float16:
torch_q = torch_q.half()
torch_o = torch_o.half()
elif preferred_dtype() == torch.bfloat16:
torch_q = torch_q.bfloat16()
torch_o = torch_o.bfloat16()
device = get_accelerator().current_device_name()
torch_q = torch_q.to(device)
torch_o = torch_o.to(device)
return torch_q, torch_o
def _compare_tp_gradients(self, model, torch_q, torch_o, input_tensor, engine):
def _get_grad(param):
if param.grad is not None:
return param.grad
return getattr(param, "grad_accum", None)
torch_q.zero_grad(set_to_none=True)
torch_o.zero_grad(set_to_none=True)
torch_q_out = torch_q(input_tensor)
torch_o_out = torch_o(torch_q_out)
torch_loss = torch_o_out.sum()
torch_loss.backward()
output = engine(input_tensor)
loss = output.sum()
engine.backward(loss)
tp_rank = groups.get_tensor_model_parallel_rank()
tp_size = engine.autotp_size()
q_proj = model.layers[0].q_proj
o_proj = model.layers[0].o_proj
torch_q_grad = torch.chunk(torch_q.weight.grad, tp_size, dim=0)[tp_rank]
torch_q_bias_grad = torch.chunk(torch_q.bias.grad, tp_size, dim=0)[tp_rank]
torch_o_grad = torch.chunk(torch_o.weight.grad, tp_size, dim=1)[tp_rank]
q_weight_grad = _get_grad(q_proj.weight)
q_bias_grad = _get_grad(q_proj.bias) if q_proj.bias is not None else None
o_weight_grad = _get_grad(o_proj.weight)
torch.testing.assert_close(q_weight_grad, torch_q_grad, atol=2e-2, rtol=2e-2)
if q_bias_grad is not None:
torch.testing.assert_close(q_bias_grad, torch_q_bias_grad, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(o_weight_grad, torch_o_grad, atol=2e-2, rtol=2e-2)
def _gather_and_compare_params(self, model, torch_q, torch_o, compare_values=True):
q_proj = model.layers[0].q_proj
o_proj = model.layers[0].o_proj
original_shards = []
for _, param in q_proj.named_parameters(recurse=False):
if is_model_parallel_parameter(param):
original_shards.append((param, param.data.detach().clone()))
for _, param in o_proj.named_parameters(recurse=False):
if is_model_parallel_parameter(param):
original_shards.append((param, param.data.detach().clone()))
for param, _ in original_shards:
param.gather_params([param])
if compare_values:
torch.testing.assert_close(q_proj.weight, torch_q.weight, atol=2e-2, rtol=2e-2)
if q_proj.bias is not None:
torch.testing.assert_close(q_proj.bias, torch_q.bias, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(o_proj.weight, torch_o.weight, atol=2e-2, rtol=2e-2)
if o_proj.bias is not None:
torch.testing.assert_close(o_proj.bias, torch_o.bias, atol=2e-2, rtol=2e-2)
for param, original in original_shards:
param._tp_partition([param])
torch.testing.assert_close(param.data, original, atol=2e-2, rtol=2e-2)
def test_tp_plan_basic_training(self):
skip_on_device()
model = self.SimpleHFModel()
if preferred_dtype() == torch.float16:
model = model.half()
elif preferred_dtype() == torch.bfloat16:
model = model.bfloat16()
torch_q, torch_o = self._setup_baseline_linears(model)
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": 2
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-4
}
},
"zero_optimization": {
"stage": 0
},
"steps_per_print": 1,
}
if preferred_dtype() == torch.float16:
ds_config["fp16"] = {"enabled": True}
elif preferred_dtype() == torch.bfloat16:
ds_config["bf16"] = {"enabled": True}
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
assert engine.autotp_size() == 2
input_tensor = torch.randn(2, 4, 64, dtype=preferred_dtype()).to(get_accelerator().current_device_name())
dist.broadcast(
input_tensor,
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group(),
)
if preferred_dtype() == torch.float16:
torch_q = torch_q.half()
torch_o = torch_o.half()
elif preferred_dtype() == torch.bfloat16:
torch_q = torch_q.bfloat16()
torch_o = torch_o.bfloat16()
self._compare_tp_gradients(model, torch_q, torch_o, input_tensor, engine)
def test_tp_plan_with_zero1(self):
skip_on_device()
model = self.SimpleHFModel()
torch_q, torch_o = self._setup_baseline_linears(model)
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": 2
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-4
}
},
"zero_optimization": {
"stage": 1
},
"steps_per_print": 1,
}
if preferred_dtype() == torch.float16:
ds_config["fp16"] = {"enabled": True}
elif preferred_dtype() == torch.bfloat16:
ds_config["bf16"] = {"enabled": True}
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
assert engine.autotp_size() == 2
for _ in range(1):
input_tensor = torch.randn(2, 4, 64, dtype=preferred_dtype()).to(get_accelerator().current_device_name())
dist.broadcast(
input_tensor,
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group(),
)
self._gather_and_compare_params(model, torch_q, torch_o, compare_values=False)
output = engine(input_tensor)
loss = output.mean()
engine.backward(loss)
engine.step()
for p in engine.parameters():
assert not torch.isnan(p).any()
def test_tp_plan_with_zero2(self):
skip_on_device()
model = self.SimpleHFModel()
torch_q, torch_o = self._setup_baseline_linears(model)
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": 2
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-4
}
},
"zero_optimization": {
"stage": 2
},
"steps_per_print": 1,
}
if preferred_dtype() == torch.float16:
ds_config["fp16"] = {"enabled": True}
elif preferred_dtype() == torch.bfloat16:
ds_config["bf16"] = {"enabled": True}
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
assert engine.autotp_size() == 2
input_tensor = torch.randn(2, 4, 64, dtype=preferred_dtype()).to(get_accelerator().current_device_name())
dist.broadcast(
input_tensor,
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group(),
)
self._gather_and_compare_params(model, torch_q, torch_o, compare_values=False)
output = engine(input_tensor)
loss = output.mean()
engine.backward(loss)
engine.step()
tests/unit/model_parallelism/test_tp_plan_real_models.py
+174
View File
@@ -0,0 +1,174 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import pytest
import torch
import deepspeed.comm as dist
import deepspeed
from deepspeed.accelerator import get_accelerator
from deepspeed.utils import groups
from unit.common import DistributedTest
def skip_on_device():
if get_accelerator().device_name() == "xpu":
pytest.skip("XPU requires a higher version for test")
class TestTPPlanRealHFModels(DistributedTest):
"""End-to-end tests using real HuggingFace models"""
world_size = 2
def test_qwen2_tp_plan_with_zero2(self):
"""Test Qwen2 model + tp_plan + ZeRO2"""
skip_on_device()
try:
from transformers import AutoModelForCausalLM, AutoConfig
except ImportError:
pytest.skip("transformers not installed")
# Create small Qwen2 config
config = AutoConfig.from_pretrained(
"Qwen/Qwen2-7B",
vocab_size=1000,
hidden_size=128,
intermediate_size=256,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=4,
)
model = AutoModelForCausalLM.from_config(config)
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": 2
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-4
}
},
"zero_optimization": {
"stage": 2
},
"bf16": {
"enabled": True
},
"steps_per_print": 1,
}
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
assert engine.autotp_size() == 2
# Train for a few steps
for _ in range(3):
input_ids = torch.randint(0, 1000, (1, 16)).to(get_accelerator().current_device_name())
dist.broadcast(
input_ids,
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group(),
)
outputs = engine(input_ids, labels=input_ids)
engine.backward(outputs.loss)
engine.step()
assert not torch.isnan(outputs.loss)
def test_custom_model_with_custom_tp_plan(self):
"""Test custom model + custom tp_plan"""
skip_on_device()
class CustomTransformerModel(torch.nn.Module):
def __init__(self, hidden_size=64):
super().__init__()
self.config = type(
"Config",
(),
{
"base_model_tp_plan": {
"encoder.*.attention.query": "colwise",
"encoder.*.attention.key": "colwise",
"encoder.*.attention.value": "colwise",
"encoder.*.attention.output": "rowwise",
"encoder.*.ffn.intermediate": "colwise",
"encoder.*.ffn.output": "rowwise",
}
},
)()
# Simple encoder layers
self.encoder = torch.nn.ModuleList([
torch.nn.ModuleDict({
"attention":
torch.nn.ModuleDict({
"query": torch.nn.Linear(hidden_size, hidden_size),
"key": torch.nn.Linear(hidden_size, hidden_size),
"value": torch.nn.Linear(hidden_size, hidden_size),
"output": torch.nn.Linear(hidden_size, hidden_size),
}),
"ffn":
torch.nn.ModuleDict({
"intermediate": torch.nn.Linear(hidden_size, hidden_size * 4),
"output": torch.nn.Linear(hidden_size * 4, hidden_size),
}),
}) for _ in range(2)
])
def forward(self, x):
for layer in self.encoder:
# Simplified attention
q = layer.attention.query(x)
k = layer.attention.key(x)
v = layer.attention.value(x)
attn_out = layer.attention.output(q + k + v)
# FFN
intermediate = torch.relu(layer.ffn.intermediate(attn_out))
x = layer.ffn.output(intermediate)
return x
model = CustomTransformerModel(hidden_size=64)
ds_config = {
"train_micro_batch_size_per_gpu": 1,
"tensor_parallel": {
"autotp_size": 2
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-4
}
},
"zero_optimization": {
"stage": 0
},
"bf16": {
"enabled": True
},
}
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
assert engine.autotp_size() == 2
# Training step
input_tensor = torch.randn(2, 4, 64, dtype=torch.bfloat16).to(get_accelerator().current_device_name())
dist.broadcast(
input_tensor,
src=groups.get_tensor_model_parallel_src_rank(),
group=groups.get_tensor_model_parallel_group(),
)
output = engine(input_tensor)
loss = output.mean()
engine.backward(loss)
engine.step()