How to launch a PyTorch script with torchrun

A distributed training entry point must receive the process identity assigned by its launcher before it can join a process group. torchrun supplies that identity to every Python worker, so a two-process CPU probe can check the launch contract before model or GPU code is added.

For a single host, --standalone creates a local rendezvous and --nproc-per-node=2 starts two worker processes. Each process receives LOCAL_RANK, while the initialized process group reports RANK and WORLD_SIZE.

The probe uses the Gloo backend so it can run without GPUs. PyTorch must already be installed, and the two rank lines may appear in either order because separate workers write to the terminal independently.

1. Create rank_probe.py with the imports and process-group initialization.

   rank_probe.py

   import os
    
   import torch.distributed as dist
    
    
   dist.init_process_group("gloo")

2. Append the worker rank lookup to rank_probe.py.

   rank = dist.get_rank()
   world_size = dist.get_world_size()
   local_rank = int(os.environ["LOCAL_RANK"])

3. Append the formatted rank output to rank_probe.py.

   print(f"rank={rank} local_rank={local_rank} world_size={world_size}")

4. Append the process-group cleanup to rank_probe.py.

   dist.destroy_process_group()

5. Set one OpenMP thread for each worker in the current shell.

   $ export OMP_NUM_THREADS=1

   The setting limits CPU thread contention during this local smoke run and lasts only for the current shell.

6. Launch two local workers with torchrun to confirm both ranks join one process group.

   $ torchrun --standalone --nproc-per-node=2 rank_probe.py
   rank=1 local_rank=1 world_size=2
   rank=0 local_rank=0 world_size=2

   The line order may vary. Distinct ranks with world_size=2 show that both workers joined the same launch.