GPU and Distributed Backends

Device Topology, Gloo and NCCL

flowchart TB
    subgraph CPU["CPU demonstration (2 ranks)"]
      C0["Rank 0: full tiny model, Gloo"] <--> Link0["Gloo collectives"]
      Link0 <--> C1["Rank 1: full tiny model, Gloo"]
    end
    subgraph GPU["Linux multi-GPU optional path (not tested here)"]
      G0["Rank 0 on GPU 0: full model"] <--> Link1["NCCL collectives"]
      Link1 <--> G1["Rank 1 on GPU 1: full model"]
    end
    S["One host and torchrun"] --> CPU
    S --> GPU

Loading

Backend selection implemented in main()

Runtime	Behavior
WORLD_SIZE=1	No process group, no DDP wrapper
CPU and multiple ranks	Gloo process group
Windows with CUDA and multiple ranks	Gloo process group in the script
Linux with CUDA and multiple ranks	NCCL process group

On CUDA the script validates LOCAL_RANK < torch.cuda.device_count(), then assigns the corresponding device using torch.cuda.set_device(local_rank). There should be one process per available GPU for a standard multi-GPU run.

Memory guidance

DDP replicates the full model on each rank. Two 4 GB GPUs do not create an 8 GB shared model-memory pool. To fit larger models, consider FSDP, tensor parallelism or memory-efficient training techniques, which are not in this code.

Local checks

python -c "import torch; print(torch.__version__); print(torch.version.cuda); print(torch.cuda.is_available()); print(torch.cuda.device_count())"

The package was tested on CPU. Physical multi-GPU behavior, throughput, speedup, and peak VRAM were not measured.