Torch Distributed

torchrun

But how does torchrun actually work? what does it do under the hood? It just launches multiple processes with the right environment variables. It also gets more complicated when we have multiple nodes, it handles rendezvous, rank assignment, etc. It also has resiliency mechanisms (--max-restarts will restart the program for all ranks when a program crashes)

It uses Elastic launch.

• https://docs.pytorch.org/docs/2.14/elastic/run.html

torchrun \
  --nnodes=2 \
  --nproc-per-node=8 \
  --rdzv-id=my_job \
  --rdzv-backend=c10d \
  --rdzv-endpoint=node0:29400 \
  train.py

torchrun
   |
   +---- train.py   LOCAL_RANK=0
   +---- train.py   LOCAL_RANK=1
   +---- train.py   LOCAL_RANK=2
   +---- train.py   LOCAL_RANK=3

it gives each process environment variables, including

RANK
LOCAL_RANK
WORLD_SIZE
LOCAL_WORLD_SIZE
MASTER_ADDR
MASTER_PORT
...

• See more https://docs.pytorch.org/docs/2.14/elastic/run.html

So our script is something like

def main():
    local_rank = int(os.environ["LOCAL_RANK"])
 
    torch.cuda.set_device(local_rank)
 
    dist.init_process_group("nccl")
 
    model = ToyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])
 
if __name__ == "__main__":
    main()