第74课 - GPU 训练与分布式训练基础

学习目标

完成本课后，你将能够：

理解GPU在深度学习训练中为何能大幅提升速度。
掌握在单机单GPU（或单个计算卡）上进行模型训练的基本代码。
了解使用多GPU进行数据并行训练的基本原理与代码实现。
初步认识分布式训练的概念，特别是数据并行（Data Parallelism）策略。

核心概念

1. 为什么需要GPU？

想象一下，你的模型就像一个拥有成千上万次计算任务的大工厂。

CPU（中央处理器）：像一位技艺高超但只有几双手的工匠。它擅长处理复杂的逻辑和串行任务，但同时处理的任务数量有限。
GPU（图形处理器）：像一支由成千上万名普通工人组成的大军。单个工人的能力可能不如CPU工匠，但胜在人数众多，可以同时进行大量简单、重复的计算。

深度学习的训练核心是大量的矩阵乘法、卷积等并行计算。这正是GPU大军的用武之地。使用GPU，可以将原本需要数天或数周的训练时间，缩短到几小时甚至几分钟。

2. 单GPU训练

这是最基础的使用方式。我们需要做两件事：

指定计算设备：明确告诉程序，我们的模型和数据要在GPU上运行。
迁移数据与模型：将模型参数和每次训练的输入数据，从内存（CPU）搬到显存（GPU）上。

在PyTorch中，核心命令是 .to(device) 或 .cuda()。

3. 多GPU与分布式训练

当单块GPU的显存或算力仍然不够时，我们可以使用多块GPU。

数据并行（Data Parallelism） 是最简单、最常用的策略：

将一个大的数据批次（Batch）拆分成若干个小的子批次。
将相同的模型副本放到不同的GPU上。
每个GPU用自己分到的小数据子集独立进行前向和反向传播，计算出梯度。
最后，将所有GPU上计算出的梯度汇总平均，并用这个平均梯度来更新主模型（或所有模型副本）的参数。

这就像是把一叠文件复印多份，分给多个团队同时处理，最后汇总报告。

代码示例

示例1：单GPU训练配置

import torch
import torch.nn as nn
import torch.optim as optim

# 1. 设置设备：优先使用GPU，如果没有则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 2. 定义一个简单的模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(10, 2)

    def forward(self, x):
        return self.fc(x)

# 3. 实例化模型，并将其迁移到GPU
model = SimpleNet().to(device)
print(f"模型已加载到 {device}")

# 4. 创建模拟数据和标签，并迁移到GPU
input_data = torch.randn(5, 10).to(device)  # 形状为(5, 10)的随机数据
labels = torch.randint(0, 2, (5,)).to(device) # 形状为(5,)的随机标签

# 5. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 6. 进行一步训练
optimizer.zero_grad()          # 清空梯度
outputs = model(input_data)    # 前向传播（在GPU上）
loss = criterion(outputs, labels) # 计算损失
loss.backward()                # 反向传播
optimizer.step()               # 更新参数

print(f"训练完成，损失: {loss.item()}")

示例2：多GPU数据并行训练 (使用`DataParallel`)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import DataParallel

# 检查可用GPU数量
num_gpus = torch.cuda.device_count()
print(f"检测到 {num_gpus} 个GPU")

# 定义模型
model = SimpleNet() # 使用示例1中的SimpleNet

# 如果有多于一个GPU，使用DataParallel包装模型
if num_gpus > 1:
    print(f"使用 DataParallel 进行多GPU训练")
    model = DataParallel(model)

# 将模型移动到主GPU（通常是‘cuda:0’）
device = torch.device("cuda:0")
model = model.to(device)

# 训练过程与单GPU几乎完全相同，PyTorch会自动处理数据分配和梯度同步
input_data = torch.randn(16, 10).to(device) # 注意：批次大小可以更大
labels = torch.randint(0, 2, (16,)).to(device)

optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

optimizer.zero_grad()
outputs = model(input_data)    # DataParallel 会自动拆分输入
loss = criterion(outputs, labels)
loss.backward()                # 梯度会自动同步
optimizer.step()

print("多GPU训练一步完成。")

示例3：分布式数据并行基础 (使用`DistributedDataParallel`)

这是更现代、更高效的方式，适用于多机多卡或单机多卡复杂场景。

# 注意：此示例需要使用 `torch.distributed.launch` 或 `torchrun` 命令启动
# 例如：torchrun --nproc_per_node=2 你的脚本名.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def train(rank, world_size):
    # 1. 初始化进程组（所有进程需要共同完成）
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    device_id = rank % torch.cuda.device_count()
    torch.cuda.set_device(device_id)

    # 2. 创建模型并移至当前GPU
    model = SimpleNet().to(device_id)

    # 3. 用DDP包装模型，这是关键！
    ddp_model = DDP(model, device_ids=[device_id])

    # 4. 创建优化器和数据
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()
    # 注意：在分布式训练中，每个进程应加载不同的数据子集（由DistributedSampler处理）
    # 这里为简化，使用相同数据，但实际项目中请使用Sampler。

    # 5. 训练循环
    for epoch in range(2):
        optimizer.zero_grad()
        input_data = torch.randn(4, 10).to(device_id)
        labels = torch.randint(0, 2, (4,)).to(device_id)

        outputs = ddp_model(input_data)
        loss = criterion(outputs, labels)
        loss.backward()  # 梯度会在所有进程间自动同步
        optimizer.step()

        print(f"进程 {rank}, 周期 {epoch}, 损失: {loss.item()}")

    # 6. 清理
    dist.destroy_process_group()

if __name__ == "__main__":
    # 世界大小（总进程数，这里假设为2）
    world_size = 2
    # 此示例需要由外部启动器（如torchrun）启动，这里仅为说明结构
    # 实际运行时，每个进程会独立执行此脚本，rank由启动器分配
    print("请使用 `torchrun --nproc_per_node=2 此脚本名.py` 来启动分布式训练。")

实践练习

练习1：设备检测与迁移

要求：编写代码，检测当前环境是否有CUDA GPU。如果有，创建一个形状为(3, 4)的随机张量，并将其移动到该GPU上。打印出张量的设备信息。
预期输出示例：张量设备: cuda:0 或 张量设备: cpu (如果无GPU)

练习2：尝试多GPU训练

要求：假设你有2块GPU，修改示例1中的SimpleNet模型，使用DataParallel使其支持双卡训练。创建一个批次大小为8的输入，并完成一次前向传播。观察并解释：模型输出张量的形状是否改变？
提示：关注DataParallel在不同维度上的拆分与合并行为。

练习3：理解分布式数据并行

要求：解释以下问题（无需代码）：
1. DataParallel (DP) 和 DistributedDataParallel (DDP) 在实现“数据并行”的核心思想上有什么共同点？
2. 为什么在实践中，DDP通常比DP更高效、更推荐使用？（提示：从通信方式和Python GIL锁的角度思考）

常见错误

忘记迁移模型和数据：这是最常犯的错误。尝试在GPU上运行一个在CPU上定义的模型，或反之，会导致RuntimeError: Expected all tensors to be on the same device。
```
# 错误示例
model = SimpleNet()  # 在CPU上
input_gpu = torch.randn(2, 10).cuda()  # 数据在GPU上
output = model(input_gpu)  # 报错！模型参数在CPU，输入在GPU
```
混淆批次维度与设备并行：使用DataParallel时，默认会在批次维度（dim=0） 上拆分数据。如果你的模型输入张量没有批次维度（如某些图像任务），需要手动调整。
分布式训练中未正确初始化进程组：使用DDP前必须调用dist.init_process_group，否则后续操作会失败。
忽略随机种子同步：在分布式训练中，为了保证各进程初始化模型权重一致，需要在初始化模型前同步随机种子。
```
# 在每个进程中
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)
```
过度依赖默认设备：明确用.to(device)指定设备是好习惯，这使你的代码在GPU和CPU环境间具有更好的可移植性。

小结

GPU的力量：深度学习训练的计算密集型特性，与GPU大规模并行架构高度契合，是速度提升的关键。
单GPU训练：核心是.to(device)，将模型和数据放在同一计算设备上。
多GPU训练：DataParallel提供了最简单的多卡训练方式，但效率并非最优。它在单进程内管理多个GPU。
分布式训练基础：DistributedDataParallel是更强大、更灵活的多机多卡训练方案。它通过多进程来避免GIL限制，通信效率更高，是工业级训练的标准选择。
核心思想：数据并行的本质是“分而治之”——拆分数据，同步梯度。理解了这一点，就能更好地掌握各种并行训练策略。

练习编辑器

rust