wehub-resource-sync
494 lines
20 KiB
Python
494 lines
20 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
from typing import Callable
|
|
import torch
|
|
from torch.optim import Optimizer, Adam, AdamW
|
|
from torch.optim.lr_scheduler import _LRScheduler, LambdaLR
|
|
|
|
from unit.simple_model import SimpleModel, random_dataloader
|
|
from unit.common import DistributedTest
|
|
from unit.util import bf16_required_version_check, required_amp_check
|
|
|
|
import deepspeed
|
|
from deepspeed.ops.adam import FusedAdam
|
|
from deepspeed.runtime.lr_schedules import WARMUP_LR, WarmupLR
|
|
from deepspeed.runtime.config import ADAM_OPTIMIZER
|
|
from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
|
|
from deepspeed.runtime.utils import see_memory_usage
|
|
from deepspeed.utils.torch import required_torch_version
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.ops.op_builder import FusedAdamBuilder
|
|
|
|
|
|
# Ensure client multiprocessing is not broken by deepspeed import
|
|
@pytest.mark.parametrize('method', ['spawn', 'fork', 'forkserver'])
|
|
def test_start_method_safety(method):
|
|
import torch.multiprocessing as mp
|
|
mp.set_start_method(method, force=True)
|
|
|
|
|
|
@pytest.mark.parametrize('zero_stage', [0, 3])
|
|
class TestNoOptim(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, zero_stage):
|
|
if zero_stage == 3 and not required_torch_version(min_version=1.8):
|
|
pytest.skip("zero-3 param offload requires at least torch 1.8")
|
|
|
|
ds_config = {
|
|
'train_batch_size': self.world_size,
|
|
'zero_optimization': {
|
|
"stage": zero_stage,
|
|
"offload_param": {
|
|
"device": "cpu"
|
|
}
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
ds_config["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
ds_config["fp16"] = {"enabled": True}
|
|
# 20B test
|
|
#hidden_dim = 16 * 1024
|
|
hidden_dim = 4
|
|
|
|
with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=ds_config):
|
|
model = SimpleModel(hidden_dim, nlayers=78)
|
|
see_memory_usage('pre-init', force=True)
|
|
model, _, _, _ = deepspeed.initialize(model=model, config=ds_config)
|
|
see_memory_usage('post-init', force=True)
|
|
data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device)
|
|
for batch in data_loader:
|
|
model(batch[0], batch[1])
|
|
see_memory_usage('post-fwds', force=True)
|
|
|
|
|
|
@pytest.mark.parametrize('optimizer_type', [None, Optimizer, Callable])
|
|
class TestClientOptimizer(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, optimizer_type):
|
|
|
|
def _optimizer_callable(params) -> Optimizer:
|
|
return AdamW(params=params)
|
|
|
|
if (optimizer_type is None) and (not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME]):
|
|
pytest.skip("FusedAdam is not compatible")
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
config_dict = {'train_batch_size': 1}
|
|
if optimizer_type is None:
|
|
client_optimizer = None
|
|
config_dict['optimizer'] = {'type': ADAM_OPTIMIZER}
|
|
elif optimizer_type is Optimizer:
|
|
client_optimizer = Adam(model.parameters())
|
|
else:
|
|
client_optimizer = _optimizer_callable
|
|
|
|
_, ds_optimizer, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer)
|
|
if client_optimizer is None:
|
|
assert isinstance(ds_optimizer, FusedAdam)
|
|
elif isinstance(client_optimizer, Optimizer):
|
|
assert ds_optimizer == client_optimizer
|
|
else:
|
|
assert isinstance(ds_optimizer, AdamW)
|
|
|
|
|
|
@pytest.mark.parametrize('client_parameters', [True, False])
|
|
class TestConfigOptimizer(DistributedTest):
|
|
world_size = 1
|
|
|
|
@pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME],
|
|
reason="FusedAdam is not compatible")
|
|
def test(self, client_parameters):
|
|
ds_config = {"train_batch_size": 1, "optimizer": {"type": "Adam", "params": {"lr": 0.001}}}
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
if client_parameters:
|
|
model_parameters = list(model.parameters())
|
|
else:
|
|
model_parameters = None
|
|
|
|
_, ds_optimizer, _, _ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model_parameters)
|
|
|
|
assert isinstance(ds_optimizer, FusedAdam)
|
|
|
|
|
|
@pytest.mark.parametrize('optimizer_extension', ['zero1', 'zero2', 'zero3', 'amp', None])
|
|
@pytest.mark.parametrize('model_dtype', ['fp16', 'bf16', 'fp32'])
|
|
@pytest.mark.parametrize('grad_accum_dtype', [None, 'fp16', 'bf16', 'fp32'])
|
|
class TestOptimizerImplementation(DistributedTest):
|
|
world_size = 1
|
|
reuse_dist_env = True
|
|
|
|
def test(self, optimizer_extension, model_dtype, grad_accum_dtype):
|
|
if not get_accelerator().is_fp16_supported():
|
|
if model_dtype == 'fp16' or grad_accum_dtype == 'fp16':
|
|
pytest.skip("fp16 is not supported")
|
|
if optimizer_extension == 'zero1':
|
|
zero_stage = 1
|
|
elif optimizer_extension == 'zero2':
|
|
zero_stage = 2
|
|
elif optimizer_extension == 'zero3':
|
|
zero_stage = 3
|
|
else:
|
|
zero_stage = 0
|
|
amp = (optimizer_extension == 'amp')
|
|
fp16 = (model_dtype == 'fp16')
|
|
bf16 = (model_dtype == 'bf16')
|
|
# Skip checks
|
|
if bf16 and not bf16_required_version_check():
|
|
pytest.skip(
|
|
"DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
|
|
)
|
|
if amp and not required_amp_check():
|
|
pytest.skip("Amp is not installed can't run amp check")
|
|
# Config declaration
|
|
ds_config = {
|
|
"train_batch_size": 1,
|
|
'fp16': {
|
|
'enabled': fp16
|
|
},
|
|
'bf16': {
|
|
'enabled': bf16
|
|
},
|
|
'amp': {
|
|
'enabled': amp
|
|
},
|
|
'zero_optimization': {
|
|
"stage": zero_stage
|
|
},
|
|
"data_types": {
|
|
"grad_accum_dtype": grad_accum_dtype
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.001
|
|
}
|
|
}
|
|
}
|
|
|
|
key = (optimizer_extension, model_dtype, grad_accum_dtype)
|
|
|
|
# Enumerate supported configurations
|
|
is_supported = {}
|
|
# ZeRO 1 Wrapper
|
|
is_supported[('zero1', 'fp16', None)] = True
|
|
is_supported[('zero1', 'fp16', 'fp16')] = True
|
|
is_supported[('zero1', 'fp16', 'bf16')] = True
|
|
is_supported[('zero1', 'fp16', 'fp32')] = True
|
|
is_supported[('zero1', 'bf16', None)] = True
|
|
is_supported[('zero1', 'bf16', 'fp16')] = True
|
|
is_supported[('zero1', 'bf16', 'bf16')] = True
|
|
is_supported[('zero1', 'bf16', 'fp32')] = True
|
|
is_supported[('zero1', 'fp32', None)] = True
|
|
is_supported[('zero1', 'fp32', 'fp16')] = True
|
|
is_supported[('zero1', 'fp32', 'bf16')] = True
|
|
is_supported[('zero1', 'fp32', 'fp32')] = True
|
|
# ZeRO 2 Wrapper
|
|
is_supported[('zero2', 'fp16', None)] = True
|
|
is_supported[('zero2', 'fp16', 'fp16')] = True
|
|
is_supported[('zero2', 'fp16', 'bf16')] = True
|
|
is_supported[('zero2', 'fp16', 'fp32')] = True
|
|
is_supported[('zero2', 'bf16', None)] = True
|
|
is_supported[('zero2', 'bf16', 'fp16')] = True
|
|
is_supported[('zero2', 'bf16', 'bf16')] = True
|
|
is_supported[('zero2', 'bf16', 'fp32')] = True
|
|
is_supported[('zero2', 'fp32', None)] = True
|
|
is_supported[('zero2', 'fp32', 'fp16')] = True
|
|
is_supported[('zero2', 'fp32', 'bf16')] = True
|
|
is_supported[('zero2', 'fp32', 'fp32')] = True
|
|
# ZeRO 3 Wrapper
|
|
is_supported[('zero3', 'fp16', None)] = True
|
|
is_supported[('zero3', 'fp16', 'fp16')] = True
|
|
is_supported[('zero3', 'fp16', 'bf16')] = True
|
|
is_supported[('zero3', 'fp16', 'fp32')] = True
|
|
is_supported[('zero3', 'bf16', None)] = True
|
|
is_supported[('zero3', 'bf16', 'fp16')] = True
|
|
is_supported[('zero3', 'bf16', 'bf16')] = True
|
|
is_supported[('zero3', 'bf16', 'fp32')] = True
|
|
is_supported[('zero3', 'fp32', None)] = True
|
|
is_supported[('zero3', 'fp32', 'fp16')] = True
|
|
is_supported[('zero3', 'fp32', 'bf16')] = True
|
|
is_supported[('zero3', 'fp32', 'fp32')] = True
|
|
# Amp Wrapper
|
|
is_supported[('amp', 'fp32', None)] = True
|
|
is_supported[('amp', 'fp32', 'fp32')] = True
|
|
# FP16 Wrapper
|
|
is_supported[(None, 'fp16', None)] = True
|
|
is_supported[(None, 'fp16', 'fp16')] = True
|
|
# BF16 Wrapper
|
|
is_supported[(None, 'bf16', 'bf16')] = True
|
|
is_supported[(None, 'bf16', None)] = True
|
|
# No Wrapper
|
|
is_supported[(None, 'fp32', None)] = True
|
|
is_supported[(None, 'fp32', 'fp32')] = True
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
model_parameters = list(model.parameters())
|
|
|
|
if key in is_supported:
|
|
_, ds_optimizer, _, _ = deepspeed.initialize(config=ds_config,
|
|
model=model,
|
|
model_parameters=model_parameters)
|
|
assert True
|
|
else:
|
|
with pytest.raises(NotImplementedError):
|
|
_, ds_optimizer, _, _ = deepspeed.initialize(config=ds_config,
|
|
model=model,
|
|
model_parameters=model_parameters)
|
|
|
|
|
|
class TestBf16ZeRO0UnfusedOptimizer(DistributedTest):
|
|
world_size = 1
|
|
reuse_dist_env = True
|
|
|
|
def test_static_scale_and_zero_grad_after_step(self):
|
|
if not bf16_required_version_check():
|
|
pytest.skip(
|
|
"DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
|
|
)
|
|
|
|
hidden_dim = 16
|
|
model = SimpleModel(hidden_dim)
|
|
client_optimizer = AdamW(model.parameters(), lr=1e-4)
|
|
ds_config = {
|
|
"train_batch_size": 1,
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0
|
|
},
|
|
}
|
|
|
|
engine, _, _, _ = deepspeed.initialize(config=ds_config,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer)
|
|
|
|
assert isinstance(engine.optimizer, FP16_UnfusedOptimizer)
|
|
assert engine.optimizer.low_precision_dtype == torch.bfloat16
|
|
assert engine.optimizer.loss_scale_config.dynamic_loss_scale is False
|
|
assert engine.optimizer.loss_scale_config.cur_scale == 1
|
|
|
|
data_loader = random_dataloader(model=engine,
|
|
total_samples=1,
|
|
hidden_dim=hidden_dim,
|
|
device=engine.device,
|
|
dtype=torch.bfloat16)
|
|
batch = next(iter(data_loader))
|
|
|
|
loss = engine(batch[0], batch[1])
|
|
engine.backward(loss)
|
|
assert any(param.grad is not None for param in engine.module.parameters() if param.requires_grad)
|
|
|
|
engine.step()
|
|
assert all(param.grad is None for param in engine.module.parameters() if param.requires_grad)
|
|
|
|
|
|
@pytest.mark.parametrize("scheduler_type", [None, _LRScheduler, Callable])
|
|
@pytest.mark.parametrize("optimizer_type", [None, Optimizer, Callable])
|
|
class TestClientLrScheduler(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, scheduler_type, optimizer_type):
|
|
|
|
def _my_lambda(epoch):
|
|
return epoch // 10
|
|
|
|
def _optimizer_callable(params) -> Optimizer:
|
|
return torch.optim.AdamW(params=params)
|
|
|
|
def _lr_scheduler_callable(optimizer) -> _LRScheduler:
|
|
return LambdaLR(optimizer, _my_lambda)
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
config_dict = {'train_batch_size': 1}
|
|
|
|
client_optimizer = None
|
|
client_scheduler = None
|
|
|
|
if optimizer_type is None:
|
|
config_dict['optimizer'] = {'type': ADAM_OPTIMIZER}
|
|
elif optimizer_type is Optimizer:
|
|
client_optimizer = torch.optim.Adam(model.parameters())
|
|
else:
|
|
client_optimizer = _optimizer_callable
|
|
|
|
if scheduler_type is None:
|
|
config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
|
|
elif scheduler_type == _LRScheduler:
|
|
if isinstance(client_optimizer, Optimizer):
|
|
client_scheduler = LambdaLR(client_optimizer, _my_lambda)
|
|
else:
|
|
# Verify invalid combination is correctly handled
|
|
client_scheduler = LambdaLR(torch.optim.Adam(model.parameters()), _my_lambda)
|
|
else:
|
|
client_scheduler = _lr_scheduler_callable
|
|
|
|
if isinstance(client_scheduler, _LRScheduler) and not isinstance(client_optimizer, Optimizer):
|
|
with pytest.raises(AssertionError):
|
|
_, _, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer,
|
|
lr_scheduler=client_scheduler)
|
|
else:
|
|
_, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer,
|
|
lr_scheduler=client_scheduler)
|
|
if client_scheduler is None:
|
|
assert isinstance(ds_lr_scheduler, WarmupLR)
|
|
elif isinstance(client_scheduler, _LRScheduler):
|
|
assert ds_lr_scheduler == client_scheduler
|
|
else:
|
|
assert isinstance(ds_lr_scheduler, LambdaLR)
|
|
|
|
|
|
@pytest.mark.parametrize("scheduler_type", [None, _LRScheduler, Callable])
|
|
class TestClientLrSchedulerInit(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test_same_lrscheler_and_callable(self, scheduler_type):
|
|
"""
|
|
Expect behavior
|
|
|
|
if lr scheduler is defined in code and passed into initialize as arg,
|
|
it will be used even this is a lr scheduler has been defined in config.
|
|
|
|
Initialize lr scheduler from config when no lr scheduler is defined in code.
|
|
"""
|
|
|
|
def _my_lambda(epoch):
|
|
return epoch // 10
|
|
|
|
def _lr_scheduler_callable(optimizer) -> _LRScheduler:
|
|
return LambdaLR(optimizer, _my_lambda)
|
|
|
|
config_dict = {'train_batch_size': 1}
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
|
|
|
|
if scheduler_type is None:
|
|
config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
|
|
client_scheduler = None
|
|
elif scheduler_type == _LRScheduler:
|
|
client_scheduler = LambdaLR(client_optimizer, _my_lambda)
|
|
else:
|
|
client_scheduler = _lr_scheduler_callable
|
|
|
|
_, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer,
|
|
lr_scheduler=client_scheduler)
|
|
if scheduler_type is None:
|
|
# in this case, we initialize from config
|
|
assert not isinstance(ds_lr_scheduler, LambdaLR)
|
|
assert isinstance(ds_lr_scheduler, WarmupLR)
|
|
else:
|
|
# in this case, we initialize from passed-in scheduler
|
|
assert isinstance(ds_lr_scheduler, LambdaLR)
|
|
assert not isinstance(ds_lr_scheduler, WarmupLR)
|
|
|
|
def test_diff_lrscheler_and_callable(self, scheduler_type):
|
|
"""
|
|
In this test,
|
|
the LambdaLR will be used for lrscheduler type
|
|
and the StepLR will be used for callable type
|
|
"""
|
|
|
|
from torch.optim.lr_scheduler import StepLR
|
|
|
|
def _my_lambda(epoch):
|
|
return epoch // 10
|
|
|
|
def _lr_scheduler_callable(optimizer) -> _LRScheduler:
|
|
return StepLR(optimizer, step_size=30)
|
|
|
|
config_dict = {'train_batch_size': 1}
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
|
|
|
|
if scheduler_type is None:
|
|
config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
|
|
client_scheduler = None
|
|
elif scheduler_type == _LRScheduler:
|
|
client_scheduler = LambdaLR(client_optimizer, _my_lambda)
|
|
else:
|
|
client_scheduler = _lr_scheduler_callable
|
|
|
|
_, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer,
|
|
lr_scheduler=client_scheduler)
|
|
if scheduler_type is None:
|
|
assert isinstance(ds_lr_scheduler, WarmupLR)
|
|
elif scheduler_type == _LRScheduler:
|
|
assert isinstance(ds_lr_scheduler, LambdaLR)
|
|
else:
|
|
# callable
|
|
assert isinstance(ds_lr_scheduler, StepLR)
|
|
|
|
def test_diff_lrscheler_and_callable_onecyclelr_steplr(self, scheduler_type):
|
|
|
|
from deepspeed.runtime.lr_schedules import OneCycle, ONE_CYCLE, CYCLE_MIN_LR, CYCLE_MAX_LR
|
|
from torch.optim.lr_scheduler import OneCycleLR, StepLR
|
|
|
|
def _lr_scheduler_callable(optimizer) -> _LRScheduler:
|
|
return OneCycleLR(optimizer, max_lr=0.01, total_steps=200)
|
|
|
|
config_dict = {'train_batch_size': 1}
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
|
|
|
|
if scheduler_type is None:
|
|
config_dict['scheduler'] = {'type': ONE_CYCLE, 'params': {CYCLE_MIN_LR: 0, CYCLE_MAX_LR: 0.1}}
|
|
client_scheduler = None
|
|
elif scheduler_type == _LRScheduler:
|
|
client_scheduler = StepLR(client_optimizer, step_size=30)
|
|
else:
|
|
client_scheduler = _lr_scheduler_callable
|
|
|
|
_, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=list(model.parameters()),
|
|
optimizer=client_optimizer,
|
|
lr_scheduler=client_scheduler)
|
|
if scheduler_type is None:
|
|
assert isinstance(ds_lr_scheduler, OneCycle)
|
|
elif scheduler_type == _LRScheduler:
|
|
assert isinstance(ds_lr_scheduler, StepLR)
|
|
else:
|
|
# callable
|
|
assert isinstance(ds_lr_scheduler, OneCycleLR)
|