wehub/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
deepspeedai--deepspeed/tests/unit/ops/adam/test_cpu_adam.py
T
wehub-resource-sync 4ececc111a chore: import upstream snapshot with attribution
2026-07-13 13:18:33 +08:00

504 lines
22 KiB
Python
Raw Permalink Blame History

# 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()["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'])
Reference in New Issue View Git Blame Copy Permalink