test_cpu_adam.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os
import sys
import torch
import numpy as np
import pytest
from cpuinfo import get_cpu_info

import deepspeed
from deepspeed.accelerator import get_accelerator
from deepspeed.ops.adam import FusedAdam
from deepspeed.ops.op_builder import CPUAdamBuilder, FusedAdamBuilder
from unit.common import DistributedTest

if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
    pytest.skip("cpu-adam is not compatible", allow_module_level=True)

pytest.cpu_vendor = get_cpu_info().get("vendor_id_raw", "").lower()


def check_equal(first, second, atol=1e-2, verbose=False):
    x = first.detach().float().numpy()
    y = second.detach().float().numpy()
    print("ATOL", atol)
    if verbose:
        print("x = {}".format(x.flatten()))
        print("y = {}".format(y.flatten()))
        print('-' * 80)
    np.testing.assert_allclose(x, y, err_msg="param-update mismatch!", atol=atol)


def _compare_optimizers(model_size, param1, optimizer1, param2, optimizer2):
    for i in range(10):
        param1.grad = torch.randn(model_size, device=param1.device).to(param1.dtype)
        param2.grad = param1.grad.clone().detach().to(device=param2.device, dtype=param2.dtype)

        optimizer1.step()
        optimizer2.step()

    tolerance = param1.float().norm().detach().numpy() * 1e-2
    check_equal(param1.float().norm(), param2.float().cpu().norm(), atol=tolerance, verbose=True)


@pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16, torch.float], ids=["fp16", "bf16", "fp32"])
@pytest.mark.parametrize('model_size',
                         [
                             (64),
                             (22),
                             #(55),
                             (128),
                             (1024),
                             (1048576),
                         ]) # yapf: disable
class TestCPUAdam(DistributedTest):
    world_size = 1
    reuse_dist_env = True
    requires_cuda_env = False
    if not get_accelerator().is_available():
        init_distributed = False
        set_dist_env = False

    @pytest.mark.skipif(not get_accelerator().is_available(), reason="only supported in CUDA environments.")
    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME],
                        reason="FusedAdam is not compatible")
    def test_fused_adam_equal(self, dtype, model_size):
        if dtype not in get_accelerator().supported_dtypes():
            pytest.skip(f"dtype {dtype} not supported in current accelerator")

        if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
            pytest.skip("cpu-adam with half precision not supported on AMD CPUs")

        from deepspeed.ops.adam import DeepSpeedCPUAdam

        cpu_data = torch.randn(model_size, device='cpu').to(dtype)
        cpu_param = torch.nn.Parameter(cpu_data)
        cuda_param = torch.nn.Parameter(cpu_data.to(get_accelerator().device_name()))

        # tolerance = cpu_param.float().norm().detach().numpy() * 1e-2
        # check_equal(cpu_param.float().norm(),
        #             cuda_param.float().cpu().norm(),
        #             atol=tolerance,
        #             verbose=True)

        cpu_optimizer = DeepSpeedCPUAdam([cpu_param])
        cuda_optimizer = FusedAdam([cuda_param])

        _compare_optimizers(model_size=model_size,
                            param1=cpu_param,
                            optimizer1=cpu_optimizer,
                            param2=cuda_param,
                            optimizer2=cuda_optimizer)

    def test_torch_adamw_equal(self, dtype, model_size):
        if get_accelerator().is_available():
            if dtype == torch.half:
                pytest.skip("torch.optim.AdamW with half precision inf/nan output.")
            if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
                pytest.skip("cpu-adam with half precision not supported on AMD CPUs")
            ref_param_device = get_accelerator().device_name()
        else:
            if dtype == torch.half:
                pytest.skip("torch.optim.AdamW with half precision only supported in CUDA environments.")
            ref_param_device = 'cpu'

        from deepspeed.ops.adam import DeepSpeedCPUAdam

        cpu_data = torch.randn(model_size, device='cpu').to(dtype)
        cpu_param = torch.nn.Parameter(cpu_data)
        ref_param = torch.nn.Parameter(cpu_data.to(ref_param_device))

        cpu_optimizer = DeepSpeedCPUAdam([cpu_param])
        ref_optimizer = torch.optim.AdamW([ref_param])

        _compare_optimizers(model_size=model_size,
                            param1=cpu_param,
                            optimizer1=cpu_optimizer,
                            param2=ref_param,
                            optimizer2=ref_optimizer)


