【大模型】- 扩展和分布式训练

扩展和分布式训练

训练更大的模型，跨多个设备

类型： 学习　｜　语言： Python　｜　🏷 前置：《预训练 Mini-GPT》（本系列第 4 篇）

学习目标

• 理解数据并行和模型并行
• 实现分布式数据并行 (DDP)
• 使用混合精度训练加速计算
• 处理检查点和容错
• 优化内存使用以训练更大的模型

为什么需要分布式训练

现代 LLM 有数十亿参数，需要在 TB 级数据上训练。单个 GPU 无法处理：

• 内存：模型参数、梯度和优化器状态需要比单个 GPU 更多的内存
• 时间：训练需要数周或数月
• 数据：数据量太大，无法在合理时间内处理

并行策略

数据并行

数据并行是将数据分成多份，在多个设备上复制模型：

GPU 0: Model Copy → Batch 0 → Gradients 0 →┐
GPU 1: Model Copy → Batch 1 → Gradients 1 → ├→ Average → Update
GPU 2: Model Copy → Batch 2 → Gradients 2 → │
GPU 3: Model Copy → Batch 3 → Gradients 3 →┘

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

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train_ddp(rank, world_size):
    setup(rank, world_size)
    
    model = GPT(config).to(rank)
    model = DDP(model, device_ids=[rank])
    
    # 训练循环
    for iter in range(max_iters):
        xb, yb = get_batch("train")
        xb, yb = xb.to(rank), yb.to(rank)
        
        logits, loss = model(xb, yb)
        loss.backward()
        
        # DDP 自动平均梯度
        optimizer.step()
        optimizer.zero_grad()
    
    cleanup()

模型并行

模型并行是将模型分成多份：

GPU 0: Layers 0-11 → Hidden States →
GPU 1: Layers 12-23 → Hidden States →
GPU 2: Layers 24-35 → Output

class ModelParallel(nn.Module):
    def __init__(self, model, device_map):
        super().__init__()
        self.devices = list(device_map.keys())
        
        # 分割模型
        self.layers = nn.ModuleDict()
        for name, layer in model.named_children():
            device = device_map[name]
            self.layers[name] = layer.to(device)
    
    def forward(self, x):
        for name, layer in self.layers.items():
            x = x.to(next(layer.parameters()).device)
            x = layer(x)
        return x

张量并行

张量并行是将单个层分割到多个设备：

class TensorParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.world_size = world_size
        
        # 分割权重
        self.weight = nn.Parameter(
            torch.randn(out_features // world_size, in_features)
        )
    
    def forward(self, x):
        # 每个设备计算部分输出
        output = x @ self.weight.t()
        
        # 所有设备收集完整输出
        dist.all_reduce(output)
        return output

混合精度训练

使用 FP16 或 BF16 加速计算并减少内存：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for iter in range(max_iters):
    xb, yb = get_batch("train")
    
    with autocast():
        logits, loss = model(xb, yb)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

梯度累积

当批次大小太大无法放入内存时：

accumulation_steps = 4  # 等效批次大小 = 32 * 4 = 128

for iter in range(max_iters // accumulation_steps):
    optimizer.zero_grad()
    
    for micro_step in range(accumulation_steps):
        xb, yb = get_batch("train")
        
        with autocast():
            logits, loss = model(xb, yb)
            loss = loss / accumulation_steps
        
        scaler.scale(loss).backward()
    
    scaler.step(optimizer)
    scaler.update()

梯度检查点

用计算换内存：

from torch.utils.checkpoint import checkpoint_sequential

class GPTWithCheckpointing(nn.Module):
    def forward(self, idx):
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(idx.size(1)))
        x = tok_emb + pos_emb
        
        # 使用梯度检查点
        x = checkpoint_sequential(self.blocks, segments=4, input=x)
        
        x = self.ln_f(x)
        return self.head(x)

DeepSpeed ZeRO

分割优化器状态和梯度：

import deepspeed

ds_config = {
    "train_batch_size": 128,
    "gradient_accumulation_steps": 4,
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,  # 分割参数、梯度和优化器状态
        "offload_optimizer": {"device": "cpu"},
        "offload_param": {"device": "cpu"},
    }
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=ds_config
)

