Set device in dataloaders

[edknv]

Fixes #651

Goals ⚽

Fix multi-gpu training notebook.

Implementation Details 🚧

• Depends on Put row lengths on the same device on gpu dataloader#113.
• device is set identical to local_rank.
• Without dropping the last batch (dataloader_drop_last=True), recsys_trainer.evaluate hangs. Probably need to investigate this because this didn't happen before the list column refactoring in merlin-dataloader (see ticket).
• torch.distributed.launch is replaced with torchrun because the former has been deprecated.

Testing Details 🔍

Manually tested in nvcr.io/nvidia/merlin/merlin-pytorch:23.02 by installing the main branch of all Merlin libraries via pip install . --no-deps. Run 01 notebook first and then run 03 notebook.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/Transformers4Rec/review/pr-654

[rnyak]

@edknv thanks for the quick fix. I have just one comment. In this doc is says if use torchrun Change your training script to read from the LOCAL_RANK environment variable as demonstrated by the following code snippet:

import os
local_rank = int(os.environ["LOCAL_RANK"])

what do you think? does it make a big difference or not?

[edknv]

@edknv thanks for the quick fix. I have just one comment. In this doc is says if use torchrun Change your training script to read from the LOCAL_RANK environment variable as demonstrated by the following code snippet:

import os
local_rank = int(os.environ["LOCAL_RANK"])

what do you think? does it make a big difference or not?

In our case, it doesn't look like it makes a difference. From what I can tell, torchrun seems to make use of the local rank automatically. I think the doc is saying, if you need to use the local_rank variable in your script, use local_rank = int(os.environ["LOCAL_RANK"]). Our script does not make use of this variable, so I didn't include it.