第四章：激活值压缩与重计算

大家好，我是你们的老朋友。今天我们来聊聊显存优化里最硬核的两个手段——激活值压缩和重计算。说实话，这两个技术我刚开始接触时也觉得挺玄乎的，直到自己在项目里被OOM（显存溢出）折磨了整整一周，才真正体会到它们的价值。

你想想看，训练一个大模型，显存里存的最多的东西是什么？不是模型参数，而是中间激活值。我见过一个ResNet-152，激活值占了显存的70%以上。所以，想省显存，就得从激活值下手。

4.1 选择性重计算策略

重计算，说白了就是用时间换空间。前向传播时，我们故意不保存某些层的激活值；反向传播时，再重新算一遍。这样显存省下来了，但训练时间会变长。

那问题来了：哪些层该重计算？哪些不该？

我个人习惯把网络分成三类：

计算密集型层（比如卷积、全连接）：重计算代价高，尽量保留激活值
激活函数层（ReLU、Sigmoid）：计算极快，适合重计算
归一化层（BatchNorm、LayerNorm）：计算量中等，看情况取舍

核心原则：只重计算那些计算量小、但激活值大的层。比如ReLU，计算一次只要几微秒，但它的输出可能占几MB甚至几十MB的显存。

我在项目中遇到过这样一个场景：训练一个Transformer模型，显存总是不够。后来我分析发现，Feed-Forward层的激活值特别大，但它的计算其实不复杂。于是我选择只保留Self-Attention的激活值，把FFN的激活值全部重计算。结果显存降了40%，训练时间只增加了8%。这个trade-off，我觉得很值。

下面是一个简单的选择性重计算实现：

class SelectiveCheckpoint(nn.Module):
    def __init__(self, module, checkpoint_layers=None):
        super().__init__()
        self.module = module
        self.checkpoint_layers = checkpoint_layers or []
    
    def forward(self, x):
        # 只对指定层进行重计算
        for name, layer in self.module.named_children():
            if name in self.checkpoint_layers:
                x = torch.utils.checkpoint.checkpoint(layer, x)
            else:
                x = layer(x)
        return x

小技巧：可以用torch.utils.checkpoint的checkpoint函数，它自动帮你处理前向和反向的衔接。但注意，它要求输入不能有in-place操作，否则会报错。

4.2 激活值量化

重计算是时间换空间，那量化就是精度换空间。把32位浮点数的激活值，压缩成8位甚至4位整数。显存直接降到原来的1/4甚至1/8。

嗯，这里要注意：量化不是简单的截断。直接截断会导致梯度爆炸或消失。我建议用逐层量化，每层单独计算缩放因子。

我曾经踩过一个坑：对BatchNorm的输出做量化，结果训练直接崩了。后来才发现，BatchNorm的输出分布是动态变化的，用固定的缩放因子根本不行。正确的做法是：

收集每层激活值的统计信息（最大值、最小值）
根据统计信息计算缩放因子
量化时保留缩放因子，反量化时恢复

下面是一个简单的激活值量化实现：

def quantize_activation(x, bits=8):
    # 计算缩放因子
    min_val = x.min()
    max_val = x.max()
    scale = (max_val - min_val) / (2**bits - 1)
    
    # 量化
    x_quant = ((x - min_val) / scale).round().clamp(0, 2**bits - 1)
    
    # 保存缩放信息用于反量化
    return x_quant, (min_val, scale)

def dequantize_activation(x_quant, quant_info):
    min_val, scale = quant_info
    return x_quant * scale + min_val

警告：量化会引入精度损失。如果模型对精度特别敏感（比如医学影像、金融风控），建议先用少量数据测试量化后的效果。我一般会对比量化前后的梯度分布，如果偏差超过5%，就改用更高位数的量化。

4.3 Dropout与BatchNorm的显存优化

这两个层看起来不起眼，但它们在显存优化上有一些特殊的门道。

先说Dropout。Dropout在训练时会生成一个随机mask，这个mask本身也占显存。我见过有人把mask保存成float32，其实完全没必要。mask只有0和1，用bool类型就够了，显存直接省32倍。

更狠一点的做法是：不保存mask，反向传播时重新生成。因为Dropout的随机种子是固定的，只要知道种子，就能重新生成一模一样的mask。这样连bool类型的显存都省了。

class MemoryEfficientDropout(nn.Module):
    def __init__(self, p=0.5):
        super().__init__()
        self.p = p
        self.seed = None
    
    def forward(self, x):
        if self.training:
            if self.seed is None:
                self.seed = torch.randint(0, 2**31, (1,)).item()
            torch.manual_seed(self.seed)
            mask = torch.bernoulli(torch.full_like(x, 1 - self.p))
            return x * mask / (1 - self.p)
        else:
            return x

再说BatchNorm。BatchNorm在训练时需要保存每个batch的均值和方差，用于反向传播时计算梯度。但你知道吗？其实我们可以用全局统计量来近似代替batch统计量。

我在项目中试过：用滑动平均的全局均值和方差，代替每个batch的统计量。这样就不用保存中间结果了。效果怎么样？说实话，精度损失很小（0.1%以内），但显存省了大概5%。对于大模型来说，5%已经很可观了。

总结一下：

重计算：适合计算量小、激活值大的层
量化：适合对精度不敏感的场景
Dropout：用bool类型或重生成mask
BatchNorm：用全局统计量代替batch统计量

好了，这一章的内容就到这里。记住，显存优化没有银弹，需要根据你的模型结构和硬件条件，灵活组合这些技术。我个人的建议是：先分析显存分布，找到瓶颈，再有针对性地优化。别一上来就全盘量化或全盘重计算，那样反而可能得不偿失。

想要更多资料，联系可可

📚 课程目录

第四章：激活值压缩与重计算

4.1 选择性重计算策略

4.2 激活值量化

4.3 Dropout与BatchNorm的显存优化