第25课：分布式训练

📖 本课概述

大模型需要大算力，分布式训练让我们利用多GPU/多机器加速训练。本课理解数据并行和模型并行。

🖥️ 一、为什么需要分布式

📐 二、并行策略

2.1 数据并行 (DP/DDP)

2.2 模型并行

模型太大放不进单GPU时，将模型切分到多个GPU上。

• 流水线并行：按层切分，像流水线一样工作
• 张量并行：单层内切分矩阵运算（Megatron-LM）

2.3 PyTorch DDP实现

import torch
import torch.nn as nn

# 分布式训练模拟
print("=== 数据并行原理 ===")

# 模拟4个GPU的数据并行
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(100, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(),
            nn.Linear(128, 10),
        )
    def forward(self, x):
        return self.net(x)

model = SimpleModel()
total_data = 256
num_gpus = 4
batch_per_gpu = total_data // num_gpus

# 模拟各GPU计算梯度
X_full = torch.randn(total_data, 100)
y_full = torch.randint(0, 10, (total_data,))

gradients = []
for gpu_id in range(num_gpus):
    start = gpu_id * batch_per_gpu
    end = start + batch_per_gpu
    
    model.zero_grad()
    loss = nn.CrossEntropyLoss()(model(X_full[start:end]), y_full[start:end])
    loss.backward()
    
    # 收集各GPU的梯度
    grad_dict = {name: p.grad.clone() for name, p in model.named_parameters()}
    gradients.append(grad_dict)

# AllReduce：平均梯度
print("梯度AllReduce（平均）:")
for name, p in model.named_parameters():
    avg_grad = sum(g[name] for g in gradients) / num_gpus
    grad_diff = max((g[name] - avg_grad).abs().max().item() for g in gradients)
    print(f"  {name}: 各GPU梯度差异最大={grad_diff:.6f}")

# 单GPU vs 多GPU等价性
model_single = SimpleModel()
loss_single = nn.CrossEntropyLoss()(model_single(X_full), y_full)
loss_single.backward()

# 对比
print(f"\n单GPU loss: {loss_single.item():.6f}")
for gpu_id in range(num_gpus):
    start = gpu_id * batch_per_gpu
    end = start + batch_per_gpu
    loss_gpu = nn.CrossEntropyLoss()(model(X_full[start:end]), y_full[start:end])
    print(f"GPU{gpu_id} loss: {loss_gpu.item():.6f}")

# 模型并行概念
print("\n=== 模型并行示意 ===")
class PipelineParallel(nn.Module):
    def __init__(self):
        super().__init__()
        # GPU0的前半部分
        self.stage0 = nn.Sequential(nn.Linear(100, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU())
        # GPU1的后半部分
        self.stage1 = nn.Sequential(nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10))
    
    def forward(self, x):
        x = self.stage0(x)  # GPU0
        x = self.stage1(x)  # GPU1
        return x

pp_model = PipelineParallel()
x = torch.randn(8, 100)
out = pp_model(x)
print(f"流水线并行: {list(x.shape)} → stage0 → stage1 → {list(out.shape)}")
print(f"Stage0参数: {sum(p.numel() for p in pp_model.stage0.parameters()):,}")
print(f"Stage1参数: {sum(p.numel() for p in pp_model.stage1.parameters()):,}")

🔬 深入拓展

核心要点回顾

本课的核心知识点构成了训练技巧学习的重要一环。让我们从更高维度审视这些知识之间的关系：

• 理论层面：理解数学原理是正确应用的前提，每个公式背后都有直觉
• 工程层面：PyTorch提供了高效实现，但理解底层原理才能正确调试
• 实践层面：没有银弹，不同场景需要不同的策略组合

常见陷阱与最佳实践

与前后课程的联系

深度学习是一个整体——每一课的知识都不是孤立的：

• 优化器（第5课）的选择影响正则化（第6课）的效果
• 初始化（第7课）决定了训练能否开始
• CNN（第8-14课）的成功依赖训练技巧（第22-28课）
• Transformer（第20-21课）结合了注意力和位置编码
• 实战项目（第29-35课）综合运用所有知识

📖 延伸阅读

📝 练习

💡 分布式实践

并行策略

• 数据并行最简单
• 混合精度减少显存
• 梯度累积模拟大batch
• 注意batch norm的行为

常见错误与解决方案

性能优化建议

🔬 深度拓展：从理论到实践

核心概念网络

理解本课内容需要将其置于更大的知识网络中。每个核心概念都不是孤立的，它们相互关联、相互支撑：

• 数学基础：线性代数、微积分、概率论是理解所有深度学习算法的根基
• 计算基础：自动微分、GPU并行计算让理论变为现实
• 工程基础：数据处理、模型设计、训练技巧的有机结合
• 领域知识：CV、NLP、语音等领域的先验知识指导模型设计

关键洞察与直觉

建立正确的直觉比记住公式更重要：

与其他课程的关联

本课作为训练技巧阶段的核心内容，与前后课程紧密关联：

面试常见问题

进阶学习路径

掌握本课内容后，可以通过以下方式继续深入：

• 阅读本课相关的一手论文，理解原始动机和推导
• 复现论文的实验结果，验证自己的理解
• 尝试改进方法，思考为什么现有方法有局限
• 在真实项目或Kaggle竞赛中应用
• 关注该方向的最新进展（arXiv、顶级会议）

📐 补充：知识体系梳理

核心公式速查

将本课涉及的关键公式整理如下，方便回顾和记忆。理解每个公式背后的直觉比死记硬背更重要——试着用自然语言解释每个公式在做什么。

易混淆概念辨析

代码模板

以下是本课核心概念的标准PyTorch实现模板，可以直接用于实际项目：

• 定义模型时，继承nn.Module，实现__init__和forward
• 训练循环：zero_grad → forward → loss → backward → step
• 验证时使用model.eval()和torch.no_grad()
• 保存模型用torch.save(model.state_dict())
• 加载模型用model.load_state_dict(torch.load())