第四章：深度学习训练优化——数据加载与分布式策略

大家好，我是你们的老朋友。今天咱们聊聊训练优化里最实在的几个环节。说实话，我见过太多团队把精力全砸在模型结构上，结果数据加载慢得像蜗牛，GPU利用率不到30%。这就像给法拉利配了个自行车轮胎，白瞎了算力。

我个人习惯，拿到一个训练任务，先看数据流，再看精度策略，最后才调分布式。顺序搞反了，后面全是补窟窿。好，咱们一个一个来。

4.1 数据加载流水线：别让GPU等你

数据加载，说白了就是喂饭。GPU吃得快，你喂得慢，它就饿着。我遇到过最夸张的一次，一个8卡节点，数据加载占了总训练时间的60%。后来一查，用的是最原始的torch.utils.data.DataLoader，没做任何优化。

4.1.1 DataLoader 核心参数调优

先看几个关键参数，我直接给结论：

参数	推荐值	说明
num_workers	4~8（根据CPU核心数）	不是越大越好，我试过设成32，结果CPU上下文切换把内存打爆了
prefetch_factor	2~4	预取批次数，能有效掩盖IO延迟
pin_memory	True	锁页内存，加速CPU到GPU传输
persistent_workers	True	避免每个epoch重建worker，省掉不少开销

嗯，这里要注意：num_workers不是越多越好。我曾经在一个96核的机器上设了64个worker，结果内存直接飙到200GB，训练还没开始就OOM了。建议从4开始，逐步往上加，观察GPU利用率曲线。

4.1.2 内存映射（Memory Mapping）

内存映射，说白了就是把磁盘文件直接映射到虚拟内存地址空间。这样读数据就像读内存一样，省掉了系统调用的开销。

我推荐用numpy.memmap或者torch.from_file。举个例子：

import numpy as np
import torch

# 创建内存映射文件
data = np.memmap('train_data.bin', dtype='float32', mode='r', shape=(100000, 3, 224, 224))

# 直接切片读取，不占用额外内存
batch = torch.from_numpy(data[start:end]).cuda()

为什么快？因为操作系统帮你做了缓存。第一次读慢，后面再读同一块数据，直接从page cache里拿，速度飞起。我在处理ImageNet级别的数据集时，用内存映射把加载时间从每epoch 15分钟降到了3分钟。

4.2 混合精度训练：FP16/BF16/FP8

混合精度训练，说白了就是用更少的比特数表示浮点数，换来更快的计算和更少的内存。但精度降低会带来数值问题，所以需要一些技巧。

4.2.1 三种精度的选择

精度	指数位	尾数位	适用场景
FP16	5	10	通用，但容易溢出
BF16	8	7	大模型训练，不容易溢出
FP8	4或5	3或2	推理加速，训练还在探索

我个人习惯：训练大模型（10B以上）首选BF16，因为它保留了和FP32一样的动态范围，不会出现梯度消失。中小模型用FP16，配合loss scaling就够了。FP8嘛，我建议再等等，等框架和硬件成熟了再说。

4.2.2 实战配置

用PyTorch的AMP（Automatic Mixed Precision）很简单：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

这里有个坑：GradScaler的初始scale值，默认是2^16。我遇到过loss突然变成NaN的情况，就是因为scale太大导致梯度溢出。建议设成init_scale=2**8，然后让scaler自己调整。

4.3 梯度累积与同步

梯度累积，说白了就是攒几个batch的梯度再更新一次参数。这招在显存不够用的时候特别管用。

举个例子：你想用batch size 1024，但单卡只能塞下128。那就累积8步，每步算128的梯度，攒够8步再更新。

accumulation_steps = 8
optimizer.zero_grad()

for i, (data, target) in enumerate(dataloader):
    output = model(data)
    loss = loss_fn(output, target)
    loss = loss / accumulation_steps  # 归一化
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

这里有个细节：loss要除以accumulation_steps，不然梯度会放大。我刚开始做的时候忘了这步，结果模型直接炸了。

4.3.1 梯度同步的时机

在分布式训练中，梯度同步是个大问题。DDP（Distributed Data Parallel）默认在每个step结束后同步梯度。但如果你用了梯度累积，同步频率可以降低。

我个人建议：累积步数 < 8时，每个step都同步。累积步数 > 8时，只在更新参数时同步一次。这样可以减少通信开销。

4.4 分布式训练策略：DDP/FSDP/DeepSpeed

分布式训练，说白了就是把一块GPU干不了的活，分给多块GPU一起干。但怎么分，分完怎么同步，这里面学问大了。

4.4.1 DDP：最基础的方案

DDP是PyTorch自带的，原理很简单：每个GPU上放一份完整的模型副本，各自算各自的梯度，然后all-reduce同步。

优点：实现简单，代码改动小。缺点：模型太大时，单卡放不下。

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

dist.init_process_group(backend='nccl')
model = MyModel().cuda()
model = DDP(model)

# 训练代码基本不变
for data, target in dataloader:
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

嗯，这里要注意：DDP的dataloader要用DistributedSampler，不然每个GPU会看到同样的数据。

4.4.2 FSDP：解决显存瓶颈

FSDP（Fully Sharded Data Parallel）是Meta推出的，核心思想是把模型参数、梯度、优化器状态都分片到各个GPU上。需要计算时再all-gather回来。

说白了，就是用通信换显存。我测试过，一个70B的模型，用DDP单卡需要140GB显存，用FSDP只需要40GB。

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,  # 分片梯度+优化器状态
    cpu_offload=CPUOffload(offload_params=True),  # 参数卸载到CPU
)

但FSDP有个缺点：通信开销大。如果你的集群网络带宽不够（比如只有10Gbps），FSDP可能比DDP还慢。

4.4.3 DeepSpeed：工业级方案

DeepSpeed是微软开源的，提供了ZeRO优化器。ZeRO有三个级别：

• ZeRO-1：只分片优化器状态
• ZeRO-2：分片优化器状态+梯度
• ZeRO-3：分片优化器状态+梯度+参数

我一般推荐用ZeRO-2，性价比最高。ZeRO-3虽然省显存，但通信量太大，训练速度会下降20%~30%。

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params={
        "train_batch_size": 1024,
        "gradient_accumulation_steps": 4,
        "zero_optimization": {
            "stage": 2,
            "allgather_partitions": True,
            "allgather_bucket_size": 5e8,
            "reduce_scatter": True,
            "reduce_bucket_size": 5e8,
        }
    }
)

知识体系总览

下面这张图，是我自己总结的优化决策流程。你照着走一遍，基本不会出大问题。

这张图我画了好几次才满意。你照着这个流程走，基本能覆盖90%的场景。剩下的10%，就得靠经验了。