Enable CPU Offload for Intel GPU

[dbyoung18]

Background

Current CPU Offload in torchao only supports CUDA backend. We would like to add support for Intel GPU with the device option "xpu".

Details

• add "device" attribute to CPUOffloadOptimizer, default setting to "cuda"
• enhance and verify UT test_optim_cpu_offload_correctness & test_optim_cpu_offload_save_load pass on Intel GPU
• add "device" argument to benchmark_low_bit_adam.py. Users can use "--device xpu" to benchmark CPU Offload on Intel GPU. Currently it supports both full BF16 and BF16 AMP training w/ eager and compiled mode. Verified workloads on Intel GPU achieve memory saving and interleaving as expected as the description in reference PR:ao#584

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1324

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 03ac00f with merge base 478d15b ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[gau-nernst]

Thanks for the feature addition! Hopefully once the device-agnostic API support arrives, we can eliminate the if-else checks 😆

[gau-nernst]

@dbyoung18 Can you run ruff format and push the formatted code? CUDA nightly is failing because of bitsandbytes calling triton.ops (I think later versions of triton doesn't have triton.ops anymore bitsandbytes-foundation/bitsandbytes#1413). It's not related but not sure if we can merge until that is fixed 😢. I think other PRs will be affected too.

Otherwise, everything else looks good already!

[dbyoung18]

@dbyoung18 Can you run ruff format and push the formatted code? CUDA nightly is failing because of bitsandbytes calling triton.ops (I think later versions of triton doesn't have triton.ops anymore bitsandbytes-foundation/bitsandbytes#1413). It's not related but not sure if we can merge until that is fixed 😢. I think other PRs will be affected too.

Otherwise, everything else looks good already!

Done for ruff format. Hopes the bnb issue could be resolved soon. THX again for ur review and quick feedback:)

[gau-nernst]

@dbyoung18 Can you merge from main? #1343 should fix the bnb issue.

Also, can you also update the doc here? https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload

After that we are good to merge 😃

[dbyoung18]

@dbyoung18 Can you merge from main? #1343 should fix the bnb issue.

Also, can you also update the doc here? https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload

After that we are good to merge 😃

Done for both. We have a plan to gradually support torch-ao & pytorch core on Intel GPU. For this PR it covers CPU Offload only and I will look into the remain part of low-bit optimizers for next step. Since meanwhile we are also on the way to upstream FlashAttention backend to pytorch core(target v2.6 or v2.7), would like to add benchmark data to the README when it's ready. So currently, I only modify the README to make the CPU-Offload part to cover XPU scope. THX for review and I am also looking forward to make further contributions soon.😃

[gau-nernst]

Sounds good! The low-bit optimizers rely entirely on the tensor subclass + torch.compile() stack, so as long as there is a triton build that supports XPU backend, it should work out-of-the-box!