# 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)