第21课：PyTorch｜GPU多卡训练与分布式训练基础【让多卡并行成为你的加速引擎】

文章目录

- [📖 课前导读](#📖 课前导读)
- - 为什么需要多GPU训练？
- 一、知识原理：GPU加速与数据并行基础
- - [1.1 单GPU训练的标准流程](#1.1 单GPU训练的标准流程)
  - [1.2 显存占用构成](#1.2 显存占用构成)
  - [1.3 数据并行的核心机制](#1.3 数据并行的核心机制)
  - [1.4 两种多卡API的对比](#1.4 两种多卡API的对比)
- 二、环境搭建与准备
- - [2.1 检查可用GPU](#2.1 检查可用GPU)
  - [2.2 必要的环境变量](#2.2 必要的环境变量)
  - [2.3 单机多卡NCCL通信配置（可选优化）](#2.3 单机多卡NCCL通信配置（可选优化）)
- 三、代码实战：从单卡到多卡的演进
- - [3.1 单GPU基础训练（基线）](#3.1 单GPU基础训练（基线）)
  - [3.2 DataParallel（DP）简易多卡](#3.2 DataParallel（DP）简易多卡)
  - [3.3 DistributedDataParallel（DDP）标准用法](#3.3 DistributedDataParallel（DDP）标准用法)
- 四、进阶：分布式启动与调试
- - [4.1 使用torchrun启动DDP](#4.1 使用torchrun启动DDP)
  - [4.2 Backend的选择](#4.2 Backend的选择)
  - [4.3 分布式训练中的同步BN](#4.3 分布式训练中的同步BN)
- 五、显存优化：突破单卡容量瓶颈
- - [5.1 梯度累积（Gradient Accumulation）](#5.1 梯度累积（Gradient Accumulation）)
  - [5.2 混合精度训练（Automatic Mixed Precision, AMP）](#5.2 混合精度训练（Automatic Mixed Precision, AMP）)
  - [5.3 梯度检查点（Gradient Checkpointing）](#5.3 梯度检查点（Gradient Checkpointing）)
  - [5.4 显存管理最佳实践](#5.4 显存管理最佳实践)
- 六、难点解析：常见报错与解决方案
- - [6.1 CUDA out of memory](#6.1 CUDA out of memory)
  - [6.2 DDP通信超时](#6.2 DDP通信超时)
  - [6.3 DataLoader中shuffle与DistributedSampler同时使用](#6.3 DataLoader中shuffle与DistributedSampler同时使用)
  - [6.4 DataParallel与DDP混用报错](#6.4 DataParallel与DDP混用报错)
- 七、实战项目：使用DDP加速图像分类训练
- 八、课后总结
- - 核心知识点速查表
  - 最佳实践总结
  - 检查清单
- 九、课后作业
- - [作业1：DataParallel vs DDP性能对比](#作业1：DataParallel vs DDP性能对比)
  - 作业2：DDP训练脚本编写
  - 作业3：梯度累积效果验证
  - 作业4：显存优化组合实验
  - 作业5：分布式调试实践
- 十、下一课预告
[🔗《精讲25课｜PyTorch 从入门到精通》系列课程导航](#🔗《精讲25课｜PyTorch 从入门到精通》系列课程导航)

📖 课前导读

为什么需要多GPU训练？

在第11课到第20课中，我们一直在单张GPU上进行模型训练。对于MNIST这样的小数据集，单卡训练完全够用。但当模型和数据规模增长时，单GPU会面临两大瓶颈：

显存容量限制：一张消费级显卡（如RTX 3060 12GB）难以训练大模型。以GPT-2（1.5B参数）为例，仅模型参数（FP32）就需要约6GB，加上中间激活值、梯度和优化器状态，总显存需求远超单卡容量。
训练时间过长：即使显存放得下，单卡的算力也可能不足以在合理时间内完成训练。在ImageNet上训练ResNet-50约需14小时（单V100），而在CIFAR-10上训练ViT-Large可能需数天。

多GPU训练的核心思想是：将同一份任务拆解到多张GPU上并行完成。根据拆解方式不同，主要分为：

数据并行（Data Parallelism） ：每张GPU持有完整的模型副本，处理不同的数据子集。梯度在GPU间同步后统一更新。这是最常用的多卡训练范式。
模型并行（Model Parallelism） ：将模型的不同层/模块分布到不同GPU上，适合单卡无法完整容纳的巨型模型。
混合并行（Hybrid Parallelism） ：结合上述两种方式，常见于大模型训练。

本课聚焦于数据并行 ------这是绝大多数深度学习项目最直接有效的多卡加速方案。PyTorch提供了两个API来实现数据并行：nn.DataParallel（简单但有限制）和nn.parallel.DistributedDataParallel（更复杂但性能更优，官方强烈推荐）。

💡 核心认知 ：PyTorch官方明确指出，即使只在单机多卡上训练，DistributedDataParallel（DDP）也是比DataParallel更推荐的选择。DDP采用多进程架构，避免了Python多线程的GIL限制，通信效率更高，扩展性更好。

一、知识原理：GPU加速与数据并行基础

1.1 单GPU训练的标准流程

在讨论多GPU之前，先回顾单卡训练的核心操作：

python 复制代码

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

for epoch in range(epochs):
    for batch in dataloader:
        images, labels = batch[0].to(device), batch[1].to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

这一模式中，to(device)将模型和数据迁移到GPU，所有计算在GPU上完成。

1.2 显存占用构成

理解显存占用是优化多卡训练的前提。显存主要占用来源：

占用模块	占比	说明
模型参数	30%-60%	所有可训练权重
中间激活值	20%-50%	前向传播中保存的中间结果
优化器状态	10%-30%	Adam需要维护动量和二阶矩

以ResNet-50在batch size=32时的单卡训练为例，模型参数约98MB，但中间激活值可达200MB以上。训练大模型时优化器状态（Adam的2个额外状态）造成的显存开销甚至可以加倍参数量占用。

1.3 数据并行的核心机制

数据并行（Data Parallelism）的流程：将一个大batch按世界大小切分到多张GPU上 → 每张GPU独立完成前向和反向 → 各GPU的梯度通过All-Reduce通信同步 → 每张GPU用汇总后的梯度独立更新参数。

DDP在工作原理上有几个关键设计：

构造时广播初始参数：rank 0进程将初始状态广播给所有其他进程，确保所有模型副本起始状态一致。
Reducer桶组织 ：参数梯度按大小组织成桶，每个桶内的梯度在一次通信调用中完成同步，减少通信次数，bucket_cap_mb参数（默认25MB）可调节桶容量。
反向传播时异步通信：DDP在反向传播过程中就开始进行梯度同步，与计算重叠，提高效率。
All-Reduce操作：所有进程互相交换梯度，最终每个进程都获得梯度的总和（或均值），确保参数更新后所有副本保持一致。

1.4 两种多卡API的对比

维度	DataParallel (DP)	DistributedDataParallel (DDP)
并行方式	单进程多线程	多进程（每个GPU一个进程）
通信开销	主GPU收集所有梯度再广播，有一定开销	All-Reduce并行通信，效率更高
GIL限制	受Python GIL影响	无GIL问题
主GPU负载	主GPU显存占用更高（收集所有梯度）	各GPU显存占用均衡
节点扩展	仅限单机	支持多机分布式
学习率调整	有效batch size需同步调整	无需调整（batch size不变）
官方推荐	不推荐新项目使用	强烈推荐

PyTorch官方建议：即使只在单机上使用多GPU，也应该优先选择DDP。DataParallel的通信开销和主GPU负载不均衡问题在扩展到多个GPU时尤其明显。

二、环境搭建与准备

2.1 检查可用GPU

python 复制代码

import torch

print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {torch.cuda.is_available()}")
print(f"GPU数量: {torch.cuda.device_count()}")

for i in range(torch.cuda.device_count()):
    print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
    print(f"  计算能力: {torch.cuda.get_device_capability(i)}")
    total_memory = torch.cuda.get_device_properties(i).total_memory / 1e9
    print(f"  显存总量: {total_memory:.2f} GB")

硬件要求：

多GPU训练需要多张NVIDIA显卡（GTX 10系以上，建议6GB以上显存）
DDP需要在GPU之间进行NVLink或PCIe通信，建议同一PCIe交换机下的GPU互联
单卡训练（无多卡）可在任意NVIDIA显卡上运行，代码也完全兼容

2.2 必要的环境变量

DDP初始化依赖于环境变量或init_method参数。使用torchrun启动时会自动设置以下环境变量：

环境变量	含义
`MASTER_ADDR`	主节点IP地址（单机多卡可设为localhost）
`MASTER_PORT`	主节点通信端口（确保未被占用）
`WORLD_SIZE`	全局进程总数
`RANK`	当前进程的全局编号（0到world_size-1）
`LOCAL_RANK`	当前进程在本节点内的编号

2.3 单机多卡NCCL通信配置（可选优化）

使用NCCL后端时，可通过环境变量优化通信性能：

bash 复制代码

# 禁用NCCL的InfiniBand（仅用PCIe通信时）
export NCCL_IB_DISABLE=1
# 设置NCCL调试级别（调试通信问题）
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,NET
# 设置socket网络接口（多机训练时）
export NCCL_SOCKET_IFNAME=eth0

三、代码实战：从单卡到多卡的演进

3.1 单GPU基础训练（基线）

以CIFAR-10分类任务为基线，展示单卡训练的标准流程：

python 复制代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import time
import os

# 基础配置
batch_size = 64
epochs = 10
learning_rate = 0.001

# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

# 模型（简单CNN示例）
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        return self.fc2(x)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

start_time = time.time()
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}")

elapsed = time.time() - start_time
print(f"单卡训练耗时: {elapsed:.2f} 秒")

3.2 DataParallel（DP）简易多卡

DataParallel是最简单的多GPU方案，只需一行代码包装模型：

python 复制代码

from torch.nn import DataParallel

# 原模型实例化（先移到GPU）
model = SimpleCNN().to(device)

# 如果有多张GPU，用DataParallel包装
if torch.cuda.device_count() > 1:
    print(f"使用 {torch.cuda.device_count()} 张GPU")
    model = DataParallel(model)

# 其余代码与单卡完全相同
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练代码无需修改
for epoch in range(epochs):
    model.train()
    # ... 与单卡相同

DataParallel的典型性能限制：

python 复制代码

# DataParallel的主GPU显存占用测试
def test_dp_memory_imbalance():
    model = SimpleCNN().cuda()
    model = DataParallel(model)
    
    # 在训练前检查各GPU显存占用
    for i in range(torch.cuda.device_count()):
        allocated = torch.cuda.memory_allocated(i) / 1024**2
        print(f"GPU {i}: 已分配 {allocated:.2f} MB")
    # 主GPU显存占用明显高于其他GPU（因为主GPU收集梯度）

DP的限制总结：

主GPU负载不均。主GPU需要聚合其他GPU的梯度并广播，显存占用比从属GPU高；DP的速度扩展性不如DDP，且性能随GPU数量增加而下降。PyTorch官方已不再推荐在新项目中使用DataParallel。
DP使用多线程（而非多进程），Python的GIL限制了在多GPU环境中的性能。因此在多GPU训练中，DDP是官方推荐的标准方案。

3.3 DistributedDataParallel（DDP）标准用法

DDP采用多进程架构，每个进程管理一张GPU，代码组织方式与单卡训练有所不同。

python 复制代码

# ddp_train.py
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

def setup(rank, world_size):
    """初始化分布式进程组"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    # 1. 初始化进程组
    setup(rank, world_size)
    
    # 2. 将模型放到对应GPU上
    device = torch.device(f"cuda:{rank}")
    model = SimpleCNN().to(device)
    ddp_model = DDP(model, device_ids=[rank])
    
    # 3. 数据加载（使用DistributedSampler）
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])
    dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    # 注意：不需要手动乘以world_size，DDP自动切分数据
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=4)
    
    # 4. 训练循环
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(ddp_model.parameters(), lr=0.001)
    
    for epoch in range(10):
        # 关键：每个epoch需要调用set_epoch保证不同进程的shuffle不同
        sampler.set_epoch(epoch)
        ddp_model.train()
        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = ddp_model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
        
        if rank == 0:  # 只在主进程打印日志
            print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
    
    cleanup()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

DDP代码要点：

每个进程独立实例化模型，移动到对应GPU后包装为DDP
使用DistributedSampler替代shuffle=True，每个epoch需调用set_epoch以保证各进程的随机性不同
日志只在rank 0进程打印，避免重复输出

四、进阶：分布式启动与调试

4.1 使用torchrun启动DDP

从PyTorch 1.10开始，torchrun取代了torch.distributed.launch作为官方推荐的分布式启动器。torchrun能自动设置环境变量（RANK、WORLD_SIZE等），并支持worker故障自动重启。它已完全取代了标记为弃用的torch.distributed.launch。

python 复制代码

# ddp_torchrun.py
import torch
import torch.distributed as dist
import os

def main():
    # torchrun自动设置的环境变量
    local_rank = int(os.environ['LOCAL_RANK'])
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    
    # 初始化进程组（无需手动设置MASTER_ADDR/PORT，torchrun已处理）
    dist.init_process_group(backend='nccl')
    
    # 设置当前进程使用的GPU
    torch.cuda.set_device(local_rank)
    device = torch.device(f'cuda:{local_rank}')
    
    model = SimpleCNN().to(device)
    ddp_model = DDP(model, device_ids=[local_rank])
    
    # ... 训练代码
    
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

启动命令（单机8卡）：

bash 复制代码

torchrun --nproc_per_node=8 ddp_torchrun.py

多机8卡分布式训练（2个节点，每节点4卡）：

bash 复制代码

# 主节点（192.168.1.1）
torchrun --nnodes=2 --nproc_per_node=4 --rdzv_backend=c10d \
         --rdzv_endpoint=192.168.1.1:29400 ddp_torchrun.py

# 从节点自动连接，无需额外参数
torchrun --nnodes=2 --nproc_per_node=4 --rdzv_backend=c10d \
         --rdzv_endpoint=192.168.1.1:29400 ddp_torchrun.py

4.2 Backend的选择

Backend	适用场景	说明
`nccl`	NVIDIA GPU多卡	最推荐，针对NVIDIA GPU优化，支持All-Reduce等高性能通信
`gloo`	CPU分布式、多机互联	跨平台，兼容性好，但GPU场景性能低于NCCL
`mpi`	HPC环境	需要MPI支持，不常用

最佳实践：单机多卡或多机多卡（NVIDIA GPU）始终选用NCCL后端。仅在NCCL不可用时（如AMD GPU混合环境）才使用Gloo。

4.3 分布式训练中的同步BN

对于小batch size训练，标准BatchNorm的统计量在不同进程间独立计算，可能导致训练不稳定。分布式同步BN在跨卡GPU间同步均值和方差，可有效改善此问题。PyTorch 1.11+提供了torch.nn.SyncBatchNorm。

python 复制代码

# 将模型中所有BN层转换为SyncBatchNorm
model = SimpleCNN()
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
ddp_model = DDP(model.to(device), device_ids=[rank])

五、显存优化：突破单卡容量瓶颈

5.1 梯度累积（Gradient Accumulation）

当模型太大、单卡batch size无法达到目标大小时，梯度累积是一种有效手段。其原理是将多个batch的梯度累加起来，每accumulation_steps步才进行一次参数更新，从而模拟更大的batch size。

python 复制代码

# 梯度累积示例
accumulation_steps = 4  # 模拟batch size扩大4倍
optimizer.zero_grad()

for i, (images, labels) in enumerate(train_loader):
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # 梯度平均，保持有效学习率不变
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

梯度累积的注意事项：

需要将原始loss除以累积步数，以保持梯度量级与直接使用大batch一致
如果配合BatchNorm使用，BN的统计量仍然是基于小batch计算的，与实际训练大batch的分布可能略有偏差
在DDP环境下，梯度累积不会自动跨卡同步------每张卡独立累积，累积完成后各卡同时step，效果等价于全局batch size扩大

5.2 混合精度训练（Automatic Mixed Precision, AMP）

混合精度训练通过FP16与FP32混合计算，可在保持模型精度的同时大幅减少显存占用并加速训练。在ResNet-50训练中，显存占用可从9.8GB降至4.2GB，训练速度提升1.8倍。

python 复制代码

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # 梯度缩放器，防止FP16下梯度下溢

for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    optimizer.zero_grad()
    
    # 前向传播自动使用混合精度
    with autocast():
        outputs = model(images)
        loss = criterion(outputs, labels)
    
    # 反向传播：缩放损失→反向→优化器step前缩放梯度
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

AMP的硬件要求：支持Tensor Core的GPU（Volta架构及以上，如V100、A100、T4、RTX 30/40系列）效果最佳。在ResNet-50训练中，AMP可降低40%显存占用，同时保持模型精度。

5.3 梯度检查点（Gradient Checkpointing）

梯度检查点通过牺牲计算时间来换取显存空间，其核心思想是：在前向传播时只保存部分中间激活值，其余中间结果在反向传播时重新计算。对于深层网络，显存消耗可从O(n)降至O(√n)。在BERT-base中，启用检查点后显存占用可从24GB降至8GB，但计算时间会增加约20%。

python 复制代码

from torch.utils.checkpoint import checkpoint

class CheckpointModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = nn.Sequential(nn.Linear(1024, 1024), nn.ReLU())
        self.block2 = nn.Sequential(nn.Linear(1024, 1024), nn.ReLU())
    
    def forward(self, x):
        # 使用checkpoint包装block1，不保存中间激活值
        x = checkpoint(self.block1, x)
        x = self.block2(x)
        return x

梯度检查点的注意事项：

适用于深层网络（如Transformer），在每3-5个层设置一个检查点，平衡显存节省与计算开销
DDP环境下，检查点也会跟随进程独立管理，不受分布式影响

5.4 显存管理最佳实践

及时释放临时张量，使用del和torch.cuda.empty_cache()组合释放
推理时使用with torch.no_grad():禁用梯度记录，避免中间激活值累积
使用梯度检查点技术或混合精度训练（AMP）控制显存峰值

六、难点解析：常见报错与解决方案

6.1 CUDA out of memory

症状：RuntimeError: CUDA out of memory

原因分析：显存碎片化、他进程占用、模型过大、batch size过高。

解决步骤：

减小batch size是直接有效的缓解手段
尝试梯度累积：物理batch size减小，但有效batch size保持不变
检查是否有其他进程占用GPU：nvidia-smi查看
清理PyTorch缓存：torch.cuda.empty_cache()
如DDP模式下某些进程OOM，考虑启用梯度检查点以节省显存

6.2 DDP通信超时

症状：进程hang住或报NCCL timeout

解决方案：

bash 复制代码

export NCCL_DEBUG=INFO  # 打开NCCL调试日志
export NCCL_IB_DISABLE=1  # 禁用InfiniBand（仅使用PCIe时）
export NCCL_SOCKET_IFNAME=eth0  # 指定网卡
# 增加超时时间（单位：秒）
export NCCL_TIMEOUT=1800

6.3 DataLoader中shuffle与DistributedSampler同时使用

错误：DDP训练中既在DataLoader里设置shuffle=True，又使用DistributedSampler，会导致数据打乱出现重复或遗漏。

正确做法 ：使用DistributedSampler时必须设置shuffle=False，并在每个epoch调用sampler.set_epoch(epoch)以确保每个epoch数据顺序打乱且各进程独立。

python 复制代码

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, shuffle=False)
for epoch in range(epochs):
    sampler.set_epoch(epoch)  # 关键：改变随机种子
    for batch in dataloader:
        ...

6.4 DataParallel与DDP混用报错

症状：同时使用DataParallel和DDP包装模型时报错。

原因：模型被多层包装会导致梯度计算错乱。

解决方案：只使用一种并行方式。官方建议直接使用DDP。

七、实战项目：使用DDP加速图像分类训练

完整的单机多卡训练示例，对比单卡与多卡性能。

python 复制代码

# ddp_benchmark.py
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
import time

def train(rank, world_size, use_amp=False):
    # 初始化DDP
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    device = torch.device(f'cuda:{rank}')
    
    # 模型
    model = SimpleCNN().to(device)
    model = DDP(model, device_ids=[rank])
    
    # 数据（使用DistributedSampler）
    dataset = datasets.CIFAR10('./data', train=True, download=True,
        transform=transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ]))
    sampler = DistributedSampler(dataset, rank=rank, num_replicas=world_size, shuffle=True)
    
    # 有效batch_size：每个GPU上的实际batch_size
    per_gpu_batch = 64
    loader = DataLoader(dataset, batch_size=per_gpu_batch, sampler=sampler, num_workers=4)
    
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    
    # AMP设置
    scaler = torch.cuda.amp.GradScaler() if use_amp else None
    
    # 热身几个epoch（排除编译开销）
    epoch_start = time.time()
    for epoch in range(10):
        sampler.set_epoch(epoch)
        model.train()
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            
            if use_amp:
                with torch.cuda.amp.autocast():
                    outputs = model(images)
                    loss = criterion(outputs, labels)
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                outputs = model(images)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
        
        if rank == 0:
            print(f"Epoch {epoch+1} completed")
    
    if rank == 0:
        elapsed = time.time() - epoch_start
        print(f"{world_size} GPU(s), batch per GPU={per_gpu_batch}, 总耗时={elapsed:.2f}s")

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    # 对比：单卡 vs 多卡 vs 多卡+AMP
    for mode in ['single', 'multi', 'multi_amp']:
        if mode == 'single':
            # 单卡对比基准：直接调用train函数（需手动适配）
            pass
        else:
            use_amp = (mode == 'multi_amp')
            mp.spawn(train, args=(world_size, use_amp), nprocs=world_size, join=True)

八、课后总结

核心知识点速查表

概念	说明
`DataParallel`	单进程多线程，简单易用，但有性能限制
`DistributedDataParallel`	多进程，高效通信，推荐方案
`DistributedSampler`	DDP数据采样器，自动分配数据
`torchrun`	官方分布式启动器，替代`torch.distributed.launch`
`init_process_group`	初始化进程组，设置backend
梯度累积	模拟大batch size，降低显存峰值
AMP	混合精度，兼顾速度和显存
梯度检查点	以时间换空间，适合深层网络

最佳实践总结

新项目一律使用DDP：跳过DataParallel，直接从DDP开始
使用torchrun启动DDP：自动化设置环境变量，官方推荐
分布式训练中始终使用DistributedSampler配合set_epoch
日志只在rank 0进程打印：防止重复输出
显存不足优先尝试：减小batch size → 梯度累积 → 混合精度AMP → 梯度检查点
多机训练时确保各节点PyTorch和CUDA版本一致，网络互通

检查清单

理解单GPU训练的基本流程
知道如何使用DataParallel包装模型
理解DDP的多进程架构优势
能正确配置DistributedSampler
掌握torchrun启动分布式训练的方法
了解梯度累积、混合精度、梯度检查点的使用场景和代码实现
遇到显存不足能给出系统性排查和优化方案
能独立编写单机多卡的DDP训练脚本

九、课后作业

作业1：DataParallel vs DDP性能对比

在同一模型（如ResNet-18）和相同有效batch size下，分别用DataParallel和DDP训练CIFAR-10（5个epoch），测量训练总耗时和GPU利用率。对比两种方案在2卡、4卡下的加速比和显存占用均衡性。

作业2：DDP训练脚本编写

从零编写一个完整的DDP训练脚本，要求：

使用torchrun启动
支持混合精度训练（AMP）
支持检查点保存和恢复（保存模型、优化器、采样器状态）
日志和模型只在rank 0进程保存

作业3：梯度累积效果验证

在单卡显存受限情况下，分别用以下三种方式训练相同模型（目标batch size=256）：

直接batch_size=256（如果显存不足则跳过）
batch_size=64，梯度累积步数=4
batch_size=256，不累积（基线）
对比三种方式的收敛速度和最终准确率，分析梯度累积的利弊。

作业4：显存优化组合实验

使用ResNet-50在CIFAR-10上训练，测试以下配置下的显存峰值和训练速度：

FP32基线
仅AMP
仅梯度检查点
AMP + 梯度检查点组合

作业5：分布式调试实践

故意制造分布式训练中的常见错误（如未调用set_epoch、忘记init_process_group、使用不匹配的后端等），记录报错信息并写出解决方案。通过这种方式加深对DDP各组件依赖关系的理解。

十、下一课预告

第22课我们将学习PyTorch模型推理部署基础优化，内容包括：

训练模型转推理模型（model.eval()、torch.no_grad()）
模型量化压缩（PTQ动态/静态量化、QAT量化感知训练）
ONNX格式导出与推理优化
推理速度优化与批量推理
离线推理脚本编写

模型部署是将训练成果落地到生产环境的关键一步。学完第22课，你将能把训练好的模型应用于实际业务场景，并优化推理性能。

附录：分布式训练核心命令速查

操作	命令/代码
安装PyTorch（带CUDA）	`pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118`
检查GPU	`torch.cuda.device_count()`, `nvidia-smi`
DDP初始化	`dist.init_process_group("nccl", rank=rank, world_size=world_size)`
torchrun启动	`torchrun --nproc_per_node=4 train.py`
混合精度	`from torch.cuda.amp import autocast, GradScaler`
梯度检查点	`from torch.utils.checkpoint import checkpoint`
清理显存	`torch.cuda.empty_cache()`

🔗《精讲25课｜PyTorch 从入门到精通》系列课程导航

去订阅

🌟 感谢您耐心阅读到这里！

💡 如果本文对您有所启发欢迎：

👍 点赞📌 收藏 📤 分享给更多需要的伙伴。

🗣️ 期待在评论区看到您的想法, 共同进步。

🔔 关注我，持续获取更多干货内容～

🤗 我们下篇文章见～