[DataLoader] Select available CUDA or 3rd devices automatically to pin memory

[wizzniu]

This PR improves the behavior for DataLoader when pin_memory is True and pin_memory_device is not set. Instead of only checking whether CUDA is available or not, a more reasonable approach is that first check available CUDA device, and if not, get privateuse1 backend automatically to pin memory to. Otherwise, throw some warnings that pin memory doesn't take effect.

Fixes #124908

cc @ssnl @VitalyFedyunin @ejguan @dzhulgakov @ezyang

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ No Failures

As of commit 144a6e0 with merge base 769b1e6 ():
💚 Looks good so far! There are no failures yet. 💚

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

[linux-foundation-easycla]

• ✅login: wizzniu / (144a6e0)

The committers listed above are authorized under a signed CLA.

[vadimkantorov]

Would it be better to explicitly accept a pin_memory_device to a DataLoader constructor? and fix things light PyTorch Lightning / HF Trainer to pass it in (instead of mixing in the data loader device detection)?

[wizzniu]

Would it be better to explicitly accept a pin_memory_device to a DataLoader constructor? and fix things light PyTorch Lightning / HF Trainer to pass it in (instead of mixing in the data loader device detection)?

This PR maintains BC in the parameters pin_memory and pin_memory_device of API. You still can explicitly pass in a pin_memory_device to it.

[wizzniu]

@ezyang @albanD @ejguan Can you help to review this PR?

[ezyang]

I'm deferring to @albanD for custom backend stuff

[wizzniu]

I'm deferring to @albanD for custom backend stuff

Thanks for your reply. Could you give some suggestions, @albanD ?

[albanD]

Yes sorry, I wanted to give a longer answer on this but here is the short version:

• We do not want for every API that takes pin_memory: bool flag to grow a pin_memory_device argument.
• We want to use the newly introduced concept of accelerator at

  pytorch/aten/src/ATen/DeviceAccelerator.h

  Line 25 in 196a0b1

  TORCH_API std::optional<c10::DeviceType> getAccelerator(bool checked = false);

  to redefine pin memory as: "Pin memory is always pinned for the current accelerator device".
• This will allow us to have all the other factory function APIs work without the need for new arguments.

What it means for this PR is a bit more subtle:

• Long term, we most likely just want to drop the pin_memory_device argument altogether
• This will require updating the implementation of the pin memory concept for low level APIs
• If you have time, it would be best to fix the low level constructs directly. If you do not, I'm happy to have a quick fix here that consists of binding the "getAccelerator()" API to python and use that here to set the default value of pin_memory_device. This is not perfect but a step in the direction where we want to go.

WDYT?

[mengpenghui]

@albanD
I am working on IPC-related interfaces to support third-party devices. The suggestions you gave in this PR are very helpful for modifying IPC. Now I have a new plan:

1. The share_cuda_ interface is changed to share_device_
2. The backend uses getAccelerator to determine the device and distribute different devices.
3. Add the share_device implementation of third-party devices to PrivateUse1HooksInterface

I plan to modify the new solution on this PR: #125122
Do you think this is feasible? Looking forward to your reply.

[wizzniu]

Yes sorry, I wanted to give a longer answer on this but here is the short version:

• We do not want for every API that takes pin_memory: bool flag to grow a pin_memory_device argument.
• We want to use the newly introduced concept of accelerator at

  pytorch/aten/src/ATen/DeviceAccelerator.h

  Line 25 in 196a0b1

  TORCH_API std::optional<c10::DeviceType> getAccelerator(bool checked = false);

  
  to redefine pin memory as: "Pin memory is always pinned for the current accelerator device".
• This will allow us to have all the other factory function APIs work without the need for new arguments.

What it means for this PR is a bit more subtle:

• Long term, we most likely just want to drop the pin_memory_device argument altogether
• This will require updating the implementation of the pin memory concept for low level APIs
• If you have time, it would be best to fix the low level constructs directly. If you do not, I'm happy to have a quick fix here that consists of binding the "getAccelerator()" API to python and use that here to set the default value of pin_memory_device. This is not perfect but a step in the direction where we want to go.

WDYT?

I agree with you that additional device argument is unnecessary when using pin_memory related APIs, which requires updating the implementation of the pin memory concept for low level APIs. And I am happy to fix the low level constructs directly.
So here I want to confirm whether our goal is to get rid of device argument of all pin_memory related APIs(e.g. Tensor.pin_memory())? And if it is, I will take time to work on this.

[albanD]

So here I want to confirm whether our goal is to get rid of device argument of all pin_memory related APIs(e.g. Tensor.pin_memory())?

Yes! But also note that most of these APIs do NOT have a device argument today. So not much to remove for most of them haha
The main thing is that we want to enable any accelerator pinned memory but we don't want to introduce this device argument everywhere. So changing the low level constructs to leverage the Accelerator concept is the way to go.
Let me know if you think this needs more clarification!

[wizzniu]

@albanD hi, I have updated the implement of pin memory in a new PR, please see #126376