Google JAX中的梯度检查点技术详解-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01195/article/details/148360649

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

jax Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more 项目地址: https://gitcode.com/gh_mirrors/jax/jax

理解JAX的梯度检查点机制

在深度学习模型训练中，内存消耗是一个常见瓶颈。Google JAX提供的jax.checkpoint（也称为jax.remat）技术能够有效控制自动微分过程中的内存使用，通过智能地在内存和计算之间进行权衡。

自动微分与内存消耗基础

在JAX的自动微分过程中，正向传播时会存储中间结果（称为"残差"）以供反向传播使用。这种机制虽然减少了重复计算，但会显著增加内存占用：

import jax
import jax.numpy as jnp

def example_func(W1, W2, W3, x):
    x = jnp.dot(W1, x)
    x = jnp.sin(x)
    x = jnp.dot(W2, x)
    x = jnp.sin(x)
    x = jnp.dot(W3, x)
    return x

使用jax.ad_checkpoint.print_saved_residuals可以查看正向传播时保存的中间值：

from jax.ad_checkpoint import print_saved_residuals
print_saved_residuals(example_func, W1, W2, W3, x)

梯度检查点的基本用法

jax.checkpoint通过减少正向传播时保存的中间值来降低内存使用：

def checkpointed_func(W1, W2, W3, x):
    x = jax.checkpoint(lambda W, x: jnp.sin(jnp.dot(W, x)))(W1, x)
    x = jax.checkpoint(lambda W, x: jnp.sin(jnp.dot(W, x)))(W2, x)
    x = jnp.dot(W3, x)
    return x

应用检查点后，正向传播只保存必要的输入，反向传播时再重新计算需要的中间值。

高级策略与自定义控制

策略函数的使用

JAX提供了多种预定义的策略函数来控制哪些中间值可以被保存：

# 只保存无批量维度的矩阵乘法结果
policy = jax.checkpoint_policies.dots_with_no_batch_dims_saveable
checkpointed_with_policy = jax.checkpoint(example_func, policy=policy)

命名中间值进行精细控制

通过checkpoint_name可以标记特定中间值，然后使用策略函数精确控制：

from jax.ad_checkpoint import checkpoint_name

def named_func(W1, W2, W3, x):
    x = checkpoint_name(jnp.dot(W1, x), name='layer1_dot')
    x = jnp.sin(x)
    x = checkpoint_name(jnp.dot(W2, x), name='layer2_dot')
    x = jnp.sin(x)
    x = jnp.dot(W3, x)
    return x

# 只保存特定命名的中间值
policy = jax.checkpoint_policies.save_only_these_names('layer1_dot')
named_checkpointed = jax.checkpoint(named_func, policy=policy)

实际应用中的考量

与JIT编译的交互

当jax.checkpoint与jax.jit一起使用时需要注意：

JIT会优化计算图，可能影响检查点的预期行为
某些策略在编译后可能有不同的内存表现
建议先测试不同策略的实际内存节省效果

递归检查点技术

对于深度网络，递归应用检查点可以实现内存使用的对数级增长：

def recursive_checkpoint(funs):
    if len(funs) == 1:
        return funs[0]
    elif len(funs) == 2:
        f1, f2 = funs
        return lambda x: f1(f2(x))
    else:
        f1 = recursive_checkpoint(funs[:len(funs)//2])
        f2 = recursive_checkpoint(funs[len(funs)//2:])
        return lambda x: f1(jax.checkpoint(f2)(x))

这种技术虽然节省内存，但会增加计算量，需要在具体场景中权衡。