# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team import torch import deepspeed import pytest from deepspeed.ops.adam import FusedAdam from unit.common import DistributedTest from deepspeed.ops.op_builder import CPUAdamBuilder from unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader from unit.util import bf16_required_version_check from deepspeed import comm as dist from deepspeed.accelerator import get_accelerator from unit.v1.zero.test_zero_user_backward import (initialize_distributed, create_ddp_model, collect_ddp_gradients, collect_gradients_safe, compare_gradients) class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest): world_size = 1 def test(self, zero_stage=2, use_cpu_offload=False): 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" ) if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: pytest.skip("cpu-adam is not compatible") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "scheduler": { "type": "OneCycle", "params": { "cycle_first_step_size": 16000, "cycle_first_stair_count": 8000, "decay_step_size": 16000, "cycle_min_lr": 1e-06, "cycle_max_lr": 3e-05, "decay_lr_rate": 1e-07, "cycle_min_mom": 0.85, "cycle_max_mom": 0.99, "decay_mom_rate": 0.0 } }, "fp16": { "enabled": False }, "bf16": { "enabled": True }, "zero_optimization": { "stage": zero_stage, "cpu_offload": use_cpu_offload } } hidden_dim = 10 model = SimpleModel(hidden_dim) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device, dtype=torch.bfloat16) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() class TestZeroAllowUntestedOptimizer(DistributedTest): world_size = 1 def test(self, zero_stage=2, use_cpu_offload=False): 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" ) if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: pytest.skip("cpu-adam is not compatible") config_dict = { "train_micro_batch_size_per_gpu": 4, "steps_per_print": 1, "fp16": { "enabled": False, }, "bf16": { "enabled": True }, "zero_optimization": { "stage": zero_stage, "cpu_offload": use_cpu_offload }, "zero_allow_untested_optimizer": False } hidden_dim = 10 model = SimpleModel(hidden_dim) optimizer = SimpleOptimizer(model.parameters()) with pytest.raises(AssertionError): model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer, model_parameters=model.parameters()) class TestZeroEmptyPartition(DistributedTest): world_size = 3 def test(self, zero_stage=2, use_cpu_offload=False): 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" ) if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: pytest.skip("cpu-adam is not compatible") if zero_stage == 3: pytest.skip("skip for now") config_dict = { "train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 1, "fp16": { "enabled": False }, "bf16": { "enabled": True }, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "zero_optimization": { "stage": zero_stage, "cpu_offload": use_cpu_offload, "reduce_bucket_size": 100, "allgather_bucket_size": 100 } } hidden_dim = 1 model = SimpleModel(hidden_dim) # Ensure model has 2 parameters, to cause empty partition with DP=3 assert len(list(model.parameters())) == 2 model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) # Now make sure things work.. data_loader = random_dataloader(model=model, total_samples=1, hidden_dim=hidden_dim, device=model.device, dtype=torch.bfloat16) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() @pytest.mark.parametrize("optimizer_constructor", [torch.optim.Adam, FusedAdam]) class TestZeroSupportedClientOptimizer(DistributedTest): world_size = 1 def test(self, optimizer_constructor, zero_stage=2): 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" ) config_dict = { "train_micro_batch_size_per_gpu": 2, "steps_per_print": 1, "fp16": { "enabled": False }, "bf16": { "enabled": True }, "zero_optimization": { "stage": zero_stage } } hidden_dim = 10 model = SimpleModel(hidden_dim) client_optimizer = optimizer_constructor(params=model.parameters()) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=client_optimizer) class TestZero2ReduceScatterOff(DistributedTest): world_size = 2 def test(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" ) config_dict = { "train_micro_batch_size_per_gpu": 2, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "gradient_clipping": 1.0, "zero_optimization": { "stage": 2, "contiguous_gradients": True, "allgather_bucket_size": 2000000000, "reduce_bucket_size": 200000000, "overlap_comm": False, "reduce_scatter": False }, "fp16": { "enabled": False }, "bf16": { "enabled": True } } hidden_dim = 10 model = SimpleModel(hidden_dim) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device, dtype=torch.bfloat16) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() class TestZeroEmptyGrad(DistributedTest): world_size = 1 def test(self, stage=2): 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" ) config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "fp16": { "enabled": False }, "bf16": { "enabled": True }, "zero_optimization": { "stage": stage } } hidden_dim = 10 model = SimpleModel(hidden_dim) optimizer = torch.optim.Adam(model.parameters()) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer) data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device, dtype=torch.bfloat16) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() @pytest.mark.parametrize("comp_type", [torch.float16, torch.bfloat16, torch.float], ids=["fp16", "bf16", "fp32"]) @pytest.mark.parametrize("comm_type", [torch.float16, torch.bfloat16, None], ids=["fp16", "bf16", "default"]) class TestZeroDtypeCocktail(DistributedTest): world_size = 2 def test(self, comp_type, comm_type): if comp_type == torch.bfloat16 or comm_type == torch.bfloat16: 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" ) if comp_type == torch.float16 or comm_type == torch.float16: if not get_accelerator().is_fp16_supported(): pytest.skip("fp16 is not supported") type_str = {torch.float16: "fp16", torch.bfloat16: "bf16"} config_dict = { "train_micro_batch_size_per_gpu": 2, "steps_per_print": 1, "fp16": { "enabled": comp_type == torch.float16 }, "bf16": { "enabled": comp_type == torch.bfloat16 }, "zero_optimization": { "stage": 2 }, } if comm_type is not None: config_dict["communication_data_type"] = type_str[comm_type] else: comm_type = comp_type hidden_dim = 10 model = SimpleModel(hidden_dim) optimizer = torch.optim.Adam(model.parameters()) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer) data_loader = random_dataloader(model=model, total_samples=2, hidden_dim=hidden_dim, device=model.device, dtype=comp_type) def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): assert tensor.dtype == comm_type return orig_torch_reduce(tensor, dst, op, group, async_op) orig_torch_reduce = dist.reduce dist.reduce = custom_reduce for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() dist.reduce = orig_torch_reduce @pytest.mark.parametrize("bf16_optimizer_states,use_cpu_offload,zero_stage", [ pytest.param(False, True, 1, id="zero_stage_1_cpu_offload"), pytest.param(True, False, 1, id="zero_stage_1_bf16_opt_states_True"), pytest.param(True, True, 1, id="zero_stage_1_bf16_opt_states_cpu_offload"), pytest.param(False, True, 2, id="zero_stage_2_cpu_offload"), pytest.param(True, False, 2, id="zero_stage_2_bf16_opt_states_True"), pytest.param(True, True, 2, id="zero_stage_2_bf16_opt_states_cpu_offload"), pytest.param(False, True, 3, id="zero_stage_3_cpu_offload"), pytest.param(True, False, 3, id="zero_stage_3_bf16_opt_states_True"), pytest.param(True, True, 3, id="zero_stage_3_bf16_opt_states_cpu_offload"), ]) class TestBF16MasterWeightsGradients(DistributedTest): world_size = 2 def test_gradients_match_ddp(self, bf16_optimizer_states, use_cpu_offload, zero_stage): 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" ) if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: pytest.skip("cpu-adam is not compatible") hidden_dim = 6 lr = 1e-3 seed = 123 device, rank, dtype = initialize_distributed() model_ddp, optimizer_ddp = create_ddp_model(SimpleModel, device, rank, dtype, seed=seed, lr=lr, hidden_dim=hidden_dim, nlayers=2) torch.manual_seed(seed) ds_model = SimpleModel(hidden_dim, nlayers=2) bf16_config = { "enabled": True, "bf16_master_weights_and_grads": True, } if bf16_optimizer_states: bf16_config["bf16_optimizer_states"] = True zero_config = {"stage": zero_stage} if use_cpu_offload: zero_config["cpu_offload"] = True config_dict = { "train_micro_batch_size_per_gpu": 2, "gradient_accumulation_steps": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": lr } }, "bf16": bf16_config, "zero_optimization": zero_config } engine, _, _, _ = deepspeed.initialize(config=config_dict, model=ds_model, model_parameters=ds_model.parameters()) data_loader = random_dataloader(model=engine, total_samples=8, hidden_dim=hidden_dim, device=device, dtype=torch.bfloat16) batch = next(iter(data_loader)) optimizer_ddp.zero_grad() loss_ddp = model_ddp(batch[0], batch[1]) loss_ddp.backward() grads_ddp = collect_ddp_gradients(model_ddp) loss_ds = engine(batch[0], batch[1]) loss_ds.backward() grads_ds = collect_gradients_safe(engine) compare_gradients( grads_ddp, grads_ds, step_info= f"bf16_optimizer_states={bf16_optimizer_states}, cpu_offload={use_cpu_offload}, zero_stage={zero_stage}") optimizer_ddp.step() optimizer_ddp.zero_grad() engine.step() engine.zero_grad() if bf16_optimizer_states and use_cpu_offload: # With CPU offload the Adam moments must be allocated in bf16 on the host so the # offloaded optimizer-state footprint is smaller than with fp32 moments. cpu_adam_state = engine.optimizer.optimizer.state moment_tensors = [] for param_state in cpu_adam_state.values(): for moment_key in ("exp_avg", "exp_avg_sq"): if moment_key in param_state: moment_tensors.append(param_state[moment_key]) assert moment_tensors, "expected Adam moment tensors to be allocated after a step" for moment in moment_tensors: assert moment.dtype == torch.bfloat16, f"expected bf16 moment, got {moment.dtype}" assert moment.device.type == "cpu", f"expected moment on cpu, got {moment.device}" engine.destroy() @pytest.mark.parametrize("zero_stage", [1, 2, 3]) class TestBF16OptimizerStatesOffloadValidation(DistributedTest): world_size = 1 def test_user_cpu_adam_must_enable_bf16_states(self, zero_stage): """A user-provided DeepSpeedCPUAdam must be built with fp32_optimizer_states=False to combine bf16_optimizer_states with ZeRO-Offload, otherwise the moments would silently stay fp32 and the memory benefit would be lost.""" 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" ) if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: pytest.skip("cpu-adam is not compatible") from deepspeed.ops.adam import DeepSpeedCPUAdam hidden_dim = 6 model = SimpleModel(hidden_dim, nlayers=2) # fp32_optimizer_states defaults to True, which keeps fp32 moments and is # incompatible with bf16_optimizer_states under ZeRO-Offload. optimizer = DeepSpeedCPUAdam(model.parameters()) config_dict = { "train_micro_batch_size_per_gpu": 2, "steps_per_print": 1, "bf16": { "enabled": True, "bf16_master_weights_and_grads": True, "bf16_optimizer_states": True, }, # offload_optimizer is the current config key for ZeRO optimizer offload # (TestBF16MasterWeightsGradients above still uses the legacy cpu_offload alias). "zero_optimization": { "stage": zero_stage, "offload_optimizer": { "device": "cpu" }, }, } with pytest.raises(AssertionError, match="fp32_optimizer_states=False"): deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)