JAX 101 第六篇《从零拼出迷你 Transformer》使用 Python for 循环逐层调用 block。模型加深之后，数学上仍然是“把上一层结果交给下一层”，工程上却需要关心编译体积、层参数的存储方式，以及反向传播保留多少中间结果。本篇就沿着这三个问题介绍 lax.scan。

scan 是按固定顺序重复执行同一段计算，每一步拿到自己的输入，并把一份状态交给下一步。 扫描的可以是时间步，也可以是模型层；本文重点是层堆叠。文中的代码按顺序组织，使用同一组导入和函数定义。

1. Python 循环的问题

迷你 Transformer 中，层堆叠可以直接写成 Python 循环：

def model(params, x):
    for layer_params in params["blocks"]:
        x = block(layer_params, x)
    return x

层数是静态 Python 列表长度时，JAX tracing 会展开这个循环。少量层很直观，但层数很大时，展开会增加追踪和编译程序的体积。

例如，64 层就可能在被追踪的程序中出现 64 份 block 运算。虽然参数和计算本来就有 64 层，编译器还要处理展开后的大量运算节点，首轮编译可能变慢。scan 让程序用一个循环体加上迭代次数描述这件事，便于控制编译表示的规模。

jax.lax.scan 用一个循环 body 表达重复计算：

final_carry, stacked_outputs = jax.lax.scan(body, init_carry, xs)

• carry 是每一步传给下一步的循环状态；
• xs 是沿第 0 维逐步读取的输入树；
• body 返回新的 carry 和当前步输出；
• carry 的 shape 和 dtype 在每一步必须一致。

可以把接口理解成 body(carry, one_input) -> (new_carry, one_output)。carry 是下一步还要继续使用的状态，one_output 是本步希望收集的结果，两者用途不同。在 Transformer 层堆叠中，隐藏状态是 carry，每一层的参数是本步输入；如果最后只要最终隐藏状态，就不必额外收集每一层输出。接口细节见 JAX scan 文档。

scan 中存在前后依赖：下一层需要上一层的输出，所以层与层仍按顺序计算。vmap 则适合把彼此独立的计算沿一个轴批量映射。例如，可以用 vmap 独立初始化各层参数，再用 scan 顺序执行这些层。

2. 从累加序列看状态传递

import jax
import jax.numpy as jnp

def body(total, x):
    new_total = total + x
    return new_total, new_total

final_total, prefix_sums = jax.lax.scan(
    body,
    jnp.array(0, dtype=jnp.int32),
    jnp.array([1, 2, 3, 4], dtype=jnp.int32),
)

print(final_total)  # 10
print(prefix_sums)  # [1, 3, 6, 10]

这与 Python 循环的逻辑相同，但 scan 是 JAX primitive，能保持紧凑的编译表示，并能与 jit 和自动微分组合。

第一步用状态 0 加上输入 1，得到新状态 1，并记录输出 1；第二步再加 2，得到状态和输出 3。四步结束后，最终状态是 10，逐步收集的输出是 [1,3,6,10]。这个例子中状态和输出恰好取同一个数，但在其他任务里可以不同。

固定步数的序列计算适合先考虑 scan。如果停止条件依赖运行时结果，可以研究 lax.while_loop；它与 scan 在自动微分支持上也有差异，不能只按语法直接替换。Python 短循环仍然很适合小模型，是否改写应结合编译时间和维护成本决定。

3. 参数为什么要“堆叠”

Python 列表参数的树结构是“第 0 层、第 1 层……各有一套叶子”。scan 的 xs 需要从叶子的第 0 维逐层切片，因此我们把相同名称参数堆成：

Python layers:                 scanned layers:
layers[0].w [D,D]              w [L,D,D]
layers[1].w [D,D]      →
...
layers[L-1].w [D,D]

def init_scanned_layers(key, num_layers, dim):
    keys = jax.random.split(key, num_layers)

    def init_one(k):
        k1, k2 = jax.random.split(k)
        return {
            "w1": 0.02 * jax.random.normal(k1, (dim, 4 * dim)),
            "w2": 0.02 * jax.random.normal(k2, (4 * dim, dim)),
            "scale": jnp.ones((dim,)),
        }

    return jax.vmap(init_one)(keys)

vmap(init_one)(keys) 返回 PyTree，其中每个数组叶子都多了一个最前面的层维 L。

于是 w1 是 [L,D,4D]，w2 是 [L,4D,D]，scale 是 [L,D]。scan 第 i 步分别取这些数组的第 i 片，组成这一层的参数字典。堆叠只改变参数的组织方式，各层仍然有独立权重，并没有变成所有层共享同一套参数。

这也解释了一个常见限制：被同一个 scan 扫描的参数叶子，需要有一致的层内 shape 和 dtype。如果不同层使用不同宽度或不同结构，可以先把同构层分组扫描，或保留部分 Python 控制逻辑，不能直接把不同形状的权重 stack 在一起。

