We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/nn/optimizer/cpu_adam.py`
Content:
```
1 import math
2 from typing import Optional
3 
4 import torch
5 
6 from colossalai.kernel.op_builder import CPUAdamBuilder
7 
8 from .nvme_optimizer import NVMeOptimizer
9 
10 
11 class CPUAdam(NVMeOptimizer):
12     """Implements Adam algorithm.
13 
14     Supports parameters updating on both GPU and CPU, depending on the device of parameters.
15     But the parameters and gradients should on the same device:
16       * Parameters on CPU and gradients on CPU is allowed.
17       * Parameters on GPU and gradients on GPU is allowed.
18       * Parameters on GPU and gradients on CPU is **not** allowed.
19 
20     `CPUAdam` requires CUDA extensions which can be built during installation or runtime.
21 
22     This version of CPU Adam accelerates parameters updating on CPU with SIMD.
23     Support of AVX2 or AVX512 is required.
24 
25     The GPU part is implemented in an naive way.
26 
27     CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.
28 
29     :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,
30     or ``torch.optim.Adam`` with ``adamw_mode=False``
31 
32     Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.
33 
34     Arguments:
35         model_params (iterable): iterable of parameters of dicts defining
36             parameter groups.
37         lr (float, optional): learning rate. (default: 1e-3)
38         betas (Tuple[float, float], optional): coefficients used for computing
39             running averages of gradient and its square. (default: (0.9, 0.999))
40         eps (float, optional): term added to the denominator to improve
41             numerical stability. (default: 1e-8)
42         weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
43         amsgrad (boolean, optional): whether to use the AMSGrad variant of this
44             algorithm from the paper `On the Convergence of Adam and Beyond`_
45             (default: False) NOT SUPPORTED yet in CPUAdam!
46         adamw_mode (boolean, optional): Apply L2 regularization or weight decay
47             True for decoupled weight decay(also known as AdamW) (default: True)
48         simd_log (boolean, optional): whether to show if you are using SIMD to
49             accelerate. (default: False)
50         nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.
51         nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.
52             If it's ``None``, a random temporary directory will be used. Defaults to None.
53 
54     .. _Adam\: A Method for Stochastic Optimization:
55         https://arxiv.org/abs/1412.6980
56     .. _On the Convergence of Adam and Beyond:
57         https://openreview.net/forum?id=ryQu7f-RZ
58     """
59 
60     # Number of fp32 shards for per parameter
61     # Param weight, grad, momentum and variance
62     num_fp32_shards_per_param = 4
63 
64     def __init__(
65         self,
66         model_params,
67         lr=1e-3,
68         bias_correction=True,
69         betas=(0.9, 0.999),
70         eps=1e-8,
71         weight_decay=0,
72         adamw_mode=True,
73         nvme_offload_fraction: float = 0.0,
74         nvme_offload_dir: Optional[str] = None,
75     ):
76         default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction)
77         super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)
78         self.adamw_mode = adamw_mode
79         cpu_adam = CPUAdamBuilder().load()
80         self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)
81 
82     def torch_adam_update(
83         self,
84         data,
85         grad,
86         exp_avg,
87         exp_avg_sq,
88         lr,
89         beta1,
90         beta2,
91         eps,
92         weight_decay,
93         bias_correction1,
94         bias_correction2,
95         use_adamw=False,
96     ):
97         grad = grad.to(data.dtype)
98 
99         if weight_decay != 0:
100             if use_adamw:
101                 data.mul_(1 - lr * weight_decay)
102             else:
103                 grad = grad.add(data, alpha=weight_decay)
104 
105         # Decay the first and second moment running average coefficient
106         exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
107         exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
108 
109         # TODO(jiaruifang) dose not support amsgrad
110         denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
111 
112         step_size = lr / bias_correction1
113 
114         data.addcdiv_(exp_avg, denom, value=-step_size)
115 
116     @torch.no_grad()
117     def step(self, closure=None, div_scale: float = -1):
118         loss = None
119         if closure is not None:
120             with torch.enable_grad():
121                 loss = closure()
122 
123         self._pre_step("exp_avg", "exp_avg_sq")
124         for _, group in enumerate(self.param_groups):
125             for _, p in enumerate(group["params"]):
126                 if p.grad is None:
127                     continue
128 
129                 state = self.state[p]
130 
131                 target_device = p.device
132                 if len(state) == 0:
133                     state["step"] = 0
134 
135                     # FIXME(ver217): CPU adam kernel only supports fp32 states now
136                     assert p.dtype is torch.float, "CPUAdam only support fp32 parameters"
137                     # gradient momentums
138                     state["exp_avg"] = torch.zeros_like(p, device=target_device)
139                     # gradient variances
140                     state["exp_avg_sq"] = torch.zeros_like(p, device=target_device)
141                     self._post_state_init(p)
142 
143                 state["step"] += 1
144                 beta1, beta2 = group["betas"]
145 
146                 if target_device.type == "cpu":
147                     assert p.data.numel() == p.grad.data.numel(), "parameter and gradient should have the same size"
148                     assert state["exp_avg"].device.type == "cpu", "exp_avg should stay on cpu"
149                     assert state["exp_avg_sq"].device.type == "cpu", "exp_avg should stay on cpu"
150                     self._pre_update(p, "exp_avg", "exp_avg_sq")
151                     if p.grad.dtype is torch.bfloat16:
152                         # cpu adam kernel does not support bf16 now
153                         bias_correction1 = 1 - beta1 ** state["step"]
154                         bias_correction2 = 1 - beta2 ** state["step"]
155                         self.torch_adam_update(
156                             p.data,
157                             p.grad.data,
158                             state["exp_avg"],
159                             state["exp_avg_sq"],
160                             group["lr"],
161                             beta1,
162                             beta2,
163                             group["eps"],
164                             group["weight_decay"],
165                             bias_correction1,
166                             bias_correction2,
167                             self.adamw_mode,
168                         )
169                     else:
170                         self.cpu_adam_op.step(
171                             state["step"],
172                             group["lr"],
173                             beta1,
174                             beta2,
175                             group["eps"],
176                             group["weight_decay"],
177                             group["bias_correction"],
178                             p.data,
179                             p.grad.data,
180                             state["exp_avg"],
181                             state["exp_avg_sq"],
182                             div_scale,
183                         )
184                     self._post_update(p, "exp_avg", "exp_avg_sq")
185                 elif target_device.type == "cuda":
186                     assert div_scale == -1, "div_scale should remain default"
187                     assert state["exp_avg"].device.type == "cuda", "exp_avg should stay on cuda"
188                     assert state["exp_avg_sq"].device.type == "cuda", "exp_avg should stay on cuda"
189 
190                     bias_correction1 = 1 - beta1 ** state["step"]
191                     bias_correction2 = 1 - beta2 ** state["step"]
192 
193                     # adam on cuda
194                     self.torch_adam_update(
195                         p.data,
196                         p.grad.data,
197                         state["exp_avg"],
198                         state["exp_avg_sq"],
199                         group["lr"],
200                         beta1,
201                         beta2,
202                         group["eps"],
203                         group["weight_decay"],
204                         bias_correction1,
205                         bias_correction2,
206                         self.adamw_mode,
207                     )
208                 else:
209                     raise RuntimeError
210         self._post_step()
211         return loss
212
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py
--- a/colossalai/nn/optimizer/cpu_adam.py
+++ b/colossalai/nn/optimizer/cpu_adam.py
@@ -77,6 +77,7 @@
         super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)
         self.adamw_mode = adamw_mode
         cpu_adam = CPUAdamBuilder().load()
