wehub-resource-sync
134 lines
4.4 KiB
Python
134 lines
4.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import deepspeed
|
|
import deepspeed.comm as dist
|
|
import torch
|
|
from unit.common import DistributedTest
|
|
from unit.simple_model import SimpleModel, random_dataloader
|
|
|
|
|
|
def create_model(config_dict):
|
|
hidden_dim = 64
|
|
model = SimpleModel(hidden_dim)
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
return model
|
|
|
|
|
|
def train_shared_loss(num_models, config_dict, dtype):
|
|
hidden_dim = 64
|
|
|
|
models = [create_model(config_dict) for _ in range(num_models)]
|
|
data_loader = random_dataloader(model=models[0],
|
|
total_samples=4,
|
|
hidden_dim=hidden_dim,
|
|
device=models[0].device,
|
|
dtype=dtype)
|
|
dist.barrier()
|
|
for _, batch in enumerate(data_loader):
|
|
losses = [m.module(batch[0], batch[1]) for m in models]
|
|
loss = sum(l / (i + 1) for i, l in enumerate(losses))
|
|
loss.backward()
|
|
|
|
for m in models:
|
|
m._backward_epilogue()
|
|
|
|
for m in models:
|
|
m.step()
|
|
|
|
for m in models:
|
|
m.optimizer.zero_grad()
|
|
|
|
for m in models:
|
|
m.destroy()
|
|
|
|
|
|
def train_independent_loss(num_models, config_dict, dtype):
|
|
hidden_dim = 64
|
|
|
|
models = [create_model(config_dict) for _ in range(num_models)]
|
|
data_loader = random_dataloader(model=models[0],
|
|
total_samples=4,
|
|
hidden_dim=hidden_dim,
|
|
device=models[0].device,
|
|
dtype=dtype)
|
|
dist.barrier()
|
|
for _, batch in enumerate(data_loader):
|
|
losses = [m.module(batch[0], batch[1]) for m in models]
|
|
for m, loss in zip(models, losses):
|
|
m.backward(loss)
|
|
m.step()
|
|
|
|
for m in models:
|
|
m.destroy()
|
|
|
|
|
|
@pytest.mark.parametrize('num_models', [1, 2, 3])
|
|
class TestMultipleModels(DistributedTest):
|
|
world_size = 2
|
|
reuse_dist_env = True
|
|
|
|
@pytest.mark.parametrize('shared_loss', [False, True])
|
|
@pytest.mark.parametrize('zero_stage', [1, 2, 3])
|
|
@pytest.mark.parametrize('fp32_grad_accum', [False, True])
|
|
@pytest.mark.parametrize('contiguous_gradients', [False, True])
|
|
@pytest.mark.parametrize('overlap_comm', [False, True])
|
|
def test_zero_optimizer(self, num_models, shared_loss, zero_stage, fp32_grad_accum, contiguous_gradients,
|
|
overlap_comm):
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-4
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"contiguous_gradients": contiguous_gradients,
|
|
"overlap_comm": overlap_comm,
|
|
},
|
|
"fp16": {
|
|
"initial_scale_power": 8,
|
|
"enabled": True
|
|
},
|
|
}
|
|
if fp32_grad_accum:
|
|
config_dict["data_types"] = {"grad_accum_dtype": "fp32"}
|
|
|
|
if shared_loss:
|
|
train_shared_loss(num_models=num_models, config_dict=config_dict, dtype=torch.float16)
|
|
else:
|
|
train_independent_loss(num_models=num_models, config_dict=config_dict, dtype=torch.float16)
|
|
|
|
# TODO: Combination of shared_loss==True and bf16.immediate_grad_update==False is currently broken
|
|
@pytest.mark.parametrize('shared_loss', [False, True])
|
|
def test_bf16_optimizer(self, num_models, shared_loss):
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-4
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 1,
|
|
},
|
|
"bf16": {
|
|
"enabled": True,
|
|
"immediate_grad_update": True,
|
|
},
|
|
"data_types": {
|
|
"grad_accum_dtype": "fp32"
|
|
}
|
|
}
|
|
|
|
if shared_loss:
|
|
train_shared_loss(num_models=num_models, config_dict=config_dict, dtype=torch.bfloat16)
|
|
else:
|
|
train_independent_loss(num_models=num_models, config_dict=config_dict, dtype=torch.bfloat16)
|