Function check_adam_kernel

tests/test_optimizer/test_adam_kernel.py:117–147

(
    kernel: Type[AdamKernel],
    adamw: bool,
    weight_decay: float,
    p_dtype: torch.dtype,
    g_dtype: torch.dtype,
    device: torch.device,
    n_steps: int,
    rtol: float,
    atol: float,
)

def check_adam_kernel(
    kernel: Type[AdamKernel],
    adamw: bool,
    weight_decay: float,
    p_dtype: torch.dtype,
    g_dtype: torch.dtype,
    device: torch.device,
    n_steps: int,
    rtol: float,
    atol: float,
):
    lr = 1e-3
    beta1, beta2 = 0.9, 0.999
    eps = 1e-8
    torch_adam = TorchAdamKernel(lr, beta1, beta2, eps, weight_decay, adamw)
    adam_kernel = kernel(lr, beta1, beta2, eps, weight_decay, adamw)
    master_p = torch.rand(64, device=device)
    master_g = torch.rand_like(master_p)
    master_exp_avg = torch.zeros_like(master_p)
    master_exp_avg_sq = torch.zeros_like(master_p)
    p = master_p.clone().to(p_dtype)
    g = master_g.clone().to(g_dtype)
    exp_avg = master_exp_avg.clone().to(p_dtype)
    exp_avg_sq = master_exp_avg_sq.clone().to(p_dtype)

    for step in range(1, 1 + n_steps):
        torch_adam.update(step, master_p, master_g, master_exp_avg, master_exp_avg_sq)
        adam_kernel.update(step, p, g, exp_avg, exp_avg_sq)
        # if overflow, the weight won't be updated. so there will be no nan in p
        assert not torch.isnan(p).any()
        assert torch.allclose(master_p, p.float(), rtol=rtol, atol=atol)


@pytest.mark.parametrize("adamw", [False, True])