33 - 完整讲解PyTorch多GPU分布式训练代码编写

1. 单机单卡

1.1 环境配置

• 判断显卡是否可用

import torch

# 判断环境中是否有cuda
print(f"torch.cuda.is_available()={torch.cuda.is_available()}")
# torch.cuda.is_available()=True

# 通用编程语句
device = torch.device('cuda:0' if torch.cuda.is_available()else 'cpu')
print(f"device={device}")
# device=cuda:0

• 设置显卡是否在程序中可见
  (1) 代码

# 代码形式，设置环境变量，第 0 个显卡可见
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

(2) 命令行

# 命令行模式下设置显卡可见，python表示解释器python ；xxxx.py表示的是执行的py文件
CUDA_VISIBLE_DEVICES="0" python  xxxx.py

1.2 模型拷贝

# 直接调用cuda函数,不需要复制（model=model.cuda是错误的），
model.cuda()
#  nn.module类中可以查看到函数具体调用细节
# 具体流程如下：
# （1）model.cuda()
# （2）self._apply
# （3）for module in self.children():
			module._apply(fn)

1.3 数据拷贝

数据的拷贝实际上调用的是张量的cuda函数[Tensor = Tensor.cuda()]，将张量拷贝到GPU上

# 对于 X 来说
source = x = token_index ==> token_index = token_index.cuda()
target = y = target      ==> target = target.cuda()

1.4 模型保存

# 定义训练的轮数epoch
EPOCH = 5

# 参数保存后文件名
PATH = "mode.pt"

# torch.save用字典的形式保存训练相关的参数值

torch.save({
	'epoch':EPOCH,
	'model_state_dict':net.state_dict(),
	'optimizer_state_dict':optimizer.state_dict(),
	'Loss':loss,},PATH)

1.5 模型的加载

# file.pt 表示加载的模型参数文件
# 方式 1 ：加载到GPU，默认GPU0
checkpoint = torch.load(file.pt,map_location=torch.device("cuda"))
# 方式 2 ：加载到GPU1上
checkpoint = torch.load(file.pt,map_location=torch.device("cuda:1"))
# 方式 3： 加载到CPU上
checkpoint = torch.load(file.pt,map_location=torch.device("cpu"))

1.6 注意事项

当我们用logging.warning(),不能直接输出张量(影响程序性能)，需要将张量(这里特指标量)转换成python类型输出;

import logging

acc = torch.Tensor([1.0])

# 张量输出以你选哪个性能，需要转换成python变量  
logging.warning(f"acc={acc}") # 影响性能，应该避免
# 输出结果：WARNING:root:acc=tensor([1.])

# 将张量转换成python变量
logging.warning(f"acc.item()={acc.item()}") # 改善后
# 输出结果：WARNING:root:acc.item()=1.0

2. 单机多卡(nn.DataParallel)

2.1 环境配置

import logging

import torch

# 检测显卡是否可用，True 表示可用
if torch.cuda.is_available():
	logging.warning("cuda is available!")
	# 判断环境中多少个显卡。如果大于1就输出显卡个数，否则
	if torch.cuda.device_count() > 1:
		logging.warning(f"find{torch.cuda.device_count()}GPUS!")
	else:
		logging.warning("it is only one GPU!")
else:
	logging.warning("cuda is not available,exit!")
# 代码设置环境变量
# 设置环境中第0个和第1个显卡可见
os.environ['CUDA_VISIBLE_DEVICES']="0,1"

# 命令行中实现设置
# python :解释器；
# xxxx.py ：需要执行的文件
# CUDA_VISIBLE_DEVICES="0,1" python xxxx.py

2.2 模型拷贝

直接将模型用nn.DataParallel包裹起来即可；

model = nn.DataParallel(model.cuda(),device_ids=[0,1,1,3])

2.3 数据拷贝

data = data.cuda()

2.4 模型保存与加载

# 模型的保存
# Additional information
EPOCH = 5
PATH = "model.pt"
LOSS = 0.4

