大模型训练实战：从单卡到分布式集群的避坑指南（PyTorch DDP + Megatron-LM）

大模型训练实战：从单卡到分布式集群的避坑指南（PyTorch DDP + Megatron-LM）

当你第一次尝试将一个大模型塞进单张GPU时，那种显存溢出的报错信息，几乎是每个深度学习开发者必经的“成人礼”。这不仅仅是代码问题，更是我们思考方式转变的开始：从“如何让模型跑起来”到“如何让模型高效、稳定地在大规模集群上跑起来”。今天，我们不谈空洞的理论，只聚焦于实战。我将结合PyTorch DDP和Megatron-LM这两个核心工具，带你一步步走过从单卡调试到分布式集群部署的完整路径，分享那些在文档里找不到、却能让项目顺利上线的关键技巧和排坑经验。

1. 单机单卡：一切开始的基石与显存困境

在冲向分布式之前，我们必须确保模型在单卡环境下是健康且可运行的。这听起来简单，但很多分布式问题其实根源于此。

1.1 构建一个可复现的单卡训练循环

别急着用高级的Trainer封装，先从一个最朴素的训练循环开始。这能让你对每一个计算步骤、每一份显存消耗都了如指掌。

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

# 一个简单的自定义数据集
class DummyDataset(Dataset):
    def __len__(self):
        return 1000
    def __getitem__(self, idx):
        # 模拟输入，例如 seq_len=512, hidden_size=768
        return torch.randn(512, 768), torch.randint(0, 100, (512,))

# 一个简化版的Transformer块
class SimpleTransformerBlock(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=12)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )
        self.norm1 = nn.LayerNorm(hidden_size)
        self.norm2 = nn.LayerNorm(hidden_size)

    def forward(self, x):
        attn_output, _ = self.attention(x, x, x)
        x = x + attn_output
        x = self.norm1(x)
        mlp_output = self.mlp(x)
        x = x + mlp_output
        x = self.norm2(x)
        return x

model = SimpleTransformerBlock(hidden_size=768).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
dataloader = DataLoader(DummyDataset(), batch_size=4)

for batch_idx, (inputs, labels) in enumerate(dataloader):
    inputs, labels = inputs.cuda(), labels.cuda()
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = nn.CrossEntropyLoss()(outputs.view(-1, 100), labels.view(-1))
    loss.backward()
    optimizer.step()
    if batch_idx % 10 == 0:
        print(f"Step {batch_idx}, Loss: {loss.item():.4f}")
        # 关键：监控显存
        print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")

注意：在这个阶段，频繁使用 torch.cuda.memory_allocated() 和 torch.cuda.max_memory_allocated() 来记录峰值显存消耗。这是后续决定采用何种并行策略（数据并行、模型并行）的核心依据。

1.2 深入剖析显存占用：不仅仅是模型参数

很多人以为显存只被模型参数占用，其实不然。一次前向传播中，显存主要由以下几部分构成：