当前位置：首页 > article >正文

Pytorch使用教程(12)-如何进行并行训练？

article 2025/1/21 17:28:49

在使用GPU训练大模型时，往往会面临单卡显存不足的情况。这时，通过多卡并行的形式来扩大显存是一个有效的解决方案。PyTorch主要提供了两个类来实现多卡并行：数据并行torch.nn.DataParallel（DP）和模型并行torch.nn.DistributedDataParallel（DDP）。本文将详细介绍这两种方法。

一、数据并行（torch.nn.DataParallel）

基本原理
数据并行是一种简单的多GPU并行训练方式。它通过多线程的方式，将输入数据分割成多个部分，每个部分在不同的GPU上并行处理，最后将所有GPU的输出结果汇总，计算损失和梯度，更新模型参数。
使用方法
使用torch.nn.DataParallel非常简单，只需要一行代码就可以实现。以下是一个示例：

import torch
import torch.nn as nn

# 检查是否有多个GPU可用
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # 将模型转换为DataParallel对象
    model = nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))

优缺点
‌优点‌：代码简单，易于使用，对小白比较友好。
‌缺点‌：GPU会出现负载不均衡的问题，一个GPU可能占用了大部分负载，而其他GPU却负载较轻，导致显存使用不平衡。

二、模型并行（torch.nn.DistributedDataParallel）

基本原理
torch.nn.DistributedDataParallel（DDP）是一种真正的多进程并行训练方式。每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。每个进程包含独立的解释器和GIL（全局解释器锁），因此可以充分利用多GPU的优势，实现更高效的并行训练。
使用方法

使用torch.nn.DistributedDataParallel需要进行一些额外的配置，包括初始化GPU通信方式、设置随机种子点、使用DistributedSampler分配数据等。以下是一个详细的示例：

初始化环境

import torch
import torch.distributed as dist
import argparse

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    return args

def main():
    args = parse()
    torch.cuda.set_device(args.local_rank)
    dist.init_process_group(
        'nccl', init_method='env://'
    )
    device = torch.device(f'cuda:{args.local_rank}')

设置随机种子点

import numpy as np

# 固定随机种子点
seed = np.random.randint(1, 10000)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

使用DistributedSampler分配数据
python
Copy Code
from torch.utils.data.distributed import DistributedSampler

train_dataset = ...  # 你的数据集
train_sampler = DistributedSampler(train_dataset, shuffle=True)
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=opts.batch_size, sampler=train_sampler
)

初始化模型

model = mymodel().to(device)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

训练循环
python
Copy Code
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()

for ep in range(total_epoch):
    train_sampler.set_epoch(ep)
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()