From 1dcaf249bd35b1d9b17374683fe0e95d752c019a Mon Sep 17 00:00:00 2001 From: ppt0011 Date: Wed, 11 Oct 2023 17:48:21 +0800 Subject: [PATCH] [doc] add reminder for issue encountered with hybrid adam --- colossalai/nn/optimizer/cpu_adam.py | 1 + 1 file changed, 1 insertion(+) diff --git a/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py index f35dc0200..3dc729c32 100644 --- a/colossalai/nn/optimizer/cpu_adam.py +++ b/colossalai/nn/optimizer/cpu_adam.py @@ -77,6 +77,7 @@ class CPUAdam(NVMeOptimizer): 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(