Skip to content

GitLab

"python-wheel/rust/lib.rs" did not exist on "a62a862745e0fcf734d9fb33751c4cb0171f5bf9"
Unverified Commit c1fab951 authored by ppt0011's avatar ppt0011 Committed by GitHub
Browse files

Merge pull request #4889 from ppt0011/main 

[doc] add reminder for issue encountered with hybrid adam
parents ffd9a3cb 1dcaf249
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
colossalai/nn/optimizer/cpu_adam.py
View file @ c1fab951
...@@ -77,6 +77,7 @@ class CPUAdam(NVMeOptimizer): ...@@ -77,6 +77,7 @@ class CPUAdam(NVMeOptimizer):
super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir) super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)
self.adamw_mode = adamw_mode self.adamw_mode = adamw_mode
cpu_adam = CPUAdamBuilder().load() 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) self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)
def torch_adam_update( def torch_adam_update(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment