pytorch分布式训练 DDP torchrun介绍和基本使用

0. 概念界定

关于分布式训练/多卡训练，我们首先明确一些概念：

1. 并行（parallel）

多个卡训练，所有卡是一个进程或者是多个进程

2. 分布式（Distributed）

多个卡训练，但是每张卡上都是一个单独的进程

但是，在下文中，为了表述的方便性，我们不严格区分“并行”和“分布式”两个概念，不论说并行训练还是分布式训练，都指的是使用多个显卡进行训练。

1. 简介

并行可以分为数据并行、模型并行（张量并行）、流水线并行。我们这里主要介绍的就是最常用的数据并行（DataParallel）。关于另外两种并行，有兴趣的读者可以自己查阅相关资料。

简单来说，pytorch分布式数据并行训练的原理是，在每张卡上跑不同的数据，然后更新参数时通过卡间通信来同步梯度等信息，保证所有卡的参数一致。

分布式训练从算法原理上大致可以分为DP、DDP、zero stage 1/2/3这几种方式。对于这几种算法，各个框架都有自己的实现，主要包括(1). pytorch官方实现，(2). 微软deepspeed。

pytorch官方

• DP(torch.nn.DataParallel)

• DDP(torch.nn.parallel.DistributedDataParallel)

• zero-1/2/3(torch.distributed.fsdp.fully_sharded_data_parallel)。fsdp是pytorch 1.11发布的最新的分布式训练框架，支持DDP和zero系列算法。zero-0就是DDP。

微软deepspeed

• zero-0/1/2/3都在deepspeed中实现了。

若要学习分布式训练的使用方法，pytorch的tutorials有一节专门讲Parallel and Distributed Training，在docs中也有相关的Developer Notes：ddp notes，fsdp notes，这些可以帮助你更好地理解相关分布式训练方法。

在此，本人不得不推荐的还是在搜索引擎中很难搜到的两篇文章：

1. 讲解pytorch的DDP：又要丝滑又要快？浅谈torch分布式训练的常见方式：DataParrel/DistributedDataParallel/Huggingface accelerate

2. 讲解zero系列原理：分布式训练：了解Deepspeed中的ZeRO1/2/3

2. 使用方法

最简单的办法是结合Huggingface的Transformers库的Trainer使用，Trainer帮我们做好了多进程数据切分，以及并行训练引擎初始化的过程，可以让多卡和单卡训练一样简单。

假如你有一个训练脚本run_clm.py（提示一下，这里是可以点击跳转到Github查看代码的），你使用Trainer训练模型，你原始调用的脚本是python train_model.py --outputdir ...，那么，你只需要改成：

deepspeed --include "localhost:1,2,3,4" train_model.py \
	--deepspeed ds_config.json \
	--output_dir ...

这里改了两个地方，一个是将python改为了deepspeed，并指定使用1,2,3,4号显卡，另一个是给训练脚本添加了deepspeed参数，Huggingface的trainer会自动解析这个参数并使用deepspeed。

需要注意的是，deepspeed不能使用CUDA_VISIBLE_DEVICES来设置使用的GPU，务必不要设置CUDA_VISIBLE_DEVICES

当然，也可以使用torchrun：

CUDA_VISIBLE_DEVICES=1,2,3,4 torchrun --nproc-per-node 4 run_clm.py \
	--deepspeed ds_config.json \
	--output_dir ...

=================================

分割线

=================================

这里简单记录一下，并非教程。

pytorch ddp用法:

master的端口默认是29500，如果被占用了就换一个

torchrun --master_port 61234 --nproc_per_node $gpu_num train.py ...

shell脚本：

export CUDA_VISIBLE_DEVICES=$1
gpu_num=$(echo $CUDA_VISIBLE_DEVICES | awk -F ',' '{print NF}')

torchrun --master_port 61234 --nproc_per_node $gpu_num hf_train.py ...