jax.remat / jax.checkpoint 更改：你需要了解的内容

jax.remat / jax.checkpoint 更改：你需要了解的内容#

内容#

发生了什么事？#

从 #11830 开始，我们正在启用 jax.checkpoint()（也称为 jax.remat()，这两个名称是彼此的别名）的新实现。**对于大多数代码，将不会有任何更改。**但在极少数情况下可能会有一些可观察到的差异；请参阅升级后可能出现哪些问题？

我现在如何禁用此更改并恢复旧行为？#

如果您遇到此更改的问题，在版本 jax==0.3.16 之前，可以通过以下任意一种方式将 jax_new_checkpoint 配置选项设置为 False 来关闭新实现。

  1. 设置 shell 环境变量 JAX_NEW_CHECKPOINT=0;

  2. 执行 jax.config.update('jax_new_checkpoint', False);

  3. 如果您使用 absl 解析标志，请传递 --jax_new_checkpoint=False 选项。

如果您需要恢复到旧的实现，请在 GitHub issue 上联系我们，以便我们能够使新的实现为您工作。

从 jax==0.3.17 开始，jax_new_checkpoint 配置选项不再可用。如果您遇到任何问题，请在 问题跟踪器 上联系我们，以便我们帮助您解决！

我们为什么要这样做？#

在撰写本文时，JAX 具有两种并行的 jax.checkpoint 实现。新的实现已经使用了几个月（例如，Pax 和 Flaxformer/T5X），并且是可选的。但它一直没有默认启用。

我们希望将新实现设置为默认启用，然后删除旧实现。使用新实现并删除旧实现可以为用户带来诸多好处。

用户可自定义的重计算策略#

新实现的主要优势是对应于 policy 参数的新功能。其理念是让用户精确控制在自动微分的前向传递过程中哪些中间结果会被保存（与重计算）。通过对内存使用与重新计算之间的权衡进行控制，用户可以获得显著的性能提升，尤其是在大型模型和我们的 LLM MLPerf 提交中！

此功能的完整文档仍在编写中，但这里有一个快速示例

from functools import partial
import jax

def apply_layer(W, x):
  return jnp.sin(jnp.dot(W, x))

@partial(jax.checkpoint, policy=jax.checkpoint_policies.checkpoint_dots)
def predict(params, x):
  for W in params[:-1]:
    x = apply_layer(W, x)
  return jnp.dot(params[-1], x)

通过在此处使用 policy=jax.checkpoint_policies.checkpoint_dots 应用 jax.checkpoint，我们确保只有矩阵乘法的结果才能在正向传递期间保存。来自 cos 应用的雅可比系数值以及计算它们所需的 sin 应用的值不会从正向传递中保存，而是在反向传递期间重新计算。（此类策略在 TPU 上可能有效，其中逐元素计算实际上是免费的，但矩阵单元的结果值得保存。）

能够重计算常量，而不仅仅是与参数存在数据依赖关系的操作#

旧的 jax.checkpoint 实现实际上无法重计算与修饰函数的参数没有数据依赖关系的计算。考虑这个玩具示例

@jax.checkpoint
def f(x):
  a = some_function(jnp.arange(10_000_000))  # `a` does not depend on `x`
  return a * x

旧的 jax.checkpoint 实现被迫保存 a 的值，这可能需要大量内存。新的 jax.checkpoint 实现可以重计算而不是保存 a 的值。

在某些情况下，Python 开销显著减少#

新的 jax.checkpoint 在某些情况下会产生更少的 Python 开销。 简单的开销基准测试 速度提高了 10 倍。这些开销仅出现在急切的操作逐个执行中，因此在使用 jax.jit 或类似操作下的 jax.checkpoint 的常见情况下，加速不相关。但无论如何，这很好！

通过简化内部结构启用新的 JAX 功能#

此更改也为未来用户解锁了巨大的好处，例如自定义批处理规则（custom_vjp 的 vmap 类比）以及对 custom_vjp 的前向可微分升级。它还显著降低了 JAX 代码库中部分内容的复杂性，这将有利于维护性和整体的错误修复。

升级后可能出现哪些问题？#

无害的数值变化#

由于新实现可以重计算更多的计算，包括可能的大常量的计算，因此某些代码可能会看到小的数值变化。任何数值变化的幅度都应该在我们预期范围内，这些变化来自更改编译器优化，例如浮点运算的重新排序。但是，一些过于严格的测试容差可能需要稍微放宽。

concrete=True 选项已移除。#

旧的 jax.checkpoint 实现有一个布尔值 concrete 选项，它允许在具体的 Python 值上进行跟踪（而不是延迟所有计算，只在抽象值上进行跟踪）。该选项很少使用，并且在使用它的那些情况下，有更简单的替代方案。因此，我们在新的 jax.checkpoint 中删除了该选项。

例如，Google 代码中绝大多数使用 concrete=True 的情况是为了支持传递类似 is_training 的参数

@partial(jax.checkpoint, concrete=True)  # OLD jax.checkpoint API
def foo(x, is_training):
  if is_training:
    return g(x)
  else:
    return h(x)

使用新的 jax.checkpoint 实现，我们可以使用 static\_argnums 选项来实现相同的功能

@partial(jax.checkpoint, static_argnums=(1,))  # NEW jax.checkpoint API
def foo(x, is_training):
  if is_training:
    ...

如果需要对静态参数执行 jax.numpy 操作，并在 Python 跟踪期间计算其数值结果而不是延迟，我们可以将 static_argnums 与 jax.ensure_compile_time_eval() 一起使用。但这似乎不太可能需要！

目录