optimizer CPU offload doesn't work outside of CUDA

[bghira]

The CPUOptimizerOffload class is very clever, but overly relies on CUDA Streams, which aren't available w/o a CUDA device.

should use torch.cpu.Stream and torch.cpu.current_stream instead.

additionally, pin_memory=True if torch.cuda.is_available() else False as MPS is a unified mem arch.

[gau-nernst]

Yes, CPU offload optimizer only works on CUDA. There is no point to support other devices. Perhaps we should make it clearer 🤔

• If you train with CPU, well, you don't need CPU offload
• If you train with MPS, since it is unified memory, again, you don't need CPU offload.

I'm not too familiar with other hardware, so not sure how this works with AMD and Intel GPUs.

[bghira]

pytorch isn't taking much advantage of MPS in a unified memory way, and, i agree, there isn't a whole lot of point to it on MPS other than to be able to run the same code and ensure a consistent experience.

for ROCm, i'm almost certain it just masquerades as CUDA and is invisible. the Intel and other systems like Ascend NPU rely on XPU or NPU extensions in pytorch, and TPUs require XLA.

[msaroufim]

Yeah agree with @gau-nernst I'd just maybe prioritize sanity checking whether the offloader works on an AMD GPU

[petrex]

This (cpu offload on AMD) is on our roadmap, stay tuned.

[gau-nernst]

@petrex Sounds awesome! Actually our single-GPU CPU-offload optimizer only relies on PyTorch's CUDA stream API, so as long as PyTorch ROCm supports those API, it should work on AMD GPUs too! I don't have access to AMD GPUs so I can't test it out though.

https://github.com/pytorch/ao/blob/v0.6.1/torchao/prototype/low_bit_optim/cpu_offload.py