# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team import pytest import torch from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum from deepspeed.utils.torch import required_torch_version from deepspeed.accelerator import get_accelerator from unit.v1.compile.util import compare_loss from unit.common import DistributedTest from unit.simple_model import SimpleModel from unit.util import bf16_required_version_check, skip_on_arch import deepspeed from deepspeed.ops.aio import AsyncIOBuilder pytestmark = pytest.mark.skipif(not required_torch_version(min_version=2.1), reason="Compile tests requires Pytorch version 2.1 or above") class TestZeRO(DistributedTest): world_size = 2 non_daemonic_procs = True @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32]) @pytest.mark.parametrize('zero_stage', [1, 2, 3]) @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme]) def test_compile_zero(self, tmpdir, zero_stage, dtype, offload_device): if dtype == torch.bfloat16: skip_on_arch(min_arch=8) if dtype == torch.bfloat16 and not bf16_required_version_check(): pytest.skip( "DeepSpeed BFloat16 tests need NCCL >= 2.10.3, CUDA >=11.0, and HW support for BFloat16 to run correctly" ) if get_accelerator().device_name() == "cpu": pytest.skip("CPU does not support this test yet") if offload_device == OffloadDeviceEnum.nvme: if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: pytest.skip('Skip tests since async-io is not compatible') if zero_stage != 3: pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "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} compare_loss(self, config_dict, dtype) class TestDeepCompile(DistributedTest): world_size = 2 non_daemonic_procs = True @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32]) @pytest.mark.parametrize('zero_stage', [1, 3]) @pytest.mark.parametrize('deepcompile', [True]) # deepcompile==False is included in test_compile_zero def test(self, zero_stage, dtype, deepcompile): if not required_torch_version(min_version=2.6): pytest.skip("DeepCompile requires PyTorch >= v2.6") if dtype == torch.bfloat16: skip_on_arch(min_arch=8) if dtype == torch.bfloat16 and not bf16_required_version_check(): pytest.skip( "DeepSpeed BFloat16 tests need NCCL >= 2.10.3, CUDA >=11.0, and HW support for BFloat16 to run correctly" ) if get_accelerator().device_name() == "cpu": pytest.skip("CPU does not support this test yet") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "zero_optimization": { "stage": zero_stage, }, "compile": { "deepcompile": deepcompile } } if dtype == torch.float16: config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8} elif dtype == torch.bfloat16: config_dict["bf16"] = {"enabled": True} # Need warmup steps compare_loss(self, config_dict, dtype, iteration=10) def test_zero1_releases_grad_buffers_after_optimizer_step(self): if not required_torch_version(min_version=2.6): pytest.skip("DeepCompile requires PyTorch >= v2.6") if get_accelerator().device_name() == "cpu": pytest.skip("CPU does not support this test yet") dtype = torch.float32 hidden_dim = 10 config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "zero_optimization": { "stage": 1, }, "compile": { "deepcompile": True } } model = SimpleModel(hidden_dim) engine, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) engine.compile() device = torch.device(get_accelerator().current_device_name()) x = torch.randn(config_dict["train_micro_batch_size_per_gpu"], hidden_dim, device=device, dtype=dtype) y = torch.randn_like(x) loss = engine(x, y) engine.backward(loss) optimizer = engine.optimizer current_grad_buffers = optimizer.averaged_gradients assert current_grad_buffers assert all(group_buffers is not None for group_buffers in current_grad_buffers.values()) assert any(buffer.numel() > 0 for group_buffers in current_grad_buffers.values() for buffer in group_buffers) for group_idx, group_buffers in current_grad_buffers.items(): assert group_buffers.flat_partition.numel() == optimizer.partition_size[group_idx] assert callable(group_buffers.release_grad_buffers) engine.step() assert all(group_buffers is None for group_buffers in optimizer.averaged_gradients.values()) engine.destroy() @pytest.mark.parametrize('dtype', [torch.float32]) @pytest.mark.parametrize('zero_stage', [3]) def test_padded_shard_handling(self, zero_stage, dtype): """Test that parameters with padding (uneven division) work correctly with DeepCompile""" if not required_torch_version(min_version=2.6): pytest.skip("DeepCompile requires PyTorch >= v2.6") if get_accelerator().device_name() == "cpu": pytest.skip("CPU does not support this test yet") # Use a hidden dimension that requires padding when divided across ranks # With world_size=2, a hidden_dim of 13 creates parameters that need padding config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "zero_optimization": { "stage": zero_stage, }, "compile": { "deepcompile": True } } # This should work correctly with our padding-aware implementation # The test verifies that padded parameters are handled properly compare_loss(self, config_dict, dtype, iteration=1, hidden_dim_override=13) @pytest.mark.parametrize('dtype', [torch.float32]) @pytest.mark.parametrize('zero_stage', [1, 3]) def test_free_activation_mode(self, zero_stage, dtype): """Test that eagerly free activations work correctly and the threshold is configurable""" if not required_torch_version(min_version=2.6): pytest.skip("DeepCompile requires PyTorch >= v2.6") if get_accelerator().device_name() == "cpu": pytest.skip("CPU does not support this test yet") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "zero_optimization": { "stage": zero_stage, }, "compile": { "deepcompile": True, "free_activation": True, "free_activation_threshold": 0, } } compare_loss(self, config_dict, dtype) @pytest.mark.parametrize('dtype', ["bfloat16", "float16"]) @pytest.mark.parametrize('zero_stage', [3]) def test_fusing_allgather_and_autocast(self, zero_stage, dtype): """Test that allgather and autocast can be correctly fused with DeepCompile""" if not required_torch_version(min_version=2.6): pytest.skip("DeepCompile requires PyTorch >= v2.6") if get_accelerator().device_name() == "cpu": pytest.skip("CPU does not support this test yet") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "torch_autocast": { "enable": True, "dtype": dtype, }, "zero_optimization": { "stage": zero_stage, }, "compile": { "deepcompile": True } } compare_loss(self, config_dict, torch.float32)