How to launch a PyTorch script with torchrun

torchrun starts multiple Python worker processes from one PyTorch entry point, which makes it the usual launcher for local distributed smoke tests before a script moves to GPUs or a cluster. A small rank probe is enough to confirm that each process receives its own rank and the shared world-size values.

For a single-node test, --standalone creates a local rendezvous backend and --nproc-per-node=2 starts two workers on the same machine. The script can read LOCAL_RANK, RANK, WORLD_SIZE, and LOCAL_WORLD_SIZE from the environment, then initialize torch.distributed with gloo for CPU-safe verification.

The probe assumes PyTorch is already installed and torchrun is on the shell path. It keeps the run CPU-only, sets one OpenMP thread per worker to avoid local oversubscription warnings, and removes the temporary script after the launch is proven.

1. Create a rank probe script.

   rank_probe.py

   import os
    
   import torch.distributed as dist
    
    
   dist.init_process_group("gloo")
   rank = dist.get_rank()
   world_size = dist.get_world_size()
   local_rank = int(os.environ["LOCAL_RANK"])
   local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
    
   message = (
       f"rank={rank} local_rank={local_rank} "
       f"world_size={world_size} local_world_size={local_world_size}"
   )
    
   gathered = [None] * world_size if rank == 0 else None
   dist.gather_object(message, gathered, dst=0)
    
   if rank == 0:
       for line in gathered:
           print(line)
       print("torchrun_launch_ok=True")
    
   dist.destroy_process_group()

   gloo works for a CPU launch probe. If a training script parses local rank from command-line arguments, accept both --local-rank and --local_rank so older and newer launchers do not break argument parsing.

2. Limit OpenMP threads for the local smoke run.

   $ export OMP_NUM_THREADS=1

   This setting applies only to the current shell. Increase it later only after measuring CPU contention for the real training script.

3. Launch two local workers with torchrun.

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

   --standalone starts a local rendezvous for one host. --nproc-per-node=2 launches two worker processes and sets LOCAL_WORLD_SIZE to 2.

4. Confirm that the launch produced one line for each rank.

   rank=0 and rank=1 show two separate worker processes. Matching world_size=2 and local_world_size=2 confirm that both workers joined the same local launch.

5. Remove the temporary probe script.

   $ rm rank_probe.py