Pytorch 多机多卡分布式训练

SJTU 路子

已于 2024-08-16 18:05:57 修改

阅读量1.3k

点赞数 30

分类专栏： pytorch 文章标签： pytorch

于 2024-08-08 18:13:10 首次发布

本文链接：https://blog.csdn.net/qq_34826149/article/details/141033501

版权

pytorch 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

概述

在pytorch中，进行多机多卡分布式训练可以大大提高训练效率和缩短训练时间. 我们通常使用官方提供的DDP (distributed data parallel)模块来实现多机多卡训练，下面主要以介绍官方文档 https://pytorch.org/tutorials/intermediate/ddp_tutorial.html 给出的demo作为例子，来介绍DDP的使用模版以及运行流程.

单机单卡Toy Code

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)
 
    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))
 
def train_on_cpu():
	model = ToyModel()
	loss_fn = nn.MSELoss()
	optimizer = optim.SGD(model.parameters(), lr = 0.001)
	inputs = torch.randn(20, 10)
	labels = torch.randn(20, 5)
 	for i in range(10):
 		optimizer.zero_grad()
 		outputs = model(inputs)
 		loss_fn(outputs, labels).backword()
 		optimizer.step()
 
 def train_on_single_gpu():
 	# 将 model 传送到GPU 0上。当你调用 model.to("cuda") 时，模型默认被传送到 GPU 0 上。
 	# 这是因为 "cuda" 默认指向 CUDA 设备列表中的第一个设备，即 cuda:0
 	# 如果当一个机器有多卡时，通常需要指定GPU后面的索引号
 	model = ToyModel().to("cuda:0") 
	loss_fn = nn.MSELoss()
	optimizer = optim.SGD(model.parameters(), lr = 0.001)
	
	# inputs和labels数据都需要传送到对应的GPU上
	inputs = torch.randn(20, 10).to("cuda:0")
	labels = torch.randn(20, 5).to("cuda:0")
 	for i in range(10):
 		optimizer.zero_grad()
 		outputs = model(inputs)
 		loss_fn(outputs, labels).backword()
 		optimizer.step()

单机多卡Toy Code

使用DDP进行单机多卡训练时，通过多进程在多个GPU上复制模型，在每个GPU都由一个进程控制，同时需要将参数local_rank传递给进程，用于表示的当前进程使用的是哪一个GPU。
要将单机多卡训练修改为基于DDP的单机多卡训练，需要进行的修改的部分如下:

初始化进程组 dist.init_process_group
设置分布式采样器 DistributedSampler
使用DistributedDataParallel封装模型
使用torchrun 或者 mp.spawn 启动分布式训练

1. 进程组初始化

在每个进程中，都需要初始化一个进程组，这是分布式训练的基础。进程组中的每个进程都可以与其他进程进行通信。
需要设置rank，world_size, local_rank参数，并需要设置cuda device， init_process_group

import os

import torch
import torch.distributed as dist

def dist_init():
	"""
	Initialize the horovod and allocate GPUS.
	Return:
		local_rank, rank and world_size of the distributed system.
	"""
	rank, world_size, local_rank = 0, 1, 0
	if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
		rank = int(os.environ["RANK"])
		world_size = int(os.environ["WORLD_SIZE"))
		local_rank = int(os.environ["LOCAL_RANK"))
		os.environ["LOCAL_SIZE"] = str(torch.cuda.device_count())
	dist_url = "env://"
	torch.cuda.set_device(local_rank)
	dist.init_process_group(backend="nccl", init_method=dist_url, world_size=world_size, rank=rank)
	return rank, world_size, local_rank

Question 1. rank, local_rank, world_size的定义是什么？
1. rank表示所有分布式训练进程中唯一的全局编号，范围为[0, world_size - 1], 这对进程相互识别和通信非常重要
2. local_rank表示当前进程在其运行的物理机器或虚拟机上的本地编号，也就是单个节点上的进程标志，这对进程使用哪个GPU很重要
3. world_size表示参与当前分布式训练任务的总进程数，在启动分布式训练任务的时候系统会自动配置

Question 2. 环境变量“RANK”, “WORLD_SIZE”, "LOCAL_RANK"等字段是在什么时候被设置的？
这些字段是有分布式训练框架（torch.distributed.launch）或者作业调度系统（Kubernetes, Slurm, MPI等）在启动每个训练进程时自动设置的。

Question 3. torch.cuda.set_device(local_rank)的用处是什么?
用处是指定当前pytorch进程使用哪个GPU设备。确保当在同一个节点运行多进程时，每个进程都绑定了一个独立的GPU，避免了资源冲突和效率低下。

Question 4. 为什么dist.init_process_group(backend=“nccl”, init_method=dist_url, world_size=world_size, rank=rank)传入的是rank，而不是local_rank. backend和init_method参数的含义是什么？
这里主要进行的是进程组初始化，使得它们能够进行同步操作，比如同步梯度，广播模型参数等，这对于保持模型在各个进程上的一致性至关重要。进程与进程之间通信时同时全局唯一标识符rank而不是通过local_rank来识别对应的进程的。
backend指定了通信所用的后端技术吗，对于NVIDIA GPU环境，通常选择“nccl”后端。 init_method 指定了如何配置进程间的初始连接，常见的方法包括使用环境变量（env://）