检查点和容错

def save_checkpoint(model, optimizer, epoch, loss):
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.module.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, f"checkpoint-{epoch}.pt")

def load_checkpoint(model, optimizer, path):
    checkpoint = torch.load(path)
    model.module.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint["epoch"]

内存优化技巧

# 1. 梯度减少精度
torch.set_float32_matmul_precision("high")

# 2. 使用 flash attention
from torch.nn.functional import scaled_dot_product_attention

# 3. 分页注意力（vLLM）
# 使用 vLLM 库处理长序列

# 4. 卸载到 CPU
model = model.to("cpu")
optimizer = DeepSpeedCPUAdam(model.parameters())

监控分布式训练

# 使用 wandb 或 tensorboard
import wandb

wandb.init(project="llm-training", config=config)

for iter in range(max_iters):
    # ... 训练代码 ...
    
    if iter % log_interval == 0:
        wandb.log({
            "train/loss": loss.item(),
            "train/learning_rate": optimizer.param_groups[0]["lr"],
            "train/gradient_norm": gradient_norm,
        })

总结

分布式训练使用多种并行策略和内存优化技术。关键概念包括数据并行、模型并行、混合精度和梯度累积。选择正确的策略取决于模型大小、数据量和硬件约束。

下一步

下一课将介绍指令微调（SFT），使模型遵循人类指令。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

📝 自我检查（课程配套测验）

• 第1题 · 学前
• 第2题 · 学前
• 第3题 · 学后
• 第4题 · 学后
• 第5题 · 学后

Q1（学前） 7B 参数模型在 FP16 下仅权重需要多少 VRAM？

A. 7 GB
B. 14 GB
C. 28 GB
D. 56 GB

答案： B　解析： FP16 下每个参数 2 字节。70 亿 × 2 字节 = 14 GB。加上 Adam 优化器状态（2 份）和梯度，训练总内存约 56 GB，尚未计入激活。

Q2（学前） 分布式训练中使用的三种并行类型是什么？

A. CPU、GPU 和 TPU 并行
B. 数据并行、张量并行和流水线并行
C. batch、序列和 token 并行
D. 前向、反向和优化器并行

答案： B　解析： 数据并行在每张 GPU 上复制模型并分割数据。张量并行将单层拆到多张 GPU。流水线并行将模型层分成阶段分布到多张 GPU。

Q3（学后） FSDP（全分片数据并行）相比标准 DDP 做了什么？

A. 使用不同优化器
B. 在 GPU 间分片模型参数、梯度和优化器状态，而非每张 GPU 复制完整模型
C. 处理数据更快
D. 支持更多 GPU

答案： B　解析： 标准 DDP 在每张 GPU 复制整个模型（浪费）。FSDP 在 GPU 间分片参数，每张只持有一部分。参数按需聚合用于计算，之后释放。

Q4（学后） DeepSpeed ZeRO Stage 3 是什么？

A. 一种量化方法
B. 在 GPU 间分区优化器状态、梯度和模型参数，实现最大内存效率
C. 一种学习率调度
D. 数据预处理流水线

答案： B　解析： ZeRO Stage 1 分片优化器状态，Stage 2 增加梯度分片，Stage 3 增加参数分片。Stage 3 内存节省最大，可训练远超单 GPU 内存的模型。

Q5（学后） 为什么数据并行训练需要梯度同步？

A. 防止过拟合
B. 每张 GPU 在不同数据上计算梯度；跨 GPU 平均梯度确保所有副本一致更新
C. 减少内存
D. 加速前向传播

答案： B　解析： 数据并行中，每张 GPU 处理不同 batch 并计算不同梯度。AllReduce 在所有 GPU 上平均这些梯度，使每个副本应用相同更新并保持同步。