PyTorch Lightning for Distributed Training
This repo contains sample code for distributed training for different configurations.
Distributed Data Parallel
- Data used: MNIST
Multi-Node, Single GPUs
- Code
- T4 GPU on 2 instances Command:
MASTER_ADDR={IP of RANK 0} MASTER_PORT=29500 NODE_RANK=0 python main.py
MASTER_ADDR={IP of RANK 0} MASTER_PORT=29500 NODE_RANK=1 python main.py
Single-Node, Multiple GPUs
Code
Command:
pip install torch_tb_profiler if profiling needed; else disable it
MASTER_ADDR=localhost MASTER_PORT=29500 WORLD_SIZE=2 NODE_RANK=0 python main.py
tensorboard --logdir=./tensorboard/ --host=0.0.0.0 # to view tensorboard
Distributed Model Parallel
- Data used: Intel Image Classification from
Strategy: FSDP
Multi-Node, Single GPUs
- ViT Model is used
Caveats: - Number of devices = 2 (1 per node; 2 in total). This is in contrast with DDP training.
- PTL checkpointing doesn't work. The weights are not stored.
- Models are saved manually and then the best model is copied to the root folder.
export MASTER_PORT=29500
export MASTER_ADDR=172.31.10.239
export WORLD_SIZE=2
export NODE_RANK=0 # and 1 respectively
python -m torch.distributed.run \
--nnodes=$WORLD_SIZE \
--master_addr $MASTER_ADDR \
--master_port $MASTER_PORT \
--node_rank $NODE_RANK \
main.py
Single Node, Multiple GPUs
MASTER_ADDR=localhost MASTER_PORT=29500 WORLD_SIZE=1 NODE_RANK=0 python main.py