手动实现梯度累积：从原理到实战

梯度累积，说白了就是「攒一波再更新」。

我刚开始接触这个技巧时，觉得它挺反直觉的——明明可以一次算完，为什么要拆成好几份？后来在项目里遇到显存不够的尴尬，才明白这招有多香。

今天咱们就手撕梯度累积的四种实现方式。从最朴素的for循环，到高大上的lax.scan，再到我总结的代码模板和调试技巧。嗯，一步步来。

为什么需要梯度累积？

先想一个问题：你的batch size想设成1024，但显存只够装128。怎么办？

梯度累积的答案很简单：

  • 每次只算128个样本的梯度
  • 把梯度存起来，不更新参数
  • 攒够8次（1024/128=8），再统一更新

效果等价于一次算1024。显存省了，精度没丢。

核心公式：

有效batch size = 单次batch size × 累积步数

参数更新 = 学习率 × (梯度₁ + 梯度₂ + ... + 梯度ₙ) / 累积步数

方法一：用for循环累积梯度

这是最直观的方式。我早期做分布式训练时，就是靠这个思路撑过来的。

import jax
import jax.numpy as jnp
from jax import grad, value_and_grad

def train_step_for_loop(params, x_batches, y_batches, lr=0.01):
    """
    x_batches: list of mini-batches, 比如 [batch1, batch2, ..., batchN]
    y_batches: 对应的标签
    """
    grads_acc = None
    
    for x, y in zip(x_batches, y_batches):
        loss, grads = value_and_grad(loss_fn)(params, x, y)
        
        # 累积梯度
        if grads_acc is None:
            grads_acc = grads
        else:
            grads_acc = jax.tree_map(lambda a, b: a + b, grads_acc, grads)
    
    # 平均梯度并更新参数
    num_steps = len(x_batches)
    grads_acc = jax.tree_map(lambda g: g / num_steps, grads_acc)
    params = jax.tree_map(lambda p, g: p - lr * g, params, grads_acc)
    
    return params, loss

这段代码有个坑——grads_acc is None的判断。我在项目里第一次跑的时候，忘了处理初始状态，结果报了个NoneType错误。嗯，调试了半小时。

注意：JAX的tree_map要求所有树的形状一致。如果你的模型有不同层，累积时一定要用jax.tree_map，别自己手写循环。

方法二：用lax.scan实现梯度累积

for循环虽然直观，但有个问题——它没法被JIT编译。JAX的lax.scan就是来解决这个的。

说白了，scan就是JAX版的「带状态的for循环」。它能把整个累积过程编译成一个高效的计算图。

from jax import lax

def train_step_scan(params, x_batches, y_batches, lr=0.01):
    """
    使用lax.scan实现梯度累积
    x_batches: shape (num_steps, batch_size, ...)
    """
    def scan_fn(carry, inputs):
        params, grads_acc = carry
        x, y = inputs
        
        loss, grads = value_and_grad(loss_fn)(params, x, y)
        
        # 累积梯度
        grads_acc = jax.tree_map(lambda a, b: a + b, grads_acc, grads)
        
        return (params, grads_acc), loss
    
    # 初始化累积梯度为零
    grads_init = jax.tree_map(jnp.zeros_like, params)
    init_carry = (params, grads_init)
    
    # 执行scan
    (params, grads_acc), losses = lax.scan(
        scan_fn, init_carry, (x_batches, y_batches)
    )
    
    # 平均并更新
    num_steps = x_batches.shape[0]
    grads_acc = jax.tree_map(lambda g: g / num_steps, grads_acc)
    params = jax.tree_map(lambda p, g: p - lr * g, params, grads_acc)
    
    return params, jnp.mean(losses)

我的经验：用lax.scan时，记得把grads_init初始化为全零。我曾经偷懒用None，结果scan内部没法处理，直接报错。JAX对类型一致性要求很严格。

方法三：梯度累积的代码模板

在实际项目中，我一般会封装一个通用的梯度累积类。这样切换累积步数时，改一个参数就行。

