Google JAX中的梯度检查点技术详解

沈瑗研

于 2025-05-30 09:18:16 发布

阅读量355

点赞数 4

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00678/article/details/148328191

Google JAX中的梯度检查点技术详解

jax Python+NumPy程序的可组合变换功能：进行求导、矢量化、JIT编译至GPU/TPU及其他更多操作项目地址: https://gitcode.com/gh_mirrors/ja/jax

理解梯度检查点的核心概念

在深度学习模型训练过程中，内存消耗是一个常见瓶颈。Google JAX提供的梯度检查点技术（通过jax.checkpoint或别名jax.remat实现）是一种巧妙的内存优化策略，它通过重新计算而非存储中间结果来减少内存使用。

基本原理

传统自动微分（autodiff）在正向传播时会保存所有中间结果（称为"残差"），以便在反向传播时使用。而梯度检查点技术改变了这一行为：

默认行为：正向传播保存所有中间结果，反向传播直接使用
检查点行为：正向传播仅保存指定结果，反向传播时重新计算所需中间值

这种技术实现了内存和计算量的权衡（memory-FLOPs tradeoff），在内存受限的场景下特别有用。

梯度检查点的实际应用

基础用法

import jax
import jax.numpy as jnp

# 原始函数
def model(W, x):
    y = jnp.dot(W, x)
    return jnp.sin(y)

# 应用梯度检查点
gradient_checkpointed_model = jax.checkpoint(model)

在这个简单例子中，jax.checkpoint装饰器改变了model函数的自动微分行为。

查看保存的残差

JAX提供了实用工具来检查哪些值会被保存：

from jax.ad_checkpoint import print_saved_residuals

# 查看原始函数的残差保存情况
print_saved_residuals(model, W, x)

# 查看检查点版本的残差保存情况
print_saved_residuals(gradient_checkpointed_model, W, x)

复合函数中的检查点

对于多层模型，检查点的放置位置很有讲究：

def deep_model(W1, W2, W3, x):
    x = layer(W1, x)  # 第一层
    x = layer(W2, x)  # 第二层
    x = layer(W3, x)  # 第三层
    return x

# 有效的检查点应用方式
def checkpointed_deep_model(W1, W2, W3, x):
    x = jax.checkpoint(layer)(W1, x)  # 仅对第一层应用检查点
    x = layer(W2, x)
    x = layer(W3, x)
    return x

关键原则是：不要对整个模型应用检查点，而是选择性地应用于早期层。

高级策略：自定义保存策略

JAX提供了灵活的策略系统，无需修改模型代码即可控制哪些中间值应该保存。

内置策略示例

仅保存无批次维度的矩阵乘法结果：

from jax.checkpoint_policies import dots_with_no_batch_dims_saveable

checkpointed_fn = jax.checkpoint(
    deep_model, 
    policy=dots_with_no_batch_dims_saveable
)

基于命名的保存策略：

from jax.ad_checkpoint import checkpoint_name

def named_model(params, x):
    x = checkpoint_name(layer(params[0], x), name="layer1_out")
    x = checkpoint_name(layer(params[1], x), name="layer2_out")
    return layer(params[2], x)

# 仅保存特定命名层的输出
policy = jax.checkpoint_policies.save_only_these_names(["layer1_out"])
checkpointed_named = jax.checkpoint(named_model, policy=policy)

自定义策略

虽然JAX允许自定义策略，但建议使用内置策略以确保兼容性。内置策略包括：

everything_saveable：默认策略，保存所有可保存内容
nothing_saveable：不保存任何内容，全部重新计算
dots_saveable：保存所有矩阵乘法结果
dots_with_no_batch_dims_saveable：保存无批次维度的矩阵乘法

内存卸载策略

对于超大模型，JAX还支持将中间结果卸载到CPU内存：

# 将无批次维度的矩阵乘法结果卸载到CPU
offload_policy = jax.checkpoint_policies.offload_dot_with_no_batch_dims(
    src="device", dst="pinned_host")
    
checkpointed_offload = jax.checkpoint(model, policy=offload_policy)

更精细的控制可以通过save_and_offload_only_these_names策略实现，允许指定哪些值应保留在设备上，哪些应卸载到主机内存，哪些应重新计算。