4、函数式纯变换：jit 编译入门

为什么需要 jit？—— 从一次痛苦的训练说起

我先讲个真实经历。几年前我在做一个小型 NLP 模型，用 Python 写的训练循环。每次迭代都要逐行解释执行，一个 epoch 跑下来要 40 分钟。后来我把同样的逻辑用 jit 编译了一下，你猜怎么着？直接降到 8 分钟。5 倍的加速，代码一行没改。

说白了，Python 是解释型语言。你写个 for 循环，Python 得一条一条读、一条一条转成字节码、再一条一条执行。这中间有大量开销。而 jit 做的事情，就是把你的 Python 函数「编译」成高效的机器码，一次性执行。

我个人的习惯是：只要函数里没有动态控制流（比如依赖数据的 if/else），一律加 @jit。这几乎成了我写 JAX 代码的肌肉记忆。

核心观点： jit 不是魔法，它是把 Python 的动态执行，提前编译成静态计算图。代价是失去了部分灵活性，换来的是数量级的性能提升。

jit 的工作原理：Tracing 与 XLA

好，那 jit 到底是怎么工作的？我拆成两步来讲。

第一步：Tracing（追踪）

当你调用一个被 @jit 装饰的函数时，JAX 不会直接执行它。它会先做一次「虚拟运行」。

怎么虚拟？JAX 会构造一些特殊的「占位符」—— 叫 抽象值（Abstract Values）。这些占位符只记录形状（shape）和数据类型（dtype），不记录具体数值。然后 JAX 把函数跑一遍，记录下每一步操作。

举个例子：

import jax
import jax.numpy as jnp

@jax.jit
def f(x, y):
    z = x + y
    w = jnp.sin(z)
    return w

# 调用时，JAX 先 tracing
result = f(jnp.array([1.0, 2.0]), jnp.array([3.0, 4.0]))

在 tracing 阶段，x 和 y 被替换成抽象值。JAX 记录下：

  • 先做加法，得到一个形状为 (2,) 的 float32 张量
  • 再做 sin，同样得到 (2,) 的 float32 张量

这一步结束后，JAX 就得到了一张完整的计算图。嗯，这里要注意：tracing 只发生一次。之后你用相同形状的输入再调用，JAX 会直接复用编译好的结果。

我的经验： 如果你在 jit 函数里用了 Python 的 if/else，而且条件依赖于数据值，tracing 会出问题。因为抽象值没有具体数值，没法判断分支。我曾经在这上面踩过坑，调试了一下午才发现是动态分支导致的。

第二步：XLA 编译

计算图有了，接下来交给 XLA（Accelerated Linear Algebra）。XLA 是 Google 开源的深度学习编译器，也是 JAX 的底层引擎。

XLA 拿到计算图后，会做三件事：

  1. 图优化：合并冗余操作、消除死代码、常量折叠等。比如你写了 a + 0，XLA 会直接把它优化成 a。
  2. 算子融合：把多个小操作合并成一个内核。比如上面的 add 和 sin，XLA 会生成一个融合内核，一次性计算，减少内存读写。
  3. 代码生成：针对你的硬件（CPU、GPU、TPU）生成最优的机器码。

我举个例子说明算子融合有多重要。假设你有这么一段代码：

@jax.jit
def g(x):
    a = x * 2
    b = a + 1
    c = jnp.tanh(b)
    return c

没有融合的话，GPU 要启动 3 次内核：一次乘法、一次加法、一次 tanh。每次启动都有开销。XLA 会把它们融合成一个内核，一次启动搞定。我在项目中实测过，融合后速度能快 2-3 倍。

注意： XLA 编译本身也有开销。第一次调用 jit 函数时，需要花时间做 tracing 和编译。所以 jit 适合「编译一次，调用多次」的场景。如果你的函数只调用一次，jit 反而可能更慢。

知识体系总览

下面这张图，是我梳理的 jit 核心逻辑。你可以把它当作一个速查地图。

JIT 编译核心流程 Python 函数 + 输入 阶段 1: Tracing（追踪） 抽象值替代真实数据，记录操作序列 JAXPR 计算图（中间表示） 阶段 2: XLA 编译 图优化 → 算子融合 → 代码生成 高效机器码 关键特性 • 编译一次，多次复用 • 输入形状必须固定 • 不支持动态控制流 • 算子融合减少开销 • 硬件自适应优化 • 首次调用有编译延迟 • 适合重复调用场景

什么时候该用 jit？什么时候不该用？

我总结了一个简单的判断表，帮你快速决策：

场景 推荐做法 原因
训练循环（多次调用） 必须用 jit 编译一次，加速整个训练过程
数据预处理（一次调用） 不用 jit 编译开销可能超过执行时间
函数内有动态分支 谨慎使用 tracing 阶段无法处理数据依赖的分支
输入形状频繁变化 避免使用 每次形状变化都会触发重新编译
小规模快速实验 可不用 调试阶段保持灵活性更重要
一个小技巧： 如果你不确定某个函数能不能用 jit，先不加，跑通再说。等逻辑稳定了，再加 @jit 看效果。我经常这么干，省了很多调试时间。

避坑指南：我曾经踩过的三个坑

最后分享几个我实际遇到过的坑，希望能帮你少走弯路。

坑一： 在 jit 函数里用了 print()。tracing 阶段 print 会执行一次，但后续调用不会再执行。我当时用 print 调试，发现只打印了一次，还以为代码没跑对。后来才明白，tracing 阶段 print 的是抽象值，不是真实数据。

坑二： 把 numpy 操作和 jax.numpy 混用。在 jit 函数里，所有操作必须用 JAX 的 API。我一开始图方便，用了 np.sum()，结果 tracing 阶段直接报错。记住：jit 函数里只用 jnp，不用 np。

坑三： 忽略了编译时间。有一次我写了个推理服务，每次请求都调用 jit 函数。结果延迟反而比不用 jit 还高。后来一查，每次请求的输入形状都不一样，导致反复编译。解决方案是固定输入形状，或者用 jax.jit(static_argnums=...) 指定静态参数。

好了，关于 jit 的核心原理就讲到这里。记住三个关键词：Tracing、XLA、算子融合。理解了它们，你就掌握了 JAX 性能优化的第一把钥匙。

专注资料整理

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