Correctness of when to call `set_device` in the docs for DDP

[craymichael]

📚 The doc issue

In the docs tutorial on how to set up Multi-GPU training, it is suggested that the following is the proper way to setup each process (initializing the, e.g., NCCL, process group and then calling torch.cuda.set_device(rank)):

def ddp_setup(rank: int, world_size: int):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

However, these issues suggest that the proper way is to call set_device before initializing the process group:

• Call to CUDA function failed. with DDP using 4 GPUs pytorch#54550 (comment)
• distributed.all_gather function stuck when using NCCL backend  pytorch#18689 (comment)

Which is the correct order? Are there pauses or slowdowns if the order changes?

Suggest a potential alternative/fix

No response

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225

[wconstab]

cc @kwen2501 @shuqiangzhang

[fegin]

With TorchRun, there is a LOCAL_RANK environment variable that allows to set cuda device before initializing PG. It is safer to call set_device before initializing PG.

[wconstab]

we should update our docs to reference this though.
unless we don't think its what we want long term. @kwen2501 is proposing for users to pass 'device' into init_process_group to enable eager initialization. if we go that way, do users still need to call set_device first or would we consider doing that inside init_process_group?

[fegin]

We don't if init_process_group accepts the device. But for the time being, I'll have a PR to update the doc.

[svekars]

Completed