Lightning-AI/pytorch-lightning

Why does running Lightning on SLURM with python perform worse than with srun?

Open

#18,650 创建于 2023年9月27日

在 GitHub 查看
 (21 评论) (2 反应) (0 负责人)Python (3,233 fork)batch import
environment: slurmhelp wantedquestionver: 2.0.x

仓库指标

Star
 (26,687 star)
PR 合并指标
 (平均合并 9天 15小时) (30 天内合并 3 个 PR)

描述

Bug description

I'm training LLMs across multiple GPUs on a single node using Nvidia/NeMo. When launching via python train.py inside of an allocation I get much worse performance than when launching directly via srun. In the first case Pytorch Lightning also raises the warning: The srun command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with srun. See below:

107 tflops:

srun --unbuffered --exclusive --open-mode=append --partition=ultra --nodes=1 --ntasks-per-node=4 --gpus=4 --pty bash
source /home/niklas/miniconda/bin/activate
conda activate core
CUDA_DEVICE_MAX_CONNECTIONS=1 GPU_PER_NODE=4 python train.py

172 tflops:

source /home/niklas/miniconda/bin/activate
conda activate core
CUDA_DEVICE_MAX_CONNECTIONS=1 GPU_PER_NODE=4 srun --unbuffered --exclusive --open-mode=append --partition=ultra --nodes=1 --ntasks-per-node=4 --gpus=4 python -u train.py

Why is it the case that the first one performs worse? Maybe is there a difference in how these two strategies launch the torch.distributed process group? (https://pytorch.org/docs/stable/notes/multiprocessing.html#cuda-in-multiprocessing)

What version are you seeing the problem on?

v2.0

How to reproduce the bug

No response

Error messages and logs

No response

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @awaelchli

贡献者指南