class GradientAccumulator:
    """梯度累积器模板"""
    
    def __init__(self, model, accum_steps=4):
        self.accum_steps = accum_steps
        self.model = model
        self.grads_acc = None
        self.step_count = 0
    
    def reset(self):
        """重置累积状态"""
        self.grads_acc = None
        self.step_count = 0
    
    def accumulate(self, params, batch):
        """累积一个mini-batch的梯度"""
        loss, grads = value_and_grad(self.model.loss_fn)(params, batch)
        
        if self.grads_acc is None:
            self.grads_acc = grads
        else:
            self.grads_acc = jax.tree_map(
                lambda a, b: a + b, self.grads_acc, grads
            )
        
        self.step_count += 1
        return loss
    
    def apply_gradients(self, params, lr):
        """应用累积的梯度"""
        # 平均梯度
        grads_acc = jax.tree_map(
            lambda g: g / self.accum_steps, self.grads_acc
        )
        
        # 更新参数
        params = jax.tree_map(
            lambda p, g: p - lr * g, params, grads_acc
        )
        
        # 重置状态
        self.reset()
        return params

这个模板的好处是：

  • 你可以在训练循环里随意调用accumulate
  • 攒够了自动apply_gradients
  • 支持动态调整累积步数

方法四：梯度累积的调试技巧

调试梯度累积，说实话比普通训练要麻烦一些。我踩过不少坑，分享几个实用的技巧。

技巧1：验证梯度数值

先跑一个累积步数为1的版本，和普通训练对比loss。如果不一样，说明累积逻辑有bug。

# 验证代码
params = init_params()
batch = get_batch()

# 普通训练
loss1, grads1 = value_and_grad(loss_fn)(params, batch)

# 累积训练（步数=1）
accum = GradientAccumulator(model, accum_steps=1)
loss2 = accum.accumulate(params, batch)
params2 = accum.apply_gradients(params, lr)

# 检查loss是否一致
assert jnp.allclose(loss1, loss2), "梯度累积有bug！"

技巧2：检查梯度范数

我曾经遇到一个问题：累积8步后，梯度爆炸了。排查了半天，发现是忘了除以累积步数。

def check_grad_norm(grads, threshold=1.0):
    """检查梯度范数，防止爆炸"""
    flat_grads, _ = jax.flatten_util.ravel_pytree(grads)
    norm = jnp.linalg.norm(flat_grads)
    
    if norm > threshold:
        print(f"警告：梯度范数 {norm:.4f} 超过阈值 {threshold}")
    
    return norm

技巧3：使用jax.debug.print

JIT编译后，普通的print不生效。要用jax.debug.print来打印中间值。

from jax import debug

def scan_fn(carry, inputs):
    params, grads_acc = carry
    x, y = inputs
    
    loss, grads = value_and_grad(loss_fn)(params, x, y)
    grads_acc = jax.tree_map(lambda a, b: a + b, grads_acc, grads)
    
    # 调试打印
    debug.print("当前loss: {x}", x=loss)
    debug.print("累积步数: {x}", x=carry[1])  # 注意：这里只是示例
    
    return (params, grads_acc), loss

注意：jax.debug.print只能在JIT函数内部使用。如果你在普通Python函数里用，它会报错。

知识体系图

下面这张图总结了梯度累积的核心逻辑。我画的时候特意把「调试技巧」也加进去了，因为这部分在实际项目中真的很重要。

梯度累积知识体系 梯度累积 实现方式 for循环累积 lax.scan实现 代码模板封装 调试技巧 验证梯度数值 检查梯度范数 jax.debug.print 核心：显存换batch size，注意梯度平均

总结

梯度累积这个技巧，说白了就是「空间换时间」的变种。你牺牲一点计算效率，换来更大的有效batch size。

我个人建议：

  • 小模型用for循环，简单直观
  • 大模型必须用lax.scan，否则JIT编译不了
  • 调试时先跑步数=1的版本，确认逻辑正确
  • 别忘了除以累积步数——这个坑我至少踩过三次

嗯，今天就到这里。代码模板可以直接复制到项目里用，调试技巧建议收藏一下，以后遇到梯度问题能省不少时间。

专注资料整理

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