+        # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification
         self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)
 
     def torch_adam_update(

{"golden_diff": "diff --git a/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py\n--- a/colossalai/nn/optimizer/cpu_adam.py\n+++ b/colossalai/nn/optimizer/cpu_adam.py\n@@ -77,6 +77,7 @@\n         super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)\n         self.adamw_mode = adamw_mode\n         cpu_adam = CPUAdamBuilder().load()\n+        # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification\n         self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)\n \n     def torch_adam_update(\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import math\nfrom typing import Optional\n\nimport torch\n\nfrom colossalai.kernel.op_builder import CPUAdamBuilder\n\nfrom .nvme_optimizer import NVMeOptimizer\n\n\nclass CPUAdam(NVMeOptimizer):\n    \"\"\"Implements Adam algorithm.\n\n    Supports parameters updating on both GPU and CPU, depending on the device of parameters.\n    But the parameters and gradients should on the same device:\n      * Parameters on CPU and gradients on CPU is allowed.\n      * Parameters on GPU and gradients on GPU is allowed.\n      * Parameters on GPU and gradients on CPU is **not** allowed.\n\n    `CPUAdam` requires CUDA extensions which can be built during installation or runtime.\n\n    This version of CPU Adam accelerates parameters updating on CPU with SIMD.\n    Support of AVX2 or AVX512 is required.\n\n    The GPU part is implemented in an naive way.\n\n    CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.\n\n    :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,\n    or ``torch.optim.Adam`` with ``adamw_mode=False``\n\n    Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.\n\n    Arguments:\n        model_params (iterable): iterable of parameters of dicts defining\n            parameter groups.\n        lr (float, optional): learning rate. (default: 1e-3)\n        betas (Tuple[float, float], optional): coefficients used for computing\n            running averages of gradient and its square. (default: (0.9, 0.999))\n        eps (float, optional): term added to the denominator to improve\n            numerical stability. (default: 1e-8)\n        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n            algorithm from the paper `On the Convergence of Adam and Beyond`_\n            (default: False) NOT SUPPORTED yet in CPUAdam!\n        adamw_mode (boolean, optional): Apply L2 regularization or weight decay\n            True for decoupled weight decay(also known as AdamW) (default: True)\n        simd_log (boolean, optional): whether to show if you are using SIMD to\n            accelerate. (default: False)\n        nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.\n        nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.\n            If it's ``None``, a random temporary directory will be used. Defaults to None.\n\n    .. _Adam\\: A Method for Stochastic Optimization:\n        https://arxiv.org/abs/1412.6980\n    .. _On the Convergence of Adam and Beyond:\n        https://openreview.net/forum?id=ryQu7f-RZ\n    \"\"\"\n\n    # Number of fp32 shards for per parameter\n    # Param weight, grad, momentum and variance\n    num_fp32_shards_per_param = 4\n\n    def __init__(\n        self,\n        model_params,\n        lr=1e-3,\n        bias_correction=True,\n        betas=(0.9, 0.999),\n        eps=1e-8,\n        weight_decay=0,\n        adamw_mode=True,\n        nvme_offload_fraction: float = 0.0,\n        nvme_offload_dir: Optional[str] = None,\n    ):\n        default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction)\n        super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)\n        self.adamw_mode = adamw_mode\n        cpu_adam = CPUAdamBuilder().load()\n        self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)\n\n    def torch_adam_update(\n        self,\n        data,\n        grad,\n        exp_avg,\n        exp_avg_sq,\n        lr,\n        beta1,\n        beta2,\n        eps,\n        weight_decay,\n        bias_correction1,\n        bias_correction2,\n        use_adamw=False,\n    ):\n        grad = grad.to(data.dtype)\n\n        if weight_decay != 0:\n            if use_adamw:\n                data.mul_(1 - lr * weight_decay)\n            else:\n                grad = grad.add(data, alpha=weight_decay)\n\n        # Decay the first and second moment running average coefficient\n        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n\n        # TODO(jiaruifang) dose not support amsgrad\n        denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n\n        step_size = lr / bias_correction1\n\n        data.addcdiv_(exp_avg, denom, value=-step_size)\n\n    @torch.no_grad()\n    def step(self, closure=None, div_scale: float = -1):\n        loss = None\n        if closure is not None:\n            with torch.enable_grad():\n                loss = closure()\n\n        self._pre_step(\"exp_avg\", \"exp_avg_sq\")\n        for _, group in enumerate(self.param_groups):\n            for _, p in enumerate(group[\"params\"]):\n                if p.grad is None:\n                    continue\n\n                state = self.state[p]\n\n                target_device = p.device\n                if len(state) == 0:\n                    state[\"step\"] = 0\n\n                    # FIXME(ver217): CPU adam kernel only supports fp32 states now\n                    assert p.dtype is torch.float, \"CPUAdam only support fp32 parameters\"\n                    # gradient momentums\n                    state[\"exp_avg\"] = torch.zeros_like(p, device=target_device)\n                    # gradient variances\n                    state[\"exp_avg_sq\"] = torch.zeros_like(p, device=target_device)\n                    self._post_state_init(p)\n\n                state[\"step\"] += 1\n                beta1, beta2 = group[\"betas\"]\n\n                if target_device.type == \"cpu\":\n                    assert p.data.numel() == p.grad.data.numel(), \"parameter and gradient should have the same size\"\n                    assert state[\"exp_avg\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    self._pre_update(p, \"exp_avg\", \"exp_avg_sq\")\n                    if p.grad.dtype is torch.bfloat16:\n                        # cpu adam kernel does not support bf16 now\n                        bias_correction1 = 1 - beta1 ** state[\"step\"]\n                        bias_correction2 = 1 - beta2 ** state[\"step\"]\n                        self.torch_adam_update(\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            bias_correction1,\n                            bias_correction2,\n                            self.adamw_mode,\n                        )\n                    else:\n                        self.cpu_adam_op.step(\n                            state[\"step\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            group[\"bias_correction\"],\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            div_scale,\n                        )\n                    self._post_update(p, \"exp_avg\", \"exp_avg_sq\")\n                elif target_device.type == \"cuda\":\n                    assert div_scale == -1, \"div_scale should remain default\"\n                    assert state[\"exp_avg\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n\n                    bias_correction1 = 1 - beta1 ** state[\"step\"]\n                    bias_correction2 = 1 - beta2 ** state[\"step\"]\n\n                    # adam on cuda\n                    self.torch_adam_update(\n                        p.data,\n                        p.grad.data,\n                        state[\"exp_avg\"],\n                        state[\"exp_avg_sq\"],\n                        group[\"lr\"],\n                        beta1,\n                        beta2,\n                        group[\"eps\"],\n                        group[\"weight_decay\"],\n                        bias_correction1,\n                        bias_correction2,\n                        self.adamw_mode,\n                    )\n                else:\n                    raise RuntimeError\n        self._post_step()\n        return loss\n", "path": "colossalai/nn/optimizer/cpu_adam.py"}], "after_files": [{"content": "import math\nfrom typing import Optional\n\nimport torch\n\nfrom colossalai.kernel.op_builder import CPUAdamBuilder\n\nfrom .nvme_optimizer import NVMeOptimizer\n\n\nclass CPUAdam(NVMeOptimizer):\n    \"\"\"Implements Adam algorithm.\n\n    Supports parameters updating on both GPU and CPU, depending on the device of parameters.\n    But the parameters and gradients should on the same device:\n      * Parameters on CPU and gradients on CPU is allowed.\n      * Parameters on GPU and gradients on GPU is allowed.\n      * Parameters on GPU and gradients on CPU is **not** allowed.\n\n    `CPUAdam` requires CUDA extensions which can be built during installation or runtime.\n\n    This version of CPU Adam accelerates parameters updating on CPU with SIMD.\n    Support of AVX2 or AVX512 is required.\n\n    The GPU part is implemented in an naive way.\n\n    CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.\n\n    :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,\n    or ``torch.optim.Adam`` with ``adamw_mode=False``\n\n    Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.\n\n    Arguments:\n        model_params (iterable): iterable of parameters of dicts defining\n            parameter groups.\n        lr (float, optional): learning rate. (default: 1e-3)\n        betas (Tuple[float, float], optional): coefficients used for computing\n            running averages of gradient and its square. (default: (0.9, 0.999))\n        eps (float, optional): term added to the denominator to improve\n            numerical stability. (default: 1e-8)\n        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n            algorithm from the paper `On the Convergence of Adam and Beyond`_\n            (default: False) NOT SUPPORTED yet in CPUAdam!\n        adamw_mode (boolean, optional): Apply L2 regularization or weight decay\n            True for decoupled weight decay(also known as AdamW) (default: True)\n        simd_log (boolean, optional): whether to show if you are using SIMD to\n            accelerate. (default: False)\n        nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.\n        nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.\n            If it's ``None``, a random temporary directory will be used. Defaults to None.\n\n    .. _Adam\\: A Method for Stochastic Optimization:\n        https://arxiv.org/abs/1412.6980\n    .. _On the Convergence of Adam and Beyond:\n        https://openreview.net/forum?id=ryQu7f-RZ\n    \"\"\"\n\n    # Number of fp32 shards for per parameter\n    # Param weight, grad, momentum and variance\n    num_fp32_shards_per_param = 4\n\n    def __init__(\n        self,\n        model_params,\n        lr=1e-3,\n        bias_correction=True,\n        betas=(0.9, 0.999),\n        eps=1e-8,\n        weight_decay=0,\n        adamw_mode=True,\n        nvme_offload_fraction: float = 0.0,\n        nvme_offload_dir: Optional[str] = None,\n    ):\n        default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction)\n        super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)\n        self.adamw_mode = adamw_mode\n        cpu_adam = CPUAdamBuilder().load()\n        # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification\n        self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)\n\n    def torch_adam_update(\n        self,\n        data,\n        grad,\n        exp_avg,\n        exp_avg_sq,\n        lr,\n        beta1,\n        beta2,\n        eps,\n        weight_decay,\n        bias_correction1,\n        bias_correction2,\n        use_adamw=False,\n    ):\n        grad = grad.to(data.dtype)\n\n        if weight_decay != 0:\n            if use_adamw:\n                data.mul_(1 - lr * weight_decay)\n            else:\n                grad = grad.add(data, alpha=weight_decay)\n\n        # Decay the first and second moment running average coefficient\n        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n\n        # TODO(jiaruifang) dose not support amsgrad\n        denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n\n        step_size = lr / bias_correction1\n\n        data.addcdiv_(exp_avg, denom, value=-step_size)\n\n    @torch.no_grad()\n    def step(self, closure=None, div_scale: float = -1):\n        loss = None\n        if closure is not None:\n            with torch.enable_grad():\n                loss = closure()\n\n        self._pre_step(\"exp_avg\", \"exp_avg_sq\")\n        for _, group in enumerate(self.param_groups):\n            for _, p in enumerate(group[\"params\"]):\n                if p.grad is None:\n                    continue\n\n                state = self.state[p]\n\n                target_device = p.device\n                if len(state) == 0:\n                    state[\"step\"] = 0\n\n                    # FIXME(ver217): CPU adam kernel only supports fp32 states now\n                    assert p.dtype is torch.float, \"CPUAdam only support fp32 parameters\"\n                    # gradient momentums\n                    state[\"exp_avg\"] = torch.zeros_like(p, device=target_device)\n                    # gradient variances\n                    state[\"exp_avg_sq\"] = torch.zeros_like(p, device=target_device)\n                    self._post_state_init(p)\n\n                state[\"step\"] += 1\n                beta1, beta2 = group[\"betas\"]\n\n                if target_device.type == \"cpu\":\n                    assert p.data.numel() == p.grad.data.numel(), \"parameter and gradient should have the same size\"\n                    assert state[\"exp_avg\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    self._pre_update(p, \"exp_avg\", \"exp_avg_sq\")\n                    if p.grad.dtype is torch.bfloat16:\n                        # cpu adam kernel does not support bf16 now\n                        bias_correction1 = 1 - beta1 ** state[\"step\"]\n                        bias_correction2 = 1 - beta2 ** state[\"step\"]\n                        self.torch_adam_update(\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            bias_correction1,\n                            bias_correction2,\n                            self.adamw_mode,\n                        )\n                    else:\n                        self.cpu_adam_op.step(\n                            state[\"step\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            group[\"bias_correction\"],\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            div_scale,\n                        )\n                    self._post_update(p, \"exp_avg\", \"exp_avg_sq\")\n                elif target_device.type == \"cuda\":\n                    assert div_scale == -1, \"div_scale should remain default\"\n                    assert state[\"exp_avg\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n\n                    bias_correction1 = 1 - beta1 ** state[\"step\"]\n                    bias_correction2 = 1 - beta2 ** state[\"step\"]\n\n                    # adam on cuda\n                    self.torch_adam_update(\n                        p.data,\n                        p.grad.data,\n                        state[\"exp_avg\"],\n                        state[\"exp_avg_sq\"],\n                        group[\"lr\"],\n                        beta1,\n                        beta2,\n                        group[\"eps\"],\n                        group[\"weight_decay\"],\n                        bias_correction1,\n                        bias_correction2,\n                        self.adamw_mode,\n                    )\n                else:\n                    raise RuntimeError\n        self._post_step()\n        return loss\n", "path": "colossalai/nn/optimizer/cpu_adam.py"}]}