1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556 |
- import argparse
- import torch
- import time
- import numpy as np
- import pytest
- import copy
- import deepspeed
- #TODO: not sure how we want to skip tests now that we're building at runtime
- #if not deepspeed.ops.__installed_ops__['cpu-adam']:
- # pytest.skip("cpu-adam is not installed", allow_module_level=True)
- #else:
- # from deepspeed.ops.adam import DeepSpeedCPUAdam
- def check_equal(first, second, atol=1e-2, verbose=False):
- x = first.detach().numpy()
- y = second.detach().numpy()
- if verbose:
- print("x = {}".format(x.flatten()))
- print("y = {}".format(y.flatten()))
- print('-' * 80)
- np.testing.assert_allclose(x, y, err_msg="param-update dismatch!", atol=atol)
- @pytest.mark.parametrize('model_size',
- [
- (64),
- (22),
- (55),
- (127),
- (1024),
- (1048576),
- ]) # yapf: disable
- def test_cpu_adam_opt(model_size):
- from deepspeed.ops.adam import DeepSpeedCPUAdam
- device = 'cpu'
- rng_state = torch.get_rng_state()
- param = torch.nn.Parameter(torch.randn(model_size, device=device))
- torch.set_rng_state(rng_state)
- param1 = torch.nn.Parameter(torch.randn(model_size, device=device))
- optimizer1 = torch.optim.Adam([param1])
- optimizer = DeepSpeedCPUAdam([param])
- for i in range(10):
- rng_state = torch.get_rng_state()
- param.grad = torch.randn(model_size, device=device)
- torch.set_rng_state(rng_state)
- param1.grad = torch.randn(model_size, device=device)
- optimizer.step()
- optimizer1.step()
- check_equal(param, param1, atol=1e-2, verbose=True)
|