This project implements a simplified distributed training runtime inspired by DeepSpeed and Megatron-LM, built on PyTorch DDP and NCCL/MPI. It runs efficiently on multi-GPU or multi-node setups such as the CU Boulder Alpine HPC cluster.
- PyTorch Distributed (DDP)
- CUDA 12.x + NCCL
- MPI + Slurm
- Python 3.10+
- Profiling: torch.profiler, Nsight Systems
module purge
module load gcc/12.2.0 cuda/12.1 openmpi/4.1.5 anaconda/2024.02
conda create -y -n mini-ds python=3.10
conda activate mini-ds
pip install torch torchvision torchaudio matplotlib pandas tqdm mpi4py