pytorch分布式训练 DDP torchrun介绍和基本使用
0. 概念界定
关于分布式训练/多卡训练,我们首先明确一些概念:
-
并行(parallel)
多个卡训练,所有卡是一个进程或者是多个进程 -
分布式(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,这些可以帮助你更好地理解相关分布式训练方法。
在此,本人不得不推荐的还是在搜索引擎中很难搜到的两篇文章:
- 讲解pytorch的DDP:又要丝滑又要快?浅谈torch分布式训练的常见方式:DataParrel/DistributedDataParallel/Huggingface accelerate
- 讲解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 ...