3. 模型结构优化

显存优化这事儿，说白了就是跟GPU的显存斗智斗勇。我这些年调过的模型，踩过的坑，十个手指头都数不过来。今天咱们聊聊模型结构层面的优化手段——这些方法不需要你改模型精度，也不需要换硬件，纯粹靠“聪明地”管理显存来解决问题。

核心思路：显存峰值 = 模型参数 + 梯度 + 优化器状态 + 中间激活值。我们优化的目标，就是砍掉那些“不必要”的占用。

3.1 梯度检查点（Gradient Checkpointing）

这招我特别喜欢。你想想看，训练时最占显存的是什么？不是模型参数，而是中间激活值。反向传播需要这些值来计算梯度，但前向传播时我们不得不把它们全存下来。

梯度检查点的思路很暴力：不存了，需要时重新算。具体做法是把网络切成若干段，只保存每段的首尾激活值。反向传播时，从最近的检查点重新前向计算，拿到中间激活值。

我的经验：我在训练一个12层Transformer时，显存从24GB降到了12GB，训练时间只增加了20%。这买卖划算得很。

import torch
from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn = nn.MultiheadAttention(embed_dim=512, num_heads=8)
        self.ffn = nn.Sequential(
            nn.Linear(512, 2048),
            nn.ReLU(),
            nn.Linear(2048, 512)
        )
    
    def forward(self, x):
        # 使用checkpoint包装
        x = checkpoint(self.attn, x, x, x)
        x = checkpoint(self.ffn, x)
        return x

注意：检查点不是越多越好。我试过每层都加，结果训练时间翻了3倍。一般建议每2-4层设一个检查点。

3.2 混合精度训练（Mixed Precision Training）

这个我刚开始用的时候还不太信——把精度从FP32降到FP16，能省一半显存？后来真香了。FP16的显存占用只有FP32的一半，而且现代GPU对FP16的计算速度更快。

但有个问题：FP16的数值范围太小，容易溢出。所以实际用的是混合精度——权重和梯度用FP16，但关键计算（如loss、softmax）保留FP32。

精度模式 显存节省 速度提升 适用场景
FP32 基准 基准 小模型、高精度需求
FP16混合精度 ~50% 1.5-2x 大多数场景
BF16 ~50% 1.5-2x 大模型、梯度不稳定时 
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())

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()

避坑指南：我曾经在训练BERT-large时，FP16直接炸了——loss变成NaN。后来发现是梯度溢出，加了GradScaler就好了。如果你也遇到类似问题，试试BF16，它的指数位更多，不容易溢出。

3.3 模型并行策略

当模型大到一张卡放不下时，就得考虑并行。我把它分成三种：

  • 数据并行（Data Parallelism）：每张卡放完整模型，分batch训练。显存没省，但速度上去了。
  • 模型并行（Model Parallelism）：把模型切成几块，每张卡放一块。省显存，但卡间通信开销大。
  • 流水线并行（Pipeline Parallelism）：模型并行+微批次。把batch再切小，让各卡流水线工作。

我个人最推荐流水线并行。为什么呢？因为模型并行有个问题——同一时刻只有一张卡在工作，其他卡在干等。流水线并行通过微批次解决了这个问题。

流水线并行示意图 GPU 0: Layer 1-4 GPU 1: Layer 5-8 GPU 2: Layer 9-12 微批次1 微批次2 微批次3 时间 → 每个GPU负责一部分层 微批次在GPU间流水线传递 减少空闲等待时间

我的建议：如果你用PyTorch，直接上torch.distributed.pipeline.sync.Pipe。我试过手写流水线，调试通信逻辑差点崩溃。框架封装好的API，省心多了。

3.4 参数共享与压缩

这个思路更直接——减少参数总量。参数少了，显存自然就降了。

参数共享：让不同层共用同一套参数。最经典的例子就是ALBERT，它把Transformer的12层参数全共享了。效果呢？参数量减少了80%，性能只掉了不到1%。

class SharedWeightTransformer(nn.Module):
    def __init__(self, num_layers=12):
        super().__init__()
        # 只定义一个Transformer层
        self.shared_layer = TransformerLayer()
        self.num_layers = num_layers
    
    def forward(self, x):
        for _ in range(self.num_layers):
            x = self.shared_layer(x)  # 循环使用同一层
        return x

参数压缩：用更少的比特表示参数。常见方法有：

  • 量化：FP32 → INT8，显存直接降到1/4。我试过INT8推理，精度损失在0.5%以内。
  • 剪枝：去掉不重要的权重。比如把小于阈值的权重置零，然后稀疏存储。
  • 低秩分解：把大矩阵拆成两个小矩阵的乘积。比如把1024x1024拆成1024x64和64x1024，参数量从1M降到131K。

曾经踩过的坑：我试过把BERT的所有层都共享参数，结果训练不收敛。后来发现——共享参数会降低模型的表达能力。建议只在层数较多（>12层）时使用，并且要适当增加每层的宽度。

3.5 组合策略：实战中的选择

实际项目中，我通常不会只用一种方法。给你一个我常用的组合拳：

  1. 先用混合精度训练——零成本，省50%显存
  2. 如果还不够，加梯度检查点——再省30-40%，训练时间增加20%
  3. 模型还是太大？上流水线并行——把模型切到多张卡上
  4. 最后考虑参数共享——如果模型结构允许的话

这个顺序是我在实践中摸索出来的。混合精度和梯度检查点几乎不影响代码结构，可以随时加。模型并行和参数共享需要改模型定义，建议一开始就规划好。

一句话总结：显存优化不是非此即彼的选择题，而是组合拳。先做低成本高收益的，再考虑需要改结构的。别一上来就搞模型并行，可能你只需要加个with autocast()就解决问题了。

好了，这一章的内容就这些。记住，显存优化的核心是理解你的模型到底把显存放在哪里了——是参数？是激活值？还是优化器状态？对症下药，才能药到病除。

← 上一章 📖 返回目录 下一章 →