class TestCPUAdamBf16OptimizerStates(DistributedTest):
    world_size = 1
    reuse_dist_env = True
    requires_cuda_env = False
    if not get_accelerator().is_available():
        init_distributed = False
        set_dist_env = False

    @pytest.mark.parametrize('model_size', [64, 1024])
    def test_bf16_optimizer_states_dtype(self, model_size):
        """fp32_optimizer_states=False keeps the Adam moments in the bf16 parameter precision."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        param = torch.nn.Parameter(torch.randn(model_size, device='cpu', dtype=torch.bfloat16))
        optimizer = DeepSpeedCPUAdam([param], fp32_optimizer_states=False)
        param.grad = torch.randn(model_size, device='cpu', dtype=torch.bfloat16)
        optimizer.step()

        state = optimizer.state[param]
        assert state['exp_avg'].dtype == torch.bfloat16
        assert state['exp_avg_sq'].dtype == torch.bfloat16
        assert state['exp_avg'].device == torch.device('cpu')
        assert state['exp_avg_sq'].device == torch.device('cpu')

    @pytest.mark.parametrize('model_size', [64, 1024])
    def test_bf16_optimizer_states_match_fp32(self, model_size):
        """bf16 moments should track fp32 moments within bf16 tolerance over several steps."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        torch.manual_seed(0)
        base = torch.randn(model_size, device='cpu', dtype=torch.float32).to(torch.bfloat16)
        param_fp32_states = torch.nn.Parameter(base.clone())
        param_bf16_states = torch.nn.Parameter(base.clone())

        opt_fp32_states = DeepSpeedCPUAdam([param_fp32_states], fp32_optimizer_states=True)
        opt_bf16_states = DeepSpeedCPUAdam([param_bf16_states], fp32_optimizer_states=False)

        for _ in range(10):
            grad = torch.randn(model_size, device='cpu', dtype=torch.bfloat16)
            param_fp32_states.grad = grad.clone()
            param_bf16_states.grad = grad.clone()
            opt_fp32_states.step()
            opt_bf16_states.step()

        assert opt_fp32_states.state[param_fp32_states]['exp_avg'].dtype == torch.float32
        assert opt_bf16_states.state[param_bf16_states]['exp_avg'].dtype == torch.bfloat16

        # bf16 moments round every Adam update to an 8-bit mantissa, so over 10 steps they
        # diverge from fp32 moments more than the same-precision comparison in _compare_optimizers
        # (1e-2). A wider 5% band keeps this stable while still catching gross errors; the dtype
        # assertions above guard the precision itself. Norm comparison follows _compare_optimizers.
        tolerance = param_fp32_states.float().norm().detach().numpy() * 5e-2
        check_equal(param_fp32_states.float().norm(), param_bf16_states.float().norm(), atol=tolerance)


def _zenflow_adam_proc_worker(param, g0, g1, ea0, ea1, eq0, eq1, stale, ctrl, ready, affinity):
    op = CPUAdamBuilder().load()
    op.create_adam(0, 1e-3, 0.9, 0.999, 1e-8, 0.0, True, False)
    handle = op.zenflow_adam_create(0, affinity)
    op.zenflow_adam_register_group(handle, param, g0, g1, ea0, ea1, eq0, eq1, stale)
    ready.set()
    op.zenflow_adam_run_worker(handle, ctrl.data_ptr())  # blocks until the exit command
    op.zenflow_adam_destroy(handle)
    op.destroy_adam(0)


