# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team import pytest import deepspeed.comm as dist import torch import math from unit.common import DistributedTest from unit.simple_model import random_dataloader, SimpleModel from unit.util import bf16_required_version_check import deepspeed from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_grad, safe_set_full_optimizer_state from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state from deepspeed.utils import safe_update_full_grad_vectorized from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum from deepspeed.ops.aio import AsyncIOBuilder from deepspeed.accelerator import get_accelerator from deepspeed.runtime.swap_tensor import MIN_SWAPPABLE_BYTES WEIGHT_KEY = 'weight' FIRST_ORDER_KEY = 'exp_avg' SECOND_ORDER_KEY = 'exp_avg_sq' GRADIENT_KEY = 'gradient' def validate_tensor(model, api_type, opt_states): assert api_type in ["full", "local"] for _, lp in model.named_parameters(): param_list = [] if opt_states: param_list.append( safe_get_full_optimizer_state(lp, 'exp_avg') if api_type == "full" else safe_get_local_optimizer_state(lp, 'exp_avg')) param_list.append( safe_get_full_optimizer_state(lp, 'exp_avg_sq') if api_type == "full" else safe_get_local_optimizer_state(lp, 'exp_avg_sq')) else: param_list.append(safe_get_full_fp32_param(lp) if api_type == "full" else safe_get_local_fp32_param(lp)) param_list.append(safe_get_full_grad(lp) if api_type == "full" else safe_get_local_grad(lp)) if lp.requires_grad: assert all([p is not None for p in param_list]) else: assert all([p is None for p in param_list]) class MyModel(torch.nn.Module): def __init__(self, hidden_dim, frozen_weights): super(MyModel, self).__init__() self.act = torch.nn.ReLU() self.cel = torch.nn.CrossEntropyLoss() self.linears = torch.nn.ModuleList( [torch.nn.Linear(hidden_dim, 1), torch.nn.Linear(1, 1), torch.nn.Linear(1, hidden_dim)]) if frozen_weights: self.linears[0].weight.requires_grad = False self.linears[0].bias.requires_grad = False def forward(self, x, y): for l in self.linears: x = l(x) x = self.act(x) return self.cel(x, y) def run_fragmented_model(model, config_dict, hidden_dim, dtype, validate_after_bwd, validate_after_step): model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict) data_loader = random_dataloader(model=model, total_samples=10, hidden_dim=hidden_dim, device=model.device, dtype=dtype) dist.barrier() for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) validate_after_bwd(model) model.step() validate_after_step(model) # Needed in ZeRO 3. Not doing so can give memory leak model.destroy() @pytest.mark.parametrize('frozen_weights', [True, False]) class TestTensorFragmentGet(DistributedTest): # Need multiple gpus to test possible hanging world_size = 2 reuse_dist_env = True @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32]) @pytest.mark.parametrize('api_type', ['local', 'full']) @pytest.mark.parametrize('zero_stage', [1, 2, 3]) @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme]) def test_zero_fragments(self, tmpdir, dtype, api_type, zero_stage, offload_device, frozen_weights): if not dtype in get_accelerator().supported_dtypes(): pytest.skip(f"{get_accelerator()._name} does not support {dtype} data type") if offload_device == OffloadDeviceEnum.nvme: if zero_stage != 3: pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}") if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: pytest.skip('Skip tests since async-io is not compatible') if api_type == "local" and zero_stage != 3: pytest.skip(f"Local APIs only for zero stage 3 but current stage is {zero_stage}") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 1e-6 } }, "zero_optimization": { "stage": zero_stage, } } if dtype == torch.half: config_dict["fp16"] = {"enabled": True, "initial_scale_power": 2} elif dtype == torch.bfloat16: config_dict["bf16"] = {"enabled": True} if offload_device == OffloadDeviceEnum.cpu: config_dict["zero_optimization"]["offload_optimizer"] = {"device": offload_device} elif offload_device == OffloadDeviceEnum.nvme: config_dict["zero_optimization"]["offload_optimizer"] = { "device": offload_device, "nvme_path": str(tmpdir) } hidden_dim = MIN_SWAPPABLE_BYTES if zero_stage == 3: with deepspeed.zero.Init(config_dict_or_path=config_dict): model = MyModel(hidden_dim, frozen_weights) else: model = MyModel(hidden_dim, frozen_weights) validate_after_bwd = lambda model: validate_tensor(model, api_type, opt_states=False) validate_after_step = lambda model: validate_tensor(model, api_type, opt_states=True) run_fragmented_model(model, config_dict, hidden_dim, dtype, validate_after_bwd, validate_after_step) def test_bf16_optimizer_fragments(self, frozen_weights): if get_accelerator().device_name() == "cpu": pytest.skip("CPU accelerator does not support this test yet.") if frozen_weights: pytest.skip("TODO: Frozen weights not currently supported by BF16 Optimizer") 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, "optimizer": { "type": "Adam", "params": { "lr": 1e-6 } }, "bf16": { "enabled": True }, # Use fp32 gradient accumulation to ensure BF16_Optimizer is used # (bf16 model + bf16 grad_accum uses FP16_Optimizer which doesn't support tensor fragment APIs) "data_types": { "grad_accum_dtype": "fp32" }, "zero_optimization": { "stage": 1, } } hidden_dim = 128 model = MyModel(hidden_dim, frozen_weights) api_type = "full" validate_after_bwd = lambda model: validate_tensor(model, api_type, opt_states=False) validate_after_step = lambda model: validate_tensor(model, api_type, opt_states=True) run_fragmented_model(model, config_dict, hidden_dim, torch.bfloat16, validate_after_bwd, validate_after_step) def create_random_values(model, key_list, group, grad_dtype): param_values = {} for n, lp in model.named_parameters(): param_shape = lp.ds_shape if hasattr(lp, 'ds_id') else lp.shape param_values[n] = {} for key in key_list: dtype = grad_dtype if key == GRADIENT_KEY else torch.float32 rand_value = torch.rand(param_shape, dtype=dtype, device=model.device) dist.broadcast(rand_value, src=0, group=group) param_values[n][key] = rand_value return param_values def set_param_values_with_dict(model, value_dict): for n, lp in model.named_parameters(): for key, value_tensor in value_dict[n].items(): if key == GRADIENT_KEY: safe_set_full_grad(lp, value_tensor) elif key == WEIGHT_KEY: safe_set_full_fp32_param(lp, value_tensor) else: safe_set_full_optimizer_state(lp, value_tensor, key) def update_param_values_with_dict(model, value_dict): new_grad_values = {} for n, lp in model.named_parameters(): if GRADIENT_KEY in value_dict[n]: new_grad_values[id(lp)] = value_dict[n][GRADIENT_KEY] def update_gradient_callback(old_value, param): return new_grad_values[id(param)] update_param_list = [] for n, lp in model.named_parameters(): for key, value_tensor in value_dict[n].items(): if key == GRADIENT_KEY: update_param_list.append(lp) if len(update_param_list) > 0: safe_update_full_grad_vectorized(update_param_list, update_gradient_callback) def validate_param_values_with_dict(model, value_dict): for n, lp in model.named_parameters(): for key, expected_tensor in value_dict[n].items(): if key == GRADIENT_KEY: actual_tensor = safe_get_full_grad(lp) elif key == WEIGHT_KEY: actual_tensor = safe_get_full_fp32_param(lp) else: actual_tensor = safe_get_full_optimizer_state(lp, key) assert torch.equal(expected_tensor, actual_tensor) def create_random_values_for_local(model, key_list, group, grad_dtype): param_values = {} for n, lp in model.named_parameters(): param_shape = lp.ds_tensor.shape param_values[n] = {} for key in key_list: dtype = grad_dtype if key == GRADIENT_KEY else torch.float32 rand_value = torch.rand(param_shape, dtype=dtype, device=model.device) param_values[n][key] = rand_value return param_values def set_local_param_values_with_dict(model, value_dict): for n, lp in model.named_parameters(): for key, value_tensor in value_dict[n].items(): if key == GRADIENT_KEY: safe_set_local_grad(lp, value_tensor) elif key == WEIGHT_KEY: safe_set_local_fp32_param(lp, value_tensor) else: safe_set_local_optimizer_state(lp, value_tensor, key) def validate_local_param_values_with_dict(model, value_dict): for n, lp in model.named_parameters(): for key, expected_tensor in value_dict[n].items(): if key == GRADIENT_KEY: actual_tensor = safe_get_local_grad(lp) elif key == WEIGHT_KEY: actual_tensor = safe_get_local_fp32_param(lp) else: actual_tensor = safe_get_local_optimizer_state(lp, key) assert torch.equal(expected_tensor, actual_tensor) helper_funcs_mapping = { "full": { "create_random_values": create_random_values, "set_param_values_with_dict": set_param_values_with_dict, "update_param_values_with_dict": update_param_values_with_dict, "validate_param_values_with_dict": validate_param_values_with_dict, }, "local": { "create_random_values": create_random_values_for_local, "set_param_values_with_dict": set_local_param_values_with_dict, "validate_param_values_with_dict": validate_local_param_values_with_dict } } @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32]) class TestTensorFragmentSet(DistributedTest): # Need multiple gpus to test possible hanging world_size = 2 reuse_dist_env = True @pytest.mark.parametrize('api_type', ['local', 'full']) @pytest.mark.parametrize('zero_stage', [1, 2, 3]) @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme]) def test_zero_fragments(self, tmpdir, api_type, zero_stage, offload_device, dtype): if not dtype in get_accelerator().supported_dtypes(): pytest.skip(f"{get_accelerator()._name} does not support {dtype} data type") if dtype == torch.bfloat16 and not bf16_required_version_check(accelerator_check=False): 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 api_type == "local" and zero_stage != 3: pytest.skip(f"Local APIs only for zero stage 3 but current stage is {zero_stage}") if offload_device == OffloadDeviceEnum.nvme: if zero_stage != 3: pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}") if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: pytest.skip('Skip tests since async-io is not compatible') config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 1e-6 } }, "zero_optimization": { "stage": zero_stage, } } if offload_device == OffloadDeviceEnum.cpu: config_dict["zero_optimization"]["offload_optimizer"] = {"device": offload_device} elif offload_device == OffloadDeviceEnum.nvme: config_dict["zero_optimization"]["offload_optimizer"] = { "device": offload_device, "nvme_path": str(tmpdir) } if dtype == torch.float16: config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8} elif dtype == torch.bfloat16: config_dict["bf16"] = {"enabled": True} hidden_dim = int(math.sqrt(MIN_SWAPPABLE_BYTES)) if zero_stage == 3: config_dict["zero_optimization"]["param_persistence_threshold"] = hidden_dim with deepspeed.zero.Init(config_dict_or_path=config_dict): model = SimpleModel(hidden_dim) else: model = SimpleModel(hidden_dim) world = dist.get_world_size() group = dist.new_group(ranks=list(range(world))) dist.barrier() def after_bwd_validate_func(model): state_keys = [WEIGHT_KEY, GRADIENT_KEY] helper_funcs = helper_funcs_mapping[api_type] optim_state_values = helper_funcs["create_random_values"](model, state_keys, group, grad_dtype=dtype) helper_funcs["set_param_values_with_dict"](model, optim_state_values) helper_funcs["validate_param_values_with_dict"](model, optim_state_values) def after_step_validate_func(model): state_keys = [WEIGHT_KEY, FIRST_ORDER_KEY, SECOND_ORDER_KEY] helper_funcs = helper_funcs_mapping[api_type] optim_state_values = helper_funcs["create_random_values"](model, state_keys, group, grad_dtype=dtype) helper_funcs["set_param_values_with_dict"](model, optim_state_values) helper_funcs["validate_param_values_with_dict"](model, optim_state_values) run_fragmented_model(model, config_dict, hidden_dim, dtype, after_bwd_validate_func, after_step_validate_func) @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32]) class TestTensorFragmentUpdate(DistributedTest): # Need multiple gpus to test possible hanging world_size = 2 reuse_dist_env = True @pytest.mark.parametrize('torch_adam', [False, True]) @pytest.mark.parametrize('zero_stage', [1, 2, 3]) @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme]) def test_zero_fragments(self, tmpdir, torch_adam, zero_stage, offload_device, dtype): if not dtype in get_accelerator().supported_dtypes(): pytest.skip(f"{get_accelerator()._name} does not support {dtype} data type") if offload_device == OffloadDeviceEnum.nvme: if zero_stage != 3: pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}") if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: pytest.skip('Skip tests since async-io is not compatible') config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 1e-6, "torch_adam": torch_adam } }, "zero_optimization": { "stage": zero_stage, } } if offload_device == OffloadDeviceEnum.cpu: config_dict["zero_optimization"]["offload_optimizer"] = {"device": offload_device} elif offload_device == OffloadDeviceEnum.nvme: config_dict["zero_optimization"]["offload_optimizer"] = { "device": offload_device, "nvme_path": str(tmpdir) } if dtype == torch.float16: config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8} elif dtype == torch.bfloat16: config_dict["bf16"] = {"enabled": True} hidden_dim = int(math.sqrt(MIN_SWAPPABLE_BYTES)) if zero_stage == 3: config_dict["zero_optimization"]["param_persistence_threshold"] = hidden_dim with deepspeed.zero.Init(config_dict_or_path=config_dict): model = SimpleModel(hidden_dim) else: model = SimpleModel(hidden_dim) world = dist.get_world_size() group = dist.new_group(ranks=list(range(world))) dist.barrier() api_type = "full" def after_bwd_validate_func(model): state_keys = [GRADIENT_KEY] helper_funcs = helper_funcs_mapping[api_type] optim_state_values = helper_funcs["create_random_values"](model, state_keys, group, grad_dtype=dtype) helper_funcs["update_param_values_with_dict"](model, optim_state_values) helper_funcs["validate_param_values_with_dict"](model, optim_state_values) def after_step_validate_func(model): pass run_fragmented_model(model, config_dict, hidden_dim, dtype, after_bwd_validate_func, after_step_validate_func)