Session 24¶

Topic 3: Neural networks and PyTorch¶

• introduction to PyTorch
• automatic differentiation
• basic neural network models and training
• optimization hyperparameters
• distributed training with data parallelism
• distributed training on multiple compute nodes 👈

Today¶

Distributed training on multiple compute nodes¶

1. PyTorch DDP with multiple nodes

PyTorch DDP with multiple nodes¶

From single machine → data center¶

so far:

• single GPU
• multi-GPU (one compute node)

→ now: use all available GPUs across nodes

What DDP actually does?¶

each GPU runs:

• its own model copy (identical copies)
• its own data batch (disjoint data)

after the backward pass:

• gradients are averaged (all-reduce)

→ all models stay in sync across GPUs

Abstraction: from hardware to software¶

one GPU = one worker = one process

each worker/process: controls one GPU device

each runs:

workers are independent, to agree they average after the backward step → all-reduce communication

Hardware → software mapping¶

• num_nodes → number of compute nodes
• num_gpus → GPUs per node
• WORLD_SIZE = num_nodes × num_gpus → total number of workers

\(\Rightarrow\) total number of workers = total number of GPUs

Local vs global rank¶

every worker has two unique identities: local & global

1. local_rank (unique ID within a node)

   • 0, 1, ..., num_gpus-1
   • selects GPU: cuda:{local_rank}

1. global_rank (unique ID across all nodes)

   • 0, 1, ..., WORLD_SIZE-1

how workers are indexed:

global_rank = node_id * num_gpus + local_rank

Example: 2 nodes × 4 GPUs¶

Process Structure¶

What actually gets launched?¶

SLURM, an HPC job scheduler launches:

1 task per node

inside each node, PyTorch launches:

1 process per GPU

→ tree structure: nodes → processes → GPUs

Rank 0 = main process 🦄¶

only one worker/process:

• logs metrics
• downloads data
• coordinates communication

→ the other ranks work silently

How do workers find the main process?¶

• MASTER_ADDR → where rank 0 lives
• MASTER_PORT → its communication port

Data parallel workflow¶

the remaining workflow is the same as for multi-GPU:

1. initialize processes (init_process_group)
2. split dataset across GPUs (DistributedSampler)
3. wrap the model (DistributedDataParallel)
4. run the training loop

PyTorch handles automatically: synchronization / all-reduce of gradients

Putting it all together¶

The multi-node DDP pipeline¶

1. SLURM job scheduler allocates nodes
2. each node spawns workers
3. each worker:
   • gets a global_rank and local_rank
   • binds to one GPU
4. all workers:
   • sync via init_process_group
   • train in lockstep

Key takeaways¶

• 1 GPU = 1 worker = 1 process
• local_rank selects the GPU
• global_rank identifies the worker
• WORLD_SIZE = total parallelism (all workers)

Looking ahead¶

• practice your skills by reviewing the previous sessions
• practice tensors, autograd, optimization, and distributed training

Prepare for the next session 🌱¶

sessions/25-todo