实验随记2-Pytorch Lightning多机多卡训练

本文章主要收录笔者在阅读时发现的一些比较优质的多机多卡训练教程~~，由于pytorchLightning多机多卡示例很少，因此需要等笔者最近做完实验验证后才能更新示例。不断完善中…~~ 及多机多卡训练细节。

任务需求：实现多机多卡训练，模型可能继承torch.nn或者torch.lightning。

示例

• 实现5节点4GPU共计20卡训练。
  • 使用pytorch_lightning == 1.9.4 存在bug: 启动4机24卡时，per node only use one GPU!
  • 使用lightning.pytorch == 2.1.2 存在bug: 启动4机24卡时，After all processes are distributed, all the processes hang forvever（不同Node抢占初始化相同world_size的进程）

环境及杂项声明：

• lightning==2.1.2
• import lightning.pytorch(在一些较低的版本中，import pytorch_lightning也可以。这样会在on_optimizer_step附近传参个数有所不同，以及Trainer初始化过程中有gpus声明。对于-1和"auto"的支持成都有所不同）
• python==3.9
• dataloader: pin_memory=False(注意num_workers的数量不能太大，有可能会在distributed之后卡住）
• distributed_backend=nccl
• A100，显存80G，若干

多机多卡代码端添加内容

与pytorch lightning教程内容保持一致。注意这里需要区分环境中不同版本的lightning，查看对应版本的教程。

Trainer(...(your params)..., devices=5, gpus=4, strategy="ddp", accelerator="gpu")

添加完代码后，直接使用python scripts.py启动代码，在All distributed processes registered. Starting with xxx processes之后卡住。因此需要继续研究下面的内容。

多机多卡启动

经研究（踩坑），使用torchrun可以正确启动，分配workers正确的global rank，以及分配给worker中不同GPU正确的local rank。
根据资源数量修改GPU_NU和NUM_NODES的数值。然后chmod 777 scripts.sh，并启动该scripts。即可正常使用多机多卡。

#!/bin/bash
cd /your/path/to/code
export OMP_NUM_THREADS=8       # change value in each execution
export GPU_NUM=4
export NUM_NODES=5
export TRAIN_DATA_PATH='../data/train'
export VAL_DATA_PATH='../data/eval'
export TEST_DATA_PATH=='../data/test'

/your/path/to/environment/bin/torchrun \
    --nnodes=$NUM_NODES \
    --nproc_per_node=$GPU_NUM \
    --rdzv_id=distributed_alldata \
    --rdzv_backend=c10d \
    --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT \
scripts.py (--your training params...)

集群设置及结果

单个节点GPU运行情况

文章list

讲述如何在torch.nn类模型中实现多机多卡的文章比较多。如下：

前置必备知识

Getting Started with Distributed Data Parallel：https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
From PyTorch DDP to Accelerate to Trainer, mastery of distributed training with ease：https://huggingface.co/blog/zh/pytorch-ddp-accelerate-transformers

官方文档

TORCH DISTRIBUTED ELASTIC：https://pytorch.org/docs/stable/distributed.elastic.html
Run with TorchRun (TorchElastic)：https://lightning.ai/docs/pytorch/latest/clouds/cluster_intermediate_2.html

概念理解

多机多卡数据并行-DistributedDataParallel(DDP)：https://support.huaweicloud.com/develop-modelarts/modelarts-distributed-0008.html

PyTorch Lightning多机多卡示例

Training on Accelerators：https://lightning.ai/docs/pytorch/stable/advanced/speed.html
DDP OPTIMIZATIONS：https://lightning.ai/docs/pytorch/latest/advanced/ddp_optimizations.html#ddp-optimizations
Pytorch分布式训练的最简单攻略—Pytorch-Lightning：https://zhuanlan.zhihu.com/p/611870086

搭建全套pl.lightning框架

自定义DataModule类（继承LightningDataModule）来调用DataLoader：https://developer.aliyun.com/article/1294235
定义customed dataloader教程: https://lightning.ai/docs/pytorch/2.1.2/data/datamodule.html