@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="cross-process ZenFlowAdam is Linux-only")
def test_zenflow_adam_cross_process():
    """The optimizer-process driver (shared-memory semaphore control + native worker, the
    production path for ZenFlow stage 1/2 overlap) must match a per-parameter adam_update
    reference bit-for-bit with alternating double buffers. Run as a plain test, not
    DistributedTest, so the pytest process (non-daemonic) can spawn the optimizer process."""
    import torch.multiprocessing as mp

    op = CPUAdamBuilder().load()
    if not hasattr(op, "zenflow_adam_ctrl_size"):
        pytest.skip("cross-process ZenFlowAdam not available in this build")

    lr, beta1, beta2, eps, wd = 1e-3, 0.9, 0.999, 1e-8, 0.0
    n = 100003  # non-SIMD-aligned, exercises the scalar tail
    affinity = list(range(min(4, os.cpu_count() or 1)))

    ctrl = torch.zeros(op.zenflow_adam_ctrl_size(), dtype=torch.uint8).share_memory_()
    op.zenflow_adam_ctrl_init(ctrl.data_ptr(), 1)

    torch.manual_seed(0)
    param = torch.randn(n).share_memory_()
    g = [torch.zeros(n).share_memory_(), torch.zeros(n).share_memory_()]
    ea = [torch.zeros(n).share_memory_(), torch.zeros(n).share_memory_()]
    eq = [torch.zeros(n).share_memory_(), torch.zeros(n).share_memory_()]
    stale = torch.zeros(n).share_memory_()

    op.create_adam(1, lr, beta1, beta2, eps, wd, True, False)
    p_ref = param.clone()
    ea_ref = [ea[0].clone(), ea[1].clone()]
    eq_ref = [eq[0].clone(), eq[1].clone()]
    st_ref = stale.clone()

    ctx = mp.get_context("spawn")
    ready = ctx.Event()
    proc = ctx.Process(target=_zenflow_adam_proc_worker,
                       args=(param, g[0], g[1], ea[0], ea[1], eq[0], eq[1], stale, ctrl, ready, affinity))
    proc.start()
    try:
        assert ready.wait(timeout=60), "optimizer process did not start"
        # With no step submitted yet, a bounded wait must time out (return False) rather than
        # block -- this is what lets the training side notice a dead optimizer process.
        assert op.zenflow_adam_wait(ctrl.data_ptr(), 0.05) is False, "wait should time out when no step is pending"
        for step in range(1, 6):
            now = step & 1
            grad = torch.randn(n)
            g[now].copy_(grad)
            op.zenflow_adam_submit(ctrl.data_ptr(), now, step, [lr], [beta1], [beta2], [eps], [wd], [1])
            assert op.zenflow_adam_wait(ctrl.data_ptr(), 60.0), f"wait timed out step {step}"
            # Reference: single-tensor adam_update on the mirror, then snapshot the updated
            # param into the stale buffer -- exactly what the native worker does per group.
            op.adam_update(1, step, lr, beta1, beta2, eps, wd, True, p_ref, grad.clone(), ea_ref[now], eq_ref[now])
            st_ref.copy_(p_ref)
            assert torch.equal(param, p_ref), f"param mismatch step {step}"
            assert torch.equal(ea[now], ea_ref[now]), f"exp_avg mismatch step {step}"
            assert torch.equal(eq[now], eq_ref[now]), f"exp_avg_sq mismatch step {step}"
            assert torch.equal(stale, st_ref), f"stale mismatch step {step}"
        op.zenflow_adam_ctrl_exit(ctrl.data_ptr())
        proc.join(timeout=10)
    finally:
        if proc.is_alive():
            proc.terminate()
            proc.join(timeout=5)
        op.destroy_adam(1)


class TestCPUAdamGPUError(DistributedTest):

    def test_cpu_adam_gpu_error(self):
        model_size = 64
        from deepspeed.ops.adam import DeepSpeedCPUAdam
        device = get_accelerator().device_name(0)  # 'cuda:0' or 'xpu:0'
        param = torch.nn.Parameter(torch.randn(model_size, device=device))
        optimizer = DeepSpeedCPUAdam([param])

        param.grad = torch.randn(model_size, device=device)
        with pytest.raises(AssertionError):
            optimizer.step()


