code: 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
          )
  
imports: import math
  from abc import abstractmethod
  from typing import Type
  import pytest
  import torch
  from torch import Tensor
  from colossalai.utils import get_current_device, multi_tensor_applier
label_desc: Use image node_id 1 for calling a global function with example usage: check_adam_kernel(kernel, adamw, weight_decay, p_dtype, g_dtype, device, n_steps, rtol, atol) without return types