Question: Can DeepSpeedCPUAdam be used as a drop in replacement to torch.optim.Adam?

[ofirzaf]

Hi,

I want to use the DeepSpeedCPUAdam instead of torch.optim.Adam for reducing the RAM usage of my GPUs while training.
I was wondering if DeepSpeedCPUAdam can be just dropped in instead of torch.optim.Adam or additional steps are needed?
I tried to do exactly that and I got a segmentation fault

Thanks

[tjruwase]

@ofirzaf Thanks for using DeepSpeed.

Yes, you can use DeepSpeedCPUAdam as drop-in replacement for torch.optim.Adam, that will give up to a 6X throughput boost. Let us know of any issues.

[ofirzaf]

I tried to make it work with mnist example pytorch published: https://pastebin.pl/view/embed/d3921ab0
When I try to run it I get:

Adam Optimizer #0 is created with AVX2 arithmetic capability.
group 0 param 0 = 288
Segmentation fault (core dumped)

I am running with pytorch 1.6, cuda 10.1, deepspeed installed from master and v0.3 using DS_BUILD_CPU_ADAM=1 install.sh
Also tried to use your docker image to check if the problem is in my installation but the image doesn't have cpu adam installed, got ModuleNotFoundError on cpu_adam_op

[RezaYazdaniAminabadi]

Hi @ofirzaf

It seems that CPUAdam is instantiated correctly. However, it has crashed when trying to access the parameters on CPU memory. I see in your code, that the model is converted in line 124, and can be sent over to the gpu memory if the use_cuda flag is set in line 98. CPUAdam requires the parameter (master copy in FP32) and the optimizer states such as variance and momentum reside on CPU memory. Could you please check if such requirement is met in your test?

Thanks,
Reza

[ofirzaf]

I initialize the optimizer while the model's parameters which I want to optimize are passed to the GPU when I move the entire model to the GPU. I am not sure how should I keep the parameters and other stats on CPU while still training on GPU.
Maybe I misunderstood and this optimizer should only be used when training using CPU only?

[RezaYazdaniAminabadi]

Hi @ofirzaf ,

In this case, you probably want to run your code through DeepSpeed and turning on the ZeRO-Offload feature. You can find more information on ZeRO-offload here: https://www.deepspeed.ai/tutorials/zero-offload/. Also, to get started with DeepSpeed, please check out this tutorial: https://www.deepspeed.ai/getting-started/
Thanks.
Reza

[tjruwase]

@ofirzaf, yes you are right, DeepSpeedCPUAdam is only used for running optimizer on CPU only. I apologize that my initial response was confusing because I did not clarify that while torch.optim.Adam works correctly on both GPU and CPU, DeepSpeedCPUAdam only works on CPU. So it is only on CPU execution that DeepSpeedCPUAdam is a drop in replacement for torch.optim.Adam.

[tjruwase]

Closing for lack of activity. Please reopen as needed.

[peterukk]

So DeepSpeedCPUAdam should work without CUDA?

  File "...python3.9/site-packages/deepspeed/ops/op_builder/builder.py", line 37, in installed_cuda_version
    assert cuda_home is not None, "CUDA_HOME does not exist, unable to compile CUDA op(s)"

[tjruwase]

@peterukk, DeepSpeedCPUAdam will not work without CUDA (in theory it could). The reason is that DeepSpeedCPUAdam has a mode of execution where it also copies the updated parameters back to GPU using CUDA kernels. Do you have scenario where you want to use DeepSpeedCPUAdam outside CUDA environment? Depending on your answer can you please open a Question or reopen this one as appropriate? Thanks.

[peterukk]

@tjruwase thanks, I opened a new issue