pytorch 多卡并行训练实战

AI算法加油站

已于 2023-12-30 10:20:01 修改

阅读量1k

点赞数 9

分类专栏： python基础文章标签： pytorch 人工智能 python

于 2023-11-30 03:10:03 首次发布

本文链接：https://blog.csdn.net/jacke121/article/details/134702645

版权

python基础专栏收录该内容

397 篇文章 47 订阅

订阅专栏

使用torch.distributed.launch启动

nccl Windows多卡训练例子

并行框架

windows支持 gloo和mpi

UNDEFINED = "undefined" GLOO = "gloo" NCCL = "nccl" UCC = "ucc" MPI = "mpi" 这里面有

linux系统设置：

设置参数：

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int, default=64, help='Batch size for training')
    parser.add_argument('--local_rank', type=int, default=0, help='Local rank of the process')
    parser.add_argument('--device', type=str, default='0', help='Local rank of the process')
    args = parser.parse_args()

    torch.distributed.init_process_group(backend='nccl', init_method='env://')

    local_rank = args.local_rank
    device = torch.device('cuda', local_rank)
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    epoch_num = 500
    batch_size = args.batch_size


    if torch.cuda.is_available():
        net.to(device)
    net = torch.nn.parallel.DistributedDataParallel(net, device_ids= 
    [local_rank],find_unused_parameters=True)

训练时参数：

# 在训练循环中使用本地 GPU 设备
for batch in dataloader:
    inputs, labels = batch
    inputs = inputs.to(device)
    labels = labels.to(device)
    # 在这里进行训练
    ...

调用命令：

CUDA_VISIBLE_DEVICES=3,4 /data3/lbg/envs//aimet_3.8/bin/python3.8 -m torch.distributed.launch --master_port 49998 --nproc_per_node 2 train.py --device '' --batch_size 256

windows系统：

参考：windows下使用pytorch进行单机多卡分布式训练 - 知乎

使用Distributed进行分布式训练

首先了解一下概念：

node：主机数，单机多卡就一个主机，也就是1。

rank：当前进程的序号，用于进程之间的通讯，rank=0的主机为master节点。

local_rank：当前进程对应的GPU编号。

world_size：总的进程数。

在windows中，我们需要在py文件里面使用：

import os
os.environ["CUDA_VISIBLE_DEVICES]='0,1,3'

来指定使用的显卡。

假设现在我们使用上面的三张显卡，运行时显卡会重新按照0-N进行编号，有：

[38664] rank = 1, world_size = 3, n = 1, device_ids = [1]
[76032] rank = 0, world_size = 3, n = 1, device_ids = [0]
[23208] rank = 2, world_size = 3, n = 1, device_ids = [2]

也就是进程0使用第1张显卡，进行1使用第2张显卡，进程2使用第三张显卡。

有了上述的基本知识，再看看具体的实现。

使用torch.distributed.launch启动

使用torch.distributed.launch启动时，我们必须要在args里面添加一个local_rank参数，也就是：

parser.add_argument("--local_rank", type=int, default=0)

1、初始化：

import torch.distributed as dist

env_dict = {
        key: os.environ[key]
        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE")
}
current_work_dir = os.getcwd()
    init_method = f"file:///{os.path.join(current_work_dir, 'ddp_example')}"
dist.init_process_group(backend="gloo", init_method=init_method, rank=int(env_dict["RANK"]),
                                world_size=int(env_dict["WORLD_SIZE"]))

这里需要重点注意，这种启动方式在环境变量中会存在RANK和WORLD_SIZE，我们可以拿来用。backend必须指定为gloo，init_method必须是file:///，而且每次运行完一次，下一次再运行前都必须删除生成的ddp_example，不然会一直卡住。

2、构建模型并封装

local_rank会自己绑定值，不再是我们--local_rank指定的。

 model.cuda(args.local_rank)
 r_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=device_ids)

3、构建数据集加载器并封装

  train_dataset = dataset(file_path='data/{}/{}'.format(args.data_name, train_file))
  train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  train_loader = DataLoader(train_dataset, batch_size=args.train_batch_size,
                              collate_fn=collate.collate_fn, num_workers=4, sampler=train_sampler)

4、计算损失函数

我们把每一个GPU上的loss进行汇聚后计算。

def loss_reduce(self, loss):
        rt = loss.clone()
        dist.all_reduce(rt, op=dist.ReduceOp.SUM)
        rt /= self.args.local_world_size
        return rt

loss = self.criterion(outputs, labels)
torch.distributed.barrier()
loss = self.loss_reduce(loss)

注意打印相关信息和保存模型的时候我们通常只需要在local_rank=0时打印。同时，在需要将张量转换到GPU上时，我们需要指定使用的GPU，通过local_rank指定就行，即data.cuda(args.local_rank)，保证数据在对应的GPU上进行处理。

5、启动

在windows下需要把换行符去掉，且只变为一行。

python -m torch.distributed.launch \
--nnode=1 \
--node_rank=0 \
--nproc_per_node=3 \
main_distributed.py \
--local_world_size=3 \
--bert_dir="../model_hub/chinese-bert-wwm-ext/" \
--data_dir="./data/cnews/" \
--data_name="cnews" \
--log_dir="./logs/" \
--output_dir="./checkpoints/" \
--num_tags=10 \
--seed=123 \
--max_seq_len=512 \
--lr=3e-5 \
--train_batch_size=64 \
--train_epochs=1 \
--eval_batch_size=64 \
--do_train \
--do_predict \
--do_test

nproc_per_node、local_world_size和GPU的数目保持一致。

nccl Windows多卡训练例子

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
 
# 初始化每个进程中的GPU
def init_process(rank, world_size):
    # 设置每个进程使用的 GPU
    torch.cuda.set_device(rank)
    # 初始化进程组，这里使用 NCCL 作为后端，具体支持与否要看环境配置
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
 
# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc = nn.Linear(1000, 5)
 
    def forward(self, x):
        return self.fc(x)
 
# 训练函数
def train(rank, world_size):
    # 初始化进程
    init_process(rank, world_size)
    # 加载模型并分布到每个 GPU
    model = Model().cuda(rank)
    # 定义优化器
    optimizer = optim.SGD(model.parameters(), lr=0.01)
 
    # 数据分布到每个 GPU
    data = torch.randn(1000, 1000).cuda(rank)
    target = torch.randn(1000, 5).cuda(rank)
 
    for epoch in range(1000000):
        # 在每个 epoch 上执行训练
        optimizer.zero_grad()
        output = model(data)
        loss = nn.MSELoss()(output, target)
        loss.backward()
        optimizer.step()
 
        print(f"Rank {rank}, Epoch {epoch}, Loss {loss.item()}")
 
if __name__ == "__main__":
    import os
 
    # 设置单个环境变量
    os.environ['MASTER_ADDR'] = '127.0.0.1'
 
    os.environ.update({
        'MASTER_PORT': '1234',
    })
 
    # 设置每个进程的数量和 GPU 数量
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size,), nprocs=world_size)