torch.save({
            'epoch': EPOCH,
			# 因为包裹了一个nn.DataParallel,多了一个module
            'model_state_dict': model.module.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, PATH)
# 模型的加载
checkpoint = torch.load(resume,map_location=torch.device("cuda:0"))

2.5 nn.DataParallel的缺点

• 单进程效率慢
• 不支持多机情况
• 不支持模型并行
• 注意事项：我们需要将此处的BATCH_SIZE=64*2,这样分配到GPU0=64，GPU1=64

3. 单机多卡(nn.DistributedDataParallel)

3.1 设置环境变量

# 方式1：通过代码设置
torch.cuda.set_device(args.local_rank) # 表示用第 local_rank 张卡

# 方式2：通过命令行设置环境变量
CUDA_VISIBLE_DEVICES="0,1"

3.2 初始化进程组

torch.distributed.init_process_group("ncll",world_size=n_gpus,rank=args.local_rank)
# "ncll" : 表示GPU之间的通讯方式，推荐
# "world_size": 表示当前节点上有多少个GPU卡
# "rank":表示当前进程是在哪个GPU上
# args.local.rank是通过命令 python -m torch.distributed.launch --nproc-node=n_gpus train.py 出入的

3.3 数据的拷贝和分包

因为存在单机多卡的情况，所以我们需要将train_datasets and eval_datasets根据GPU的个数来进行不重叠分割来保证均等随机的分配到每个GPU上，这里我们用到pytorch官网上提供的类 DistributedSampler 来对train_datasets and eval_datasets进行均匀随机分割；第一要做到均匀分割到不同GPU上，第二要将train_datasets and eval_datasets里面的样本打散；

# 2.split the dataloader
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader

def prepare(rank,world_size,batch_size=32,pin_memory=False,num_workers=0):
	# Your_Dataset() 就是需要拆包的train_dataset and eval_dataset
	dataset = Your_Dataset()
	# num_replicas：GPU的个数
	# rank：表示第 rank 个 GPU
	# shuffle：是否需要打乱
	# drop_last：表示是否将最后几个丢弃掉，保证样本个数是GPU个数的整数倍
	
	sampler = DistributedSampler(dataset,num_replicas=world_size,
								 rank=rank,shuffle=False,drop_last=False)
								 
	dataloader = DataLoader(dataset,batch_size=batch_size,
							pin_memory=pin_memory,num_workers=num_workers,drop_last=False,
							shuffle=False,sampler=sampler)
	return dataloader

data = data.cuda(args.local_rank)

3.4 模型的保存和加载

• 保存
  模型的参数值只需要在GPU0上保存即可；torch.save在local_rank=0位置保存，同样注意调用model.module.state_dict()
• 加载
  torch.load 注意 map_location=torch.device(“cuda:0”)

3.5 注意事项

1. train.py 中要有接受 local_rank的参数选项，launch会传入这个参数
2. 每个进程的BATCH_SIZE应该是一个GPU所需要的batch_size大小
3. 在每个周期开始出调用 train_sampler.set_epoch(epoch)可以使得数据充分打乱

# 是为了每张卡在每个周期中得到的数据是随机的，可以查看DistributedSampler类的set_epoch函数
train_sampler.set_epoch(epoch_index)
# from DistributedSampler.__iter__函数源码
#    g = torch.Generator()
#    g.manual_seed(self.seed + self.epoch)
# 可以看出随机种子设置为self.seed+self.epoch所以我们需要在每个周期中设置set_epoch
# 这样就可以保证随机种子不一样了。

4. 有了sampler就不要在DataLoader中设置shuffle=True了

3.6 代码优化修改

# 在命令行中解析相关参数
import argparse
parser=argparse.ArgumentParser()
parser.add_argument("--local_rank",help="local deviced on current node",type=int)
args = parser.parse_args()

n_gpus = 2
torch.distributed.init_process_groups("nccl",world_size=n_gpus,rank=args.local_rank)
torch.cuda.set_device(args.local_rank)

# 函数def train(xxxx,local_rank,train_datasests)
model = nn.Parallel.DistributedDataParallel(model.cuda(local_rank),devices_ids=[localrank])

# dataset_train分割
train_sampler=DistributedSampler(train_datasets)

train_data_loader=torch.utils.data.DataLoader(train_dataset,batch_size=BATCH_SIZE,collate_fn=collate_fn,
											  sampler=train_sampler)
eval_data_loader = torch.utils.data.DataLoader(eval_dataset,batch_size=8,collate_fn=collate_fn)

# 数据拷贝到GPU中
token_index = token_index.cuda(local_rank)
target=target.cuda(local_rank)

4. 多机多卡

torch.nn.parallel.DistributedDataParallel

1. 代码编写流程跟单机多卡一致
2. 执行命令以两节点为例，每个节点出有n_gpus个GPU

python -m torch.distributed.launch -- nproc_per_node=n_gpus--nnodes=2--node_rank=0--master_addr="主节点IP"--master_port="主节点端口" train.py

python -m torch.distributed.launch -- nproc_per_node=n_gpus--nnodes=2--node_rank=1--master_addr="主节点IP"--master_port="主节点端口" train.py

5. 模型并行

模型参数太大，单个GPU无法容纳模型，需要将模型的不同层拆分到多个GPU上

6. 参考网址

A Comprehensive Tutorial to Pytorch DistributedDataParallel
PyTorch Distributed Training