wehub-resource-sync
275 lines
9.5 KiB
Python
275 lines
9.5 KiB
Python
# 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()
|