class TestCPUAdamSubgroup(DistributedTest):
    world_size = 1
    reuse_dist_env = True
    requires_cuda_env = False
    if not get_accelerator().is_available():
        init_distributed = False
        set_dist_env = False

    @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16], ids=["fp16", "bf16"])
    @pytest.mark.parametrize('model_size', [64, 128, 1024])
    def test_step_subgroup_basic(self, dtype, model_size):
        """Test basic functionality of step_subgroup method."""
        if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
            pytest.skip("cpu-adam with half precision not supported on AMD CPUs")

        from deepspeed.ops.adam import DeepSpeedCPUAdam

        # Create parameters
        cpu_data = torch.randn(model_size, device='cpu').to(dtype)
        param = torch.nn.Parameter(cpu_data)
        optimizer = DeepSpeedCPUAdam([param])

        # Set gradient
        param.grad = torch.randn(model_size, device='cpu').to(dtype)

        # Store initial parameter values
        initial_param = param.data.clone()

        # Test step_subgroup with subgroup_id=0
        subgroup_id = 0
        optimizer.step_subgroup(subgroup_id)

        # Verify parameter was updated
        assert not torch.equal(param.data, initial_param), "Parameters should be updated after step_subgroup"

        # Verify optimizer state was created for subgroup
        assert subgroup_id in optimizer.state, "Optimizer state should be created for subgroup"
        assert optimizer.state[subgroup_id]['step'] == 1, "Step count should be 1"
        assert 'exp_avg' in optimizer.state[subgroup_id], "exp_avg should be in state"
        assert 'exp_avg_sq' in optimizer.state[subgroup_id], "exp_avg_sq should be in state"

    @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16], ids=["fp16", "bf16"])
    def test_step_subgroup_multiple_calls(self, dtype):
        """Test multiple calls to step_subgroup increment step count correctly."""
        if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
            pytest.skip("cpu-adam with half precision not supported on AMD CPUs")

        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 64
        cpu_data = torch.randn(model_size, device='cpu').to(dtype)
        param = torch.nn.Parameter(cpu_data)
        optimizer = DeepSpeedCPUAdam([param])

        subgroup_id = 0

        # Perform multiple steps
        for step in range(1, 4):
            param.grad = torch.randn(model_size, device='cpu').to(dtype)
            optimizer.step_subgroup(subgroup_id)

            # Verify step count increments
            assert optimizer.state[subgroup_id]['step'] == step, f"Step count should be {step}"

    @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16], ids=["fp16", "bf16"])
    def test_rollback_subgroup_basic(self, dtype):
        """Test basic functionality of rollback_subgroup method."""
        if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
            pytest.skip("cpu-adam with half precision not supported on AMD CPUs")

        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 64
        cpu_data = torch.randn(model_size, device='cpu').to(dtype)
        param = torch.nn.Parameter(cpu_data)
        optimizer = DeepSpeedCPUAdam([param])

        subgroup_id = 0
        param.grad = torch.randn(model_size, device='cpu').to(dtype)

        # First, perform a step to initialize state
        optimizer.step_subgroup(subgroup_id)
        assert optimizer.state[subgroup_id]['step'] == 1

        # Store parameter state after step
        param_after_step = param.data.clone()
        exp_avg_after_step = optimizer.state[subgroup_id]['exp_avg'].clone()
        exp_avg_sq_after_step = optimizer.state[subgroup_id]['exp_avg_sq'].clone()

        # Now rollback
        optimizer.rollback_subgroup(subgroup_id)

        # Verify step count decremented
        assert optimizer.state[subgroup_id]['step'] == 0, "Step count should be decremented after rollback"

    def test_rollback_subgroup_uninitialized_error(self):
        """Test that rollback_subgroup raises error for uninitialized subgroup."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 64
        param = torch.nn.Parameter(torch.randn(model_size, device='cpu'))
        optimizer = DeepSpeedCPUAdam([param])

        # Try to rollback uninitialized subgroup
        with pytest.raises(RuntimeError, match="Cannot rollback optimizer state for sub_group_id 0"):
            optimizer.rollback_subgroup(0)

    def test_rollback_subgroup_zero_step_error(self):
        """Test that rollback_subgroup raises error when step count is already 0."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 64
        param = torch.nn.Parameter(torch.randn(model_size, device='cpu'))
        optimizer = DeepSpeedCPUAdam([param])

        subgroup_id = 0
        param.grad = torch.randn(model_size, device='cpu')

        # Initialize state by doing one step
        optimizer.step_subgroup(subgroup_id)

        # Rollback once (step should become 0)
        optimizer.rollback_subgroup(subgroup_id)
        assert optimizer.state[subgroup_id]['step'] == 0

        # Try to rollback again - should raise error
        with pytest.raises(RuntimeError, match="Cannot rollback sub_group_id 0: step count is 0"):
            optimizer.rollback_subgroup(subgroup_id)

    @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16], ids=["fp16", "bf16"])
    def test_step_rollback_sequence(self, dtype):
        """Test sequence of step_subgroup and rollback_subgroup operations."""
        if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
            pytest.skip("cpu-adam with half precision not supported on AMD CPUs")

        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 64
        cpu_data = torch.randn(model_size, device='cpu').to(dtype)
        param = torch.nn.Parameter(cpu_data)
        optimizer = DeepSpeedCPUAdam([param])

        subgroup_id = 0
        param.grad = torch.randn(model_size, device='cpu').to(dtype)

        # Perform multiple steps
        for step in range(1, 4):
            optimizer.step_subgroup(subgroup_id)
            assert optimizer.state[subgroup_id]['step'] == step

        # Rollback steps one by one
        for step in range(2, -1, -1):
            optimizer.rollback_subgroup(subgroup_id)
            assert optimizer.state[subgroup_id]['step'] == step

    def test_multiple_subgroups(self):
        """Test that different subgroups maintain independent state."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 64
        param = torch.nn.Parameter(torch.randn(model_size, device='cpu'))
        optimizer = DeepSpeedCPUAdam([param])

        param.grad = torch.randn(model_size, device='cpu')

        # Step different subgroups
        optimizer.step_subgroup(0)
        optimizer.step_subgroup(1)
        optimizer.step_subgroup(0)  # Step subgroup 0 again

        # Verify independent step counts
        assert optimizer.state[0]['step'] == 2, "Subgroup 0 should have step count 2"
        assert optimizer.state[1]['step'] == 1, "Subgroup 1 should have step count 1"

        # Rollback subgroup 0 only
        optimizer.rollback_subgroup(0)
        assert optimizer.state[0]['step'] == 1, "Subgroup 0 step count should be decremented"
        assert optimizer.state[1]['step'] == 1, "Subgroup 1 step count should be unchanged"

    def test_step_subgroup_same_step_idempotent_across_subgroups(self):
        """Repeated same-step subgroup updates should remain bit-identical."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 128
        steps = 4
        base = torch.randn(model_size, device='cpu', dtype=torch.float32)
        param_a = torch.nn.Parameter(base.clone())
        param_b = torch.nn.Parameter(base.clone())

        optimizer = DeepSpeedCPUAdam([param_a])
        for logical_step in range(1, steps + 1):
            grad = torch.randn(model_size, device='cpu', dtype=torch.float32)

            optimizer.param_groups[0]['params'] = [param_a]
            param_a.grad = grad.clone()
            optimizer.step_subgroup(0)

            optimizer.param_groups[0]['params'] = [param_b]
            param_b.grad = grad.clone()
            optimizer.step_subgroup(1)

            assert optimizer.state[0]['step'] == logical_step
            assert optimizer.state[1]['step'] == logical_step
            assert torch.equal(param_a.data, param_b.data)
            assert torch.equal(optimizer.state[0]['exp_avg'], optimizer.state[1]['exp_avg'])
            assert torch.equal(optimizer.state[0]['exp_avg_sq'], optimizer.state[1]['exp_avg_sq'])

    def test_step_same_step_idempotent_across_param_keys(self):
        """Repeated optimizer.step() with swapped param keys should be deterministic."""
        from deepspeed.ops.adam import DeepSpeedCPUAdam

        model_size = 128
        steps = 4
        base = torch.randn(model_size, device='cpu', dtype=torch.float32)
        param_a = torch.nn.Parameter(base.clone())
        param_b = torch.nn.Parameter(base.clone())

        optimizer = DeepSpeedCPUAdam([param_a])
        for logical_step in range(1, steps + 1):
            grad = torch.randn(model_size, device='cpu', dtype=torch.float32)

            optimizer.param_groups[0]['params'] = [param_a]
            param_a.grad = grad.clone()
            optimizer.step()

            optimizer.param_groups[0]['params'] = [param_b]
            param_b.grad = grad.clone()
            optimizer.step()

            assert optimizer.state[param_a]['step'] == logical_step
            assert optimizer.state[param_b]['step'] == logical_step
            assert torch.equal(param_a.data, param_b.data)
            assert torch.equal(optimizer.state[param_a]['exp_avg'], optimizer.state[param_b]['exp_avg'])
            assert torch.equal(optimizer.state[param_a]['exp_avg_sq'], optimizer.state[param_b]['exp_avg_sq'])