2. 数据加载

将整个数据切分成多份，使得每个进程处理各自的数据分片，一般使用torch.utils.data.distributed.DistributedSampler.

import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

class ToyDataset(Dataset):
	def __init__(self， num_samples = 100):
		self.dataset = [torch.randn(20, 10)] * num_samples
	
	def __len__(self):
		return len(self.dataset)
	
	def __getitem__(self, idx):
		return self.dataset[idx]

dataset = ToyDataset()
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) # world_size，rank由前面进程组初始化dist_init得到
loader = DataLoader(dataset, batch_size=8, sampler=sampler)

Question 1. DistributedSampler的用处是什么
DistributedSampler确保每一个进程获取的是数据集的一个子集，且各子集互不重叠

3. 数据并行，模型梯度同步

使用torch.nn.paraller.DistributedDataParallel来自动处理数据的分布和梯度的同步汇总


from torch.nn.parallel import DistributedDataParallel as DDP

model = ToyModel()
model = model.cuda(local_rank) #将模型移到对应的GPU， local_rank由前面进程组初始化dist_init得到
ddp_model = DDP(model, device_ids=[local_rank])

Question 1: 为什么这里是是local_rank而不是rank
rank是进程的唯一全局标识符，而local_rank直接关联到具体的GPU设备。

Question 2: device_ids传入的是一个list
device_ids参数接收一个列表，指明了哪些GPU（在这个例子是是单个GPU）被该进程使用

4. 循环训练

在训练中的，确保正确使用的sampler和分布式模型

for epoch in range(num_epochs):
	sampler.set_epoch(epoch)
	for data, labels in loader:
		optimizer.zero_grad()
		outputs = ddp_model(data)
		loss = loss_fn(outputs, labels)
		loss.backward()
		optimizer.step()
	torch.save(ddp_model.module().state_dict(), f'checkpoint_epoch{epoch}.pth')

Question 1: sampler.set_epoch(epoch)的用处
设置分布式采样器当前的epoch, 使得数据加载顺序在每个epoch都会有所不同，有助于增加模型的泛化能力

Question 2: 为什么在保存模型权重的时候，是用ddp_model.module().state_dict(), 而不是ddp_model.state_dict()?
由于ddp_model = DDP(model, device_ids=[local_rank])使得ddp_model相比于model封装了一层。所有parameter的name都多了前缀module。ddp对象管理者数据的分布式传递，梯度的汇总等任务，以支持分布式训练。

5. 任务启动

在单节点上，我们通常使用torch.distributed.launch来实现任务的启动

python3 -m torch.distributed.launch --nproc_per_node=4 test.py

Question 1: torch.distributed.launch介绍，有哪些参数，哪些是必须设置的
torch.distributed.lauch是用于启动Pytorch 分布式训练的工具。
nproc_per_node 是必须设置的，每个节点启动的进程数，通常设置为每个节点上的GPU数量
nnodes 表示参与训练的节点总数, 主要应用于多机多卡训练, 单机的话默认是1
node_rank表示当前节点的排名
master_addr主节点的地址，用于所有节点间的通信
master_port 主节点用于通信的端口、

在torhc1.10开始用终端命令torchrun来代替torch.distributed.launch, torchrun实现了launch的一个超集

torchrun --nproc_per_node=2 test_gpu.py

多机多卡Toy Code

多机多卡与单机多卡不同之处主要在于前者需要在多个节点上跑命令

torchrun --nnodes=2 --node_rank=0 --nproc_per_node=8 test_gpu.py
torchrun --nnodes=2 --node_rank=1 --nproc_per_node=8 test_gpu.py

Reference:
https://shomy.top/2022/01/05/torch-ddp-intro/
https://232525.github.io/posts/PyTorch%E5%8D%95%E6%9C%BA%E5%A4%9A%E5%8D%A1%E8%AE%AD%E7%BB%83%E5%A4%87%E5%BF%98%E8%AE%B0%E5%BD%95/
https://www.cnblogs.com/chentiao/p/17666330.html

其他
在实际操作中，我们通常是每个进程控制一个GPU, 操作简单，分布式数据并行（distributed data parallel）通常是以进程的视野去探究的，包括进程组初始化（进程同步和通信），每个进程指定gpu(torch.cuda.set_device()), 数据的sampler（需要传入对应的rank), 模型的数据同步（需要传送到对应的GPU上, 传入device_ids信息以指定GPU）

补充一点使用分布式做evaluation的时候，一般需要先所有进程的输出结果进行gather，再进行指标的计算，两个常用的函数:
dist.all_gather(tensor_list, tensor) : 将所有进程的tensor进行收集并拼接成新的tensorlist返回，比如:
dist.all_reduce(tensor, op) 这是对tensor的in-place的操作, 对所有进程的某个tensor进行合并操作，op可以是求和等