4. 扫描 Transformer 风格的残差 MLP

def rms_norm(scale, x, eps=1e-6):
    rms = jnp.sqrt(jnp.mean(x.astype(jnp.float32) ** 2, axis=-1,
                            keepdims=True) + eps)
    return (x / rms) * scale

def layer_body(x, layer):
    h = rms_norm(layer["scale"], x)
    h = jax.nn.gelu(h @ layer["w1"])
    h = h @ layer["w2"]
    return x + h, None

@jax.jit
def apply_layers(layers, x):
    x, _ = jax.lax.scan(layer_body, x, layers)
    return x

L, B, T, D = 4, 2, 8, 32
layers = init_scanned_layers(jax.random.key(0), L, D)
x = jnp.ones((B, T, D))
y = apply_layers(layers, x)
assert y.shape == (B, T, D)

这里 x 是 carry，层参数是 xs。因为我们不需要保存每层输出，第二个返回值写 None。

每次 layer_body 收到的 x 都是 [B,T,D]，layer 中已经没有最前面的层维。MLP 在内部扩宽，再回到 D 维，与输入做残差相加，所以输出 carry 的 shape 保持不变。这个示例只有归一化和残差 MLP，没有注意力；目的是单独验证层扫描的数据组织。

代码先采用 FP32。以后切换到 BF16 时，需要明确 RMSNorm 归约、权重和残差输出各自的 dtype；如果首步返回的 carry 从 BF16 被提升成 FP32，就会破坏 scan 的类型约定。加入 dropout 时，也应给每层准备独立 key，例如把层参数和层 key 一起放入 xs，不能在循环体中反复使用同一个 key。

5. scan 与梯度

def loss(layers, x):
    return jnp.mean(apply_layers(layers, x) ** 2)

grads = jax.grad(loss)(layers, x)
assert jax.tree.structure(grads) == jax.tree.structure(layers)

JAX 会为整个扫描计算梯度。长序列或多层模型的反向传播可能保存大量中间激活，因此还要结合 rematerialization。

梯度树同样带有层维，例如 grads["w1"] 是 [L,D,4D]，每片对应一层的梯度。只验证 shape 还不够，改写模型时应该用同一份参数与输入，比较 Python 循环和 scan 的前向结果、loss、梯度，再检查一次参数更新。这样可以尽早发现层顺序或参数切片错误。

None 只表示不返回逐层堆叠的输出。它不会自动取消反向传播需要保存的中间量。 返回哪些结果与反向保存哪些激活，是两项不同的内存开销。

6. 与 jax.checkpoint/remat 配合

remat_layer_body = jax.checkpoint(layer_body)

@jax.jit
def apply_layers_with_remat(layers, x):
    x, _ = jax.lax.scan(remat_layer_body, x, layers)
    return x

remat 会在反向传播时重新计算部分前向结果，以计算换内存。它不是免费的：减少多少内存、增加多少计算取决于 checkpoint 边界和 block 内容，应通过 profile 与显存数据判断。

这里的 rematerialization 可以理解为“需要用到某些中间结果时，再算一次”。jax.checkpoint 包住层计算后，会影响自动微分保留哪些中间结果、哪些交给反向阶段重算。它与把模型权重保存到磁盘的训练 checkpoint 是两个用途：前者控制激活内存，后者用于恢复训练。具体策略见 JAX gradient checkpointing 文档。

工程上先确认开启前后的 loss 和梯度在预期误差内一致，再比较峰值设备内存与稳定训练步耗时。若模型本来就能放下，更多重算可能拖慢训练；若节省的内存允许使用更大的 batch 或序列长度，就要结合新配置比较整体吞吐。scan 本身也不保证每步执行更快，较紧凑的编译表示与运行时性能需要分别测量。

7. 常见错误

carry 的形状发生变化。

scan 要求每步输入和输出 carry 具有相同结构、shape 和 dtype。不能在循环里不断 concatenate 让序列增长。自回归生成通常预分配 buffer，再用动态索引更新。

层参数没有堆叠。

如果参数仍是长度为 L 的 Python 列表，scan 不会自动把它理解成数组的扫描维。初始化时直接堆叠，或用 tree.map(lambda *xs: jnp.stack(xs), *layers) 转换。

把所有输出都保存下来。

不需要逐层输出时返回 None，否则会产生 [L,B,T,D] 等额外结果。

把异步派发时间当成执行时间。

计时前先做预热，等输出 .block_until_ready() 后再结束计时。首轮包含 tracing 和编译，应单独记录；后续比较使用相同参数、shape、dtype 和设备，分别观察编译耗时与稳定执行耗时。改写前后如果同时改变了模型结构，就很难判断差异来自 scan 还是其他改动。