A tutorial on pin_memory and non_blocking usage

[vmoens]

This is a draft PR on the proper usage of pin_memory and non_blocking when sending data from CPU to GPU (and GPU to CPU).

Description

There is some confusion about the proper use of pin_memory and non_blocking in the community. To be convinced about it, one can simply look for the .pin_memory.to( pattern on github (about 2k results), when this command is always slower than a simple call to to(device).

The responsibility lies in part in the pytorch doc itself where it is recommended (implicitly) to call pin_memory before calling to with non_blocking=True and several such occurrences on the pytorch forum too.

Some refs on the topic: https://discuss.pytorch.org/t/what-is-the-disadvantage-of-using-pin-memory/1702/13 https://discuss.pytorch.org/t/non-blocking-memory-transfer-to-gpu/188941 https://discuss.pytorch.org/t/should-we-set-non-blocking-to-true/38234 https://discuss.pytorch.org/t/how-is-explicit-pin-memory-different-from-just-calling-to-and-let-cuda-handle-it/197422

I use TensorDict to demonstrate how to use pin_memory across threads.

A follow-up will be to link back the pytorch doc / docstrings of to / pin_memory etc to this tutorial.

The syntax should be fixed and integrated better in the lib. Conclusion and additional resources are missing. Still, feedback is more than welcome!

The tutorial requires tensordict v0.5 which will be released in the coming days.

cc @shagunsodhani @albanD @janeyx99 @dstaay-fb @mikaylagawarecki @ptrblck

[pytorch-bot[bot]]

:link: Helpful Links

:test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/2983

• :page_facing_up: Preview Python docs built from this PR

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

:white_check_mark: No Failures

As of commit 07f9932da3694d7984ee490fce153e390dabbfa4 with merge base c3882dbd41ad26bd8b4cb009d282e911f7b1fa71 (): :green_heart: Looks good so far! There are no failures yet. :green_heart:

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

[vmoens]

@svekars Happy to get some more feedback. On my side, I only see these items to be done:

• edit figures to make them look prettier
• add additional info (links etc)
• check that we have the most demonstrative examples (ie play a bit with the setups to highlight the differences better)