Decoder 堆叠与重计算 / Offload

场景定位：Decoder 层用 nn.scan 堆叠成 While 循环，配合 remat policy 在显存与重计算之间权衡，必要时把激活 offload 到 host。

入口文件：layers/decoders.py、configs/types.py（remat_policy）

本场景原语：nn.scan / nn.remat / remat_policy / jax.checkpoint_policies（save_only_these_names、save_and_offload_only_these_names）

逐原语分析

`nn.scan`：把同构 Decoder 层提升为单个 loop primitive

1. MaxText 是如何使用 `nn.scan` 的

MaxText 的 Linen Decoder 在 scan_layers=True 时，不再用 Python for 逐层调用 decoder layer，而是在 scan_decoder_layers 中把某个 layer class 交给 nn.scan。被扫描的对象是“同构层模板”：每次迭代复用同一份 layer 计算图，但从带 layer 轴的变量树里取出当前层的参数、cache 和中间状态。

这一层包装的关键配置是：

variable_axes：把 "params" 放到 cfg.param_scan_axis 上；推理或训练 apply 阶段会使用 ScanIn(cfg.param_scan_axis)，表示参数变量输入沿该轴被 scan。"cache"、"intermediates"、"aqt"、"_overwrite_with_gradient" 等集合沿轴 0 扫描或收集。
split_rngs："params": True 表示初始化每层参数时拆分 RNG，使不同层拿到不同初始化随机数；"dropout": cfg.enable_dropout 表示启用 dropout 时每层也使用拆分后的 RNG。
in_axes：hidden state y 是 carry；attention mask、positions、segment ids 等每层共享输入通常用 nn.broadcast；需要逐层变化的 layer_idx 则放在 scan 轴上。
length：对应被这个 scan 管理的层数，不一定等于整个模型层数。Gemma3、Ling3 等异构结构会把 dense prefix、KDA/MoE prefix、MLA prefix、scannable body 拆成多段 scan。
metadata_params={nn.PARTITION_NAME: metadata_axis_name}：给 scan 轴附上类似 "layers"、"dense_layers"、"moe_layers_kda_cycle_0" 的逻辑分区名，后续逻辑轴规则可以识别 layer 维度。

所以 MaxText 使用 nn.scan 的主要目的不是改变模型数学语义，而是把“多层同构 decoder block”表达为一个 rolled loop。收益是 layer 参数多出一个显式 layer 轴、JIT 编译图不随层数线性膨胀、scan 轴可以参与 FSDP/sharding；约束是被 scan 的 body 必须同构，carry 的 pytree 结构、shape、dtype 也必须稳定。

2. Flax 是如何翻译 Linen `nn.scan` 的

nn.scan 是 Flax Linen 对 jax.lax.scan 的 lifted transform：JAX 原生 transform 只处理纯函数，而 Flax 需要把 Module 里的变量集合、RNG 序列和方法一起“函数化”后再交给 JAX。Linen 入口是 flax.linen.transforms.scan，它通过 lift_transform 调到 flax.core.lift.scan；lift.scan 再调用 flax.core.axes_scan.scan，最后落到 jax.lax.scan。

Flax 在进入 JAX 之前会把值分成三类：

scan：沿层轴逐步切片的值，例如每层不同的参数。
carry：每步更新并传给下一层的值，例如 hidden states、cache state。
broadcast：闭包常量或每层共享的值，例如部分输入、配置、mesh。

具体翻译过程可以理解为：

variable_axes 被拆成输入轴和输出轴，Flax 对被 scan 的变量集合先 remove_axis，让 body 看到的是“单层变量”。
split_rngs 先根据 length 对指定 RNG collection 做 random.split，然后把拆好的 RNG 作为 scan 输入；未拆分的 RNG 作为 broadcast 值复用。
axes_scan.scan 把任意 in_axes != 0 的输入先转置到 leading axis，因为 jax.lax.scan 只沿 leading axis 扫描；结束后再按 out_axes 转回原轴。
body 每次用当前切片变量和 RNG 重新构造临时 Scope，执行原始 Module 或方法，随后把新的变量树重新打包为 carry、scan output 或 broadcast output。
如果开启 check_constancy_invariants，Flax 会额外 trace 一次 broadcast body，确认 broadcast 输出不依赖 scan body 的逐步输入；这避免把真正逐层变化的值误标成 broadcast。
执行完 jax.lax.scan 后，Flax 对 scan 变量集合重新 add_axis，并把 metadata_params 写回 AxisMetadata，使变量树恢复为带 layer 轴的 Flax 变量。

一个具体例子更容易看清这层翻译。假设一个同构的 4 层 Decoder stack，scan_layers=True、param_scan_axis=1，没有 pipeline，也不是 DeepSeek/Ling3 这类异构拆分。MaxText 调用形态可以简化成：

python

scan_fn = nn.scan(
    RemattedDecoderLayer,
    variable_axes={
        "params": ScanIn(1),  # 初始化阶段是 1，apply 阶段是 ScanIn(1)。
        "cache": 0,
        "intermediates": 0,
        "aqt": 0,
    },
    split_rngs={"params": True, "dropout": enable_dropout},
    in_axes=(nn.broadcast, nn.broadcast, nn.broadcast),
    length=4,
    metadata_params={nn.PARTITION_NAME: "layers"},
)

y, _ = scan_fn(config=cfg, mesh=mesh, name="layers")(y, mask, positions, segment_ids)

这里 y 是 hidden states，作为 scan carry 从第 0 层传到第 3 层；mask、positions、segment_ids 是每一层都看到的 broadcast 输入。初始化阶段因为 "params" 的轴规格是 1 且 "params" RNG 被 split，Flax 会用 4 个不同 RNG 初始化 4 份 layer 参数，然后按 axis 1 stack 回一个带 layer 轴的变量树。例如单层参数和 scan 后参数大致变成：

单层参数	单层 shape	scan 后 shape	说明
attention q/kernel	`[hidden, q_dim]`	`[hidden, 4, q_dim]`	第 1 维是 `layers` scan 轴
attention o/kernel	`[o_in, hidden]`	`[o_in, 4, hidden]`	每层权重不同，但共享同一 body jaxpr
mlp wi/kernel	`[hidden, mlp]`	`[hidden, 4, mlp]`	`param_scan_axis=0` 时则会是 `[4, hidden, mlp]`
norm scale	`[hidden]`	`[hidden, 4]`	rank-1 参数会在 axis 1 stack

在这个结构中，权重的身份可以分三层看：

在 Flax 层，权重不是 Module 对象里的 Python 字段，而是 "params" collection 下的变量树；每个叶子是一个 JAX array。RemattedDecoderLayer / DecoderLayer 只是同构计算模板，真正的权重值由外层变量树持有。
在初始化阶段，MaxText 把 "params" 设成 cfg.param_scan_axis，同时 split_rngs={"params": True}。Flax 会按 scan length 拆出每层参数 RNG，让同一个 layer body 初始化多份参数，再把这些叶子按 param_scan_axis stack 回 "params" collection。因此 scan 后的权重是“带 layer 轴的一棵参数 pytree”，不是 N 个独立 Python layer 对象。
在 apply 阶段，MaxText 把 "params" 设成 ScanIn(cfg.param_scan_axis)。这表示参数不是 carry，也不是 broadcast const：它们作为 scan input 被逐层切片。carry 是 hidden state、cache 等会跨层更新的状态；broadcast 是 mask、positions、segment ids 这类每层相同的输入；参数则是每层不同但不在 body 内被更新的 scanned variable。

执行 apply 时，Flax 传给 JAX 的不再是一个 Module，而是一个纯函数近似形态：

python

def scanned_body(loop_carry, scan_vars_and_rngs):
  cache_vars, y = loop_carry
  params_i, cache_i, aqt_i, rng_i = scan_vars_and_rngs
  scope_i = make_scope(
      params=params_i,
      cache=cache_i,
      aqt=aqt_i,
      rngs=rng_i,
      broadcast_inputs=(mask, positions, segment_ids),
  )
  y_next, aux = RemattedDecoderLayer(scope_i, y, mask, positions, segment_ids)
  cache_i_next = scope_i.updated("cache")
  return (cache_i_next, y_next), aux

这个伪代码里最关键的是 params_i：虽然外层变量树保存的是 [hidden, 4, q_dim]，但 Flax 在每次 scan body 里通过 jax.lax.scan 的 leading-axis 切片机制给 body 的是单层视角。因为 param_scan_axis=1 不是 leading axis，axes_scan.scan 会先把它转到最前面，等价于把 [hidden, 4, q_dim] 临时变成 [4, hidden, q_dim]；jax.lax.scan 第 i 步切出 [hidden, q_dim]，body 执行完后再把输出变量轴转回 axis 1。

这一步对应的 Flax 代码逻辑可以概括为：

python

# lift.scan 的结构化伪代码
params_axis = split_variable_axes(variable_axes)["params"]  # ScanIn(axis)
scan_in_axes = (variable_axes_for_scanned_vars, rng_axes, user_arg_axes)
scanned = axes_scan(scan_body, in_axes=scan_in_axes, out_axes=...)

也就是说，MaxText 传入的 "params": ScanIn(cfg.param_scan_axis) 最终会进入 variable_in_axes，再作为 axes_scan.scan 的第一组 in_axes。在调用 scanned(...) 前，lift.scan 还会对 scan variable 调 meta.remove_axis(scan_group, axis, metadata_params)；这一步只更新 AxisMetadata，例如从分区名里暂时删除 "layers" 轴，让 body 里的单层参数 metadata 和单层 shape 对齐，不负责真正移动数组维度。body 执行结束后再用 meta.add_axis(...) 把 metadata 加回去。

真正移动数组维度的是 flax.core.axes_scan.scan，它的逻辑可以简化成：

python

def move_scan_axis_to_front(x, axis):
  if axis == 0:
    return x
  return transpose(x, permutation=[axis, *all_other_axes_in_original_order])

xs_for_jax = tree_map(move_scan_axis_to_front, in_axes, original_args)
carry_out, ys_leading = lax.scan(body, carry_in, xs_for_jax)
ys = tree_map(move_leading_axis_back, out_axes, ys_leading)

如果某个参数叶子 shape 是 [hidden, 4, q_dim]，且 param_scan_axis=1，这里的 perm 就是 (1, 0, 2)，转置后变成 [4, hidden, q_dim]。随后 jax.lax.scan 按自己的规则沿 xs 的 leading axis 迭代；第 i 步传给 body_fn 的 xs 已经是 leading-axis 切片，所以 params_i["attention"]["q"]["kernel"] 的 shape 是 [hidden, q_dim]。如果 param_scan_axis=0，transpose_to_front 直接返回原数组，不需要额外 transpose。

body_fn 里还有一层 broadcast 注入：

python

for each input leaf:
  if its axis marker is broadcast:
    body sees the original full value
  else:
    body sees the current scan slice

这保证 mask、positions 这类 broadcast 输入不被按层切片，而参数这种非 broadcast 输入使用 lax.scan 给出的当前切片。于是每一层看到的是单层参数 pytree，外层仍然只保留一棵带 layer 轴的参数 pytree。

降到 jax.lax.scan 后，这棵参数 pytree 表现为 xs 的一部分。语义上它不是 loop-carried state，因为第 i 步不会把 params_i 更新成 params_{i+1}；第 i+1 步只是从同一棵带 layer 轴的参数树里取下一片。再往下到 StableHLO / XLA 时，完整参数数组会作为 while body 可见的 loop-invariant 输入或 while 状态里的不变分量存在，body 根据当前 loop index 对参数数组做切片，得到当前层权重，再执行同一份 layer computation。

从分区角度看，metadata_params={nn.PARTITION_NAME: "layers"} 会把这个新增的 layer 轴标成 "layers"。如果逻辑轴规则把 "layers" 映射到物理 mesh 的 "stage"，那么 scan 后参数的那一维可以被 sharding 规则识别为 layer 维度；业务代码不需要在每个 kernel 参数上手写 mesh 细节。

换句话说，Linen nn.scan 的主要工作是处理 Module 状态、变量集合、RNG、轴元数据和 in_axes/out_axes API；真正的 loop primitive 仍然是 jax.lax.scan。

3. Flax 新的 `nnx.scan` 的差异

NNX 的对象模型和 Linen 不同。Linen 用 Scope 和变量 collection 表示 Module 状态，因此 nn.scan 需要 variable_axes、variable_carry、variable_broadcast 这些 collection filter。NNX 则把 Module 看成可拆分的对象图：nnx.split 可以把对象拆成 GraphDef、Param state 和其他 state，nnx.merge 再把它们合回可调用对象。因此 NNX 的 nnx.scan API 更接近普通函数 transform：

用 in_axes=(nnx.Carry, 0)、out_axes=nnx.Carry 标记哪些参数是 carry、哪些对象或数组沿 axis 0 扫描，而不是用 Linen 的 variable_axes collection mapping。
StateAxes 可以描述 NNX state 内部不同变量类型的轴，例如 nnx.Param 沿 param_scan_axis，其他 state 沿 0。
transform_metadata={nnx.PARTITION_NAME: "layers"} 用于给 NNX state 的 scan 轴附加分区元数据，对应 Linen 的 metadata_params={nn.PARTITION_NAME: ...}。
graph / graph_updates 控制对象图结构和变量更新是否在 transform 内传播；当需要 graph updates 时，NNX 会把 graphdef、carry、broadcast、updates 打包成纯 pytree 后调用 jax.lax.scan，再把更新应用回对象图。
NNX scan 会在进入 jax.lax.scan 前把非 0 scan 轴 moveaxis 到 0，返回后再移回，这和 Linen axes_scan.scan 的转置逻辑相似，但对象状态的拆分/合并方式不同。

MaxText 当前 NNX Decoder 路径没有直接使用 nnx.scan 装饰器，而是在 _apply_layers_sequentially 中手写了等价的低层形态：先用 nnx.vmap 创建带 layer 轴的 layers state，再 nnx.split(layers, nnx.Param, ...)，必要时把 param_scan_axis 移到 0，然后调用 jax.lax.scan(layer_fn, x_in, (params, state))。layer_fn 每轮把当前层参数和 state nnx.merge 回 Module，执行 layer，再返回新的 hidden carry 和当前层 state。这个写法避开了 Linen collection API，也让参数 host offload 时可以在每层切片后显式 jax.device_put 到 device。

4. `jax.lax.scan` 是如何工作的

jax.lax.scan 的语义是：

python

carry = init
ys = []
for x in xs:
  carry, y = f(carry, x)
  ys.append(y)
return carry, stack(ys)

这里的 y 是 scan body 每一步额外产出的 per-step output，JAX 会把每一步的 y_i 沿 scan 轴 stack 成 ys。它不是必须传到下一层的状态；真正跨 step 传递的是 carry。在 MaxText decoder stack 里，hidden state 通常走 carry，最终 decoder 输出是 final_carry；y 更多对应每层辅助输出、状态集合、intermediates，或者直接是 None。

但它不是 Python 循环展开。JAX 会先 flatten init 和 xs，推断或校验 length，然后 trace body 得到 closed jaxpr。这个 jaxpr 必须返回 (carry, y) 二元结构，并且输入 carry 与输出 carry 的 pytree 结构、shape、dtype 必须完全一致。随后 JAX 绑定 scan_p primitive；scan_p 注册了 JVP、transpose、partial eval、batching、DCE 和 MLIR lowering 规则。

scan_p 的 lowering 函数是 _scan_impl。它把参数拆成 consts、carry、xs，预先创建空的 yss 输出 buffer，并维护一个循环状态。默认 unroll=1 时，这个状态可以简化理解为 (i, carry, yss)；如果有多个 xs 或 ys leaf，则它们都是 flatten 后并列放在 loop state 中。

每个 rolled while 迭代做三件事：

根据当前 i 从每个 scanned input 中取一片 x_i。
调用 body jaxpr，得到新的 carry 和本轮输出 y_i。
把 y_i 写入 yss[i]，然后 i += 1。

reverse=True 只改变读取和写入索引的方向；unroll>1 会把 xs reshape 成 (num_trips, unroll, ...)，在一个 while body 内用 Python 小循环连续执行多步 body jaxpr，并把多个 y stack 后一次更新到 yss。unroll=True 或 unroll=0 可能完全展开，不再生成 rolled while；否则 scan 的静态长度语义不变。

反向传播不是简单对 while_loop 做 reverse-mode，因为一般动态 while_loop 不能直接反向。scan 有自己的 transpose/linearization 规则：前向 scan 可以保存或重算 residual，反向通常构造一个反向方向的 scan，把 cotangent 沿时间或层轴从后往前传播。这也是 scan 比普通 while_loop 更适合固定层数 decoder stack 的原因。

5. XLA 是如何翻译 scan 的，以及核心 op 行为

这里需要区分两条路径：

JAX lax.scan 路径：JAX 的 scan_p 先通过 _scan_impl 变成 jax.lax.while_loop；随后 while_p 的 MLIR lowering 创建 stablehlo.while。进入 OpenXLA 的 StableHLO -> HLO export 时，stablehlo.while 会被转换成 XLA While。因为 XLA While 只有一个 operand，export 会把 StableHLO while 的多个 loop-carried operands 打包成 tuple；如果 while region 有隐式捕获的外部值，也会把这些值追加到 tuple state 中并在 body 里原样 carry。
OpenXLA HLO kScan 路径：OpenXLA 自身还有一个 HLO scan opcode。它不是 JAX lax.scan 的必经形态；例如 HLO parser、XLA builder 或其他 frontend 可以生成 HloOpcode::kScan。后端 pipeline 中的 ScanExpander 会把匹配到的 kScan 展开成 kWhile。非 associative scan 在 layout-sensitive verifier 之前必须被展开；associative scan 可以在某些 pipeline 配置下保留为 kScan，由后续 pass 特殊处理。

所以对 JAX decoder scan 来说，StableHLO/XLA 层面通常看不到一个需要再由 XLA 解释的“JAX scan op”，而是已经是 while。如果入口是 OpenXLA kScan，ScanExpander 展开后的结构和 JAX lowering 的 rolled loop 很接近。

OpenXLA ScanExpander 展开的 kScan loop state 是：

text

(index, accumulators..., inputs..., output_arrays...)

其中 accumulators 对应 scan carry，inputs 是完整 scanned input，output_arrays 是预先初始化的输出 buffer。body 的实际行为是：

用 GetTupleElement 从 tuple state 里取出 index、carry、inputs 和输出 accumulator。
如果是 reverse scan，用 scan_dim_size - 1 - index 算出访问位置；否则直接用 index。
用 DynamicSlice 从每个 scanned input 的 scan dimension 切出当前 step。slice size 在 scan 维度上是 1，随后用 Reshape 删除这个 size-1 scan 维度，变成 body to_apply 需要的单步输入。
调用 to_apply computation，得到当前 step 的 outputs 和新的 carry。ScanExpander 随后会 inline 这个 call，使 body 内直接包含原 to_apply 计算。
对每个 per-step output，先用 Reshape 加回 size-1 的 scan 维度，再用 DynamicUpdateSlice 写入 output_arrays 的当前位置。
用 Add 计算 index + 1，再用 Tuple 形成下一轮完整 loop state。

结合 JAX lax.scan，核心 op 行为可以整理为：

stablehlo.while / XLA While：承载 rolled loop。XLA While 的 operand/result 通常是 tuple-shaped loop state，cond computation 返回是否继续循环，body computation 返回下一轮完整 loop state。
compare：通常是 i < length 或 i < num_trips，为 while cond 生成布尔谓词。
get-tuple-element：在 XLA tuple state、while result 或 call result 中取出 index、carry、scanned inputs、输出 accumulator、hoisted values 等分量。
dynamic-slice：从 scanned input 的 scan 轴取当前层或当前时间步。对 rank 大于 1 的张量，起始 index 是 [0, ..., i, ..., 0]，其中 scan 维度使用当前 index；slice size 在 scan 维度上通常是 1。
reshape：dynamic-slice 后删除 size-1 的 scan 维度；写输出前再把单步 y_i 扩成带 size-1 scan 维度的 slice。
call / inlined computation：调用被 trace 的 body 或 HLO to_apply。优化流水线可能把它 inline，但语义上它消费当前 carry 和 x_i，产生 carry_{i+1} 与 y_i。
dynamic-update-slice：把本轮输出 y_i 写入输出 accumulator 的当前位置。HLO 语义上它返回更新后的新数组；后续 copy insertion、buffer assignment 和后端 codegen 可以在安全时把它实现为同一 buffer 的循环内更新。
add / subtract：add 更新 loop index；reverse scan 还会用 subtract 计算反向访问 index。
tuple：把下一轮 index、carry、inputs 和输出 accumulator 重新打包为下一次 while body 的 state。

因此，Decoder scan 在 HLO 中的实际行为是“一个固定 trip count 的 while 循环”：每轮从带 layer 轴的参数和逐层输入中动态切片，执行一层 decoder body，把 hidden state 作为 carry 传给下一轮，并把需要 stack 的输出写回 accumulator。它不是复制 N 份 layer HLO；除非选择完全 unroll，层数主要体现在 while trip count 和带 layer 轴的数组形状上。

`nn.remat`：在 Module 边界上提升 `jax.checkpoint`

1. MaxText 是如何使用 `nn.remat` 的

MaxText 的 Linen Decoder 会先决定 remat_policy，再把 decoder layer class 包成 rematted layer，最后再交给 nn.scan 或普通逐层调用。这个顺序很重要：nn.remat 的边界是“单层 decoder body”，nn.scan 的边界是“多层 loop”。因此在 scan_layers=True 时，实际结构可以理解为：

python

policy = decoder.get_remat_policy()
RemattedDecoderLayer = nn.remat(
    DecoderLayer,
    policy=policy,
    prevent_cse=should_prevent_cse_in_remat(cfg),
    static_argnums=(4, 5),
)

ScanDecoderLayer = nn.scan(
    RemattedDecoderLayer,
    variable_axes={"params": ScanIn(cfg.param_scan_axis), "cache": 0},
    split_rngs={"params": True, "dropout": cfg.enable_dropout},
    length=cfg.num_decoder_layers,
)

y, _ = ScanDecoderLayer(...)(y, mask, positions, segment_ids)

在这个结构中，每一层前向仍然正常计算；nn.remat 改变的是反向传播时保存哪些 residual。MaxText 通过 checkpoint_name 给关键激活打名字，policy 再按名字决定保存、offload 或重算。典型命名点包括：

Decoder 输入：decoder_layer_input。
Attention 投影：qkv_proj、query_proj、key_proj、value_proj、attention_out、out_proj。
MLP 投影：mlpwi、mlpwi_0、mlpwi_1、mlpwo。

例如 remat_policy="minimal" 会保存少量投影和 MLP 输出，未命名或不在白名单里的中间值在反向重算；minimal_with_context 会额外保存 context；qkv_proj_offloaded 会把 Q/K/V 投影 residual 从 device offload 到 pinned host；minimal_offloaded 会 offload 更大的投影/MLP residual 集合；custom 则由 tensors_on_device 和 tensors_to_offload 明确指定。

如果 parameter_memory_host_offload=True，MaxText 在 nn.remat 之前先用 nn.map_variables 包一层参数搬运，把进入该 layer 的 "params" jax.device_put 回 device。它处理的是参数位置；activation offload 则由 remat policy 处理，二者可以叠加，但成本来源不同。

2. Flax 是如何翻译 Linen `nn.remat` 的

nn.remat 是 nn.checkpoint 的别名，Linen 入口是 flax.linen.transforms.checkpoint。它和 nn.scan 一样是 lifted transform：JAX 只知道纯函数，Flax 需要把 Module 的变量集合、RNG 和方法调用转换成可传给 jax.checkpoint 的函数。

翻译过程可以理解为：

flax.linen.transforms.checkpoint 接收 target、variables、rngs、prevent_cse、static_argnums、policy 等参数。
因为用户看到的 Linen method 第一个参数是 self，而 lifted 后交给 JAX 的纯函数不再显式包含 self，Linen 会先把用户传入的 static_argnums 减 1。
lift.checkpoint 再把变量集合和 RNG 打包为 variable_groups、rng_groups，构造临时 Scope，调用原始 Module body，最后把更新后的变量树 repack 出来。
在真正调用 jax.remat 前，lift.checkpoint 又把 static_argnums 加 2，因为底层 rematted 函数额外多了 variable_groups 和 rng_groups 两个前置参数。
jax.remat 返回的仍然是一个数值等价的函数：前向输出和未 remat 时相同，但在 autodiff partial evaluation 时改变 residual 的保存策略。

把它套回 MaxText 的单层 Decoder，可以把 Linen lifted 形态近似写成：

python

def linen_layer(scope, y, mask, positions, segment_ids, deterministic, model_mode):
  y = checkpoint_name(y, "decoder_layer_input")
  qkv = checkpoint_name(attention_qkv(y), "qkv_proj")
  out = checkpoint_name(attention_out(qkv, mask), "out_proj")
  mlp = checkpoint_name(mlp_block(out), "mlpwo")
  return y + out + mlp

def lifted_remat(variable_groups, rng_groups, y, mask, positions, segment_ids, deterministic, model_mode):
  scope = make_scope(variable_groups, rng_groups)
  y_next = linen_layer(scope, y, mask, positions, segment_ids, deterministic, model_mode)
  return y_next, repack_fn(scope)

nn.remat(..., static_argnums=(4, 5)) 的含义是：在 layer 调用参数中，deterministic 和 model_mode 是静态值，JAX trace 时要把它们当作编译期分支条件。Flax lifted 后会修正这些 argnum，使底层 jax.remat 看到的静态位置仍然对应这两个业务参数。

3. Flax NNX remat 与 MaxText NNX 路径的差异

NNX 也提供 nnx.remat，语义上是 jax.checkpoint 的 lifted 版本，但对象模型不同。Linen remat 围绕 Scope、变量 collection 和 RNG collection 打包；NNX remat 围绕对象图拆分与合并：

tree-mode 或简单路径会用 SimpleRematFn 把 NNX Module/Variable 转成 pytree，再调用 jax.checkpoint。
graph updates 路径会先 merge_inputs，在 checkpoint 内执行用户函数，再 split_inputs 把对象图更新传播出来。
graph / graph_updates 控制共享引用、对象图结构更新和变量更新是否参与 transform。

MaxText 当前 NNX Decoder 路径多数地方没有直接调用 nnx.remat 装饰器，而是先 nnx.split 得到 graphdef、params、state，写出纯函数，再直接 jax.checkpoint：

python

def layer_fn(carry, scanned_vars):
  current_params, current_state = scanned_vars
  layer = nnx.merge(graphdef, current_params, current_state)
  y_next = layer(carry, *args, **kwargs)
  return y_next, nnx.state(layer)

layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
final_y, scanned_state = jax.lax.scan(layer_fn, y, (params, state))

这和 Linen nn.remat + nn.scan 的底层形态很接近：scan 每轮切出当前层参数和 state，checkpoint 包住单层 body，policy 决定该 body 反向时保存哪些 residual。差异在于 NNX 的对象状态由 nnx.split/merge/state/update 显式处理，而不是由 Linen Scope 和 collection filter 隐式处理。

4. `jax.checkpoint` / `jax.remat` 是如何工作的

JAX 的 jax.checkpoint 与 jax.remat 是同一个机制。它先把目标函数 flatten 后 trace 成 jaxpr，再绑定 remat_p primitive。remat_p 的参数里最关键的是：

jaxpr：checkpoint 边界内的前向计算。
prevent_cse：是否在 differentiated lowering 中对部分输入加 optimization barrier，防止 HLO CSE 把前向值和反向重算合并。
policy：partial evaluation 时如何处理可作为 residual 的中间值。
differentiated：标记当前 remat primitive 是否处在微分变换产生的反向/线性化上下文里。

普通 reverse-mode autodiff 会保存许多前向线性化点；remat 的策略是只保存边界输入，以及 policy 允许保存的命名 residual，其余需要时在反向 jaxpr 中重算。policy 返回值可以理解为三类：

Saveable：把这个值作为 residual 保存到反向。
Recompute：不保存，反向重新执行产生它的计算。
Offloadable(src="device", dst="pinned_host")：前向把 residual 从 device 复制到 host memory space，反向使用前再 reload 回 device。

checkpoint_name(x, name) 在运行时是 identity，但 trace 中会给值附上名字。save_only_these_names(...) 和 save_and_offload_only_these_names(...) 都依赖这些名字做决策。因此 remat_policy 不是“是否重算”的单个开关，而是 remat 边界内的一张 residual 分配表。

5. XLA / StableHLO 中 remat、CSE barrier 和 offload 的表现

remat_p 本身不是一个长期保留到最终 HLO 的高层 op。JAX lowering 会把 remat 内部 jaxpr 正常 lower 成它包含的算子；前向数值计算通常和未 remat 时一致。差异主要出现在 autodiff 和 partial evaluation 之后：

被要求重算的值不会作为前向 residual 穿到反向，而是在反向 jaxpr 中重新出现对应的计算。
prevent_cse=True 且处在 differentiated lowering 中时，JAX 会插入 OptimizationBarrierOp。这会在 StableHLO/XLA 中表现为 optimization barrier，用来阻止优化器把“本该反向重算”的计算和前向已有计算做 CSE 合并。
在 scan 内，MaxText 通常让 should_prevent_cse_in_remat 返回 False，因为 scan 的 while 边界已经使 CSE 跨迭代/跨前反向合并的收益和风险不同，额外 barrier 往往只是开销。Ling3 的长度 1 MLA MoE prefix 是例外：while 可能被 inline，代码会重新包一层 nn.remat(..., prevent_cse=True) 让 barrier 留下来。
Offloadable residual 在 JAX jaxpr 层面会插入 device_put_p：前向把值从 device memory space copy 到 pinned_host，反向使用前再 copy 回 device。StableHLO/compiled program 中不一定表现为用户可见的 Python copy 调用，而是带 memory kind 的 host/device 数据流。

所以从 XLA 角度看，remat 的重点不是新增一个“RematOp”，而是改变前向 residual 是否进入反向、反向是否重新计算、以及某些 residual 是否经过 host/device memory space。它和 scan 组合后，常见形态是：外层 WhileOp 管理层循环，内层 layer body 在反向中按 remat policy 保存、offload 或重算激活。

独立项目：NNX scan decoder stack

这个项目可以独立于 Linen nn.scan 推进。目标不是替换 jax.lax.scan，而是把 decoder layer stack 的状态管理从 Linen collection transform 改成 NNX 的显式对象图拆分/合并，使 scan body 的输入、carry 和参数搬运路径都可控。它适合承载两类需求：

继续保留 rolled scan 的编译收益。
在每层 body 周围插入更细粒度的参数 host/device movement、FSDP all-gather prefetch、double buffer 或自定义 remat 边界。

当前 MaxText NNX 形态

MaxText NNX decoder 当前没有直接使用 nnx.scan 装饰器，而是采用“NNX 建层 + 手写 jax.lax.scan”的低层形态：

python

layers = nnx.vmap(create_layer, out_axes=StateAxes({Param: param_scan_axis, ...: 0}))
graphdef, params, state = nnx.split(layers, nnx.Param, ...)
params = move_param_scan_axis_to_front(params)

def layer_fn(hidden, scanned_vars):
  current_params, current_state = scanned_vars
  current_params = maybe_device_put_current_layer_params(current_params)
  layer = nnx.merge(graphdef, current_params, current_state)
  hidden_next = layer(hidden, ...)
  return hidden_next, nnx.state(layer)

final_hidden, scanned_state = jax.lax.scan(layer_fn, hidden, (params, state))

这里有三个关键点：

graphdef 是静态对象图，描述 layer 的 Python 结构和调用形态；它不沿 scan 轴变化。
params 和 state 是带 layer 轴的 NNX state。params 中的 nnx.Param 叶子按 param_scan_axis stack，进入 lax.scan 前会把这个轴移到 0。
每轮 body 拿到当前层的 current_params/current_state，用 nnx.merge 临时还原成一层可调用对象，执行后再用 nnx.state(layer) 抽回本层更新后的 state。

这和 Linen nn.scan 的结果很像：底层仍是一个 jax.lax.scan / while，每轮执行同一份 layer body。但控制边界不同：Linen 把变量 collection、RNG、metadata、axis 转换隐藏在 lifted transform 内；NNX 路径把参数 pytree、state pytree 和 scan carry 直接暴露给 MaxText。

`nnx.scan` 装饰器与手写 `lax.scan`

Flax NNX 的 nnx.scan 是 jax.lax.scan 的 NNX transform。它的 API 用 in_axes/out_axes 表达三类输入：

nnx.Carry：跨 step 传递并更新，例如 hidden state。
整数 axis：沿该轴扫描，例如 stack 后的 layer object/state。
None：每轮复用的辅助输入。

当输入是 NNX graph node 或 state 时，StateAxes 可以进一步描述不同变量类型的轴，例如 nnx.Param 沿 param_scan_axis，其他状态沿 0；transform_metadata={nnx.PARTITION_NAME: "layers"} 则给新增的 layer 轴附加分区元数据。开启 graph updates 时，NNX 会把 graphdef、carry、broadcast、updates 拆成纯 pytree，调用 jax.lax.scan，然后把输出合并回对象图。

因此可以把选择分成两层：

使用 nnx.scan 装饰器：更接近 NNX 官方 API，适合标准“逐层执行同构 Module”的场景。
保持 MaxText 当前手写 jax.lax.scan：代码更显式，适合参数 offload、下一层权重预取、手动 double buffer、特殊 remat 边界等系统优化。

对 MaxText decoder 来说，第二种形态更适合作为性能实验项目，因为它允许直接改变 scanned_vars 和 carry 的结构，而不需要改 Flax transform 本身。

与 Linen `nn.scan` 的核心差异

维度	Linen `nn.scan`	NNX scan 项目
状态抽象	`Scope` + variable collections	`GraphDef` + `State` + 显式 `Param` state
参数轴	`variable_axes={"params": ScanIn(axis)}`	`StateAxes({nnx.Param: axis})` 或手写 `moveaxis`
body 输入	Flax 默认只给当前层变量切片	MaxText 可以自定义 `scanned_vars`，例如同时传当前层和下一层参数
参数搬运	常通过 `nn.map_variables` 或 layer 内 `device_put`	body 内可对当前 slice 显式 `device_put` / sharding constraint / all-gather
metadata	`metadata_params` 自动 remove/add axis	`transform_metadata` 或手写 state metadata 管理
remat 边界	`nn.remat(layer)` 再 `nn.scan`	`jax.checkpoint(layer_fn)`，可进一步拆分 prefetch 与 compute
扩展性	高层稳定，但不容易改 body 协议	低层显式，更适合通信/计算 overlap

下一层权重预取的项目设计

如果目标是在第 n 层计算时提前 all-gather 第 n+1 层权重，NNX scan 项目应把“已 gather 的当前层权重”作为 carry，而不是只把参数作为 xs。原因是：只有下一轮真实使用这个值，XLA 才不能把预取当成无用计算删掉或任意重排。

可以把 scan 状态改成：

python

next_params = shift_left(params, pad_with_last_or_dummy=True)
prefetched_0 = gather_or_device_put(params[0])
carry = (hidden, prefetched_0)
xs = (params, next_params, state)

def layer_fn(carry, xs_i):
  hidden, prefetched_current = carry
  current_params, next_params_i, current_state = xs_i

  prefetched_next = gather_or_device_put(next_params_i)
  layer = nnx.merge(graphdef, prefetched_current, current_state)
  hidden_next = run_layer_compute(layer, hidden)
  current_state_next = nnx.state(layer)
  return (hidden_next, prefetched_next), current_state_next

这个结构表达的是 double buffer：

第 n 轮 compute 使用从上一轮 carry 传来的 prefetched_current。
同一轮尽早构造 prefetched_next，让 all-gather / device transfer 有机会和当前层 compute overlap。
第 n+1 轮通过 carry 使用 prefetched_next，从而形成真实依赖。

如果预取的是 host offload 参数，gather_or_device_put 可以是把当前层参数放回 device memory space；如果预取的是 FSDP-sharded 参数，它应该是把当前单层参数约束到 no-FSDP / replicated layout，让编译器插入 all-gather。后一种最好复用已有 sharding 工具中“去掉 FSDP 轴”的思路，而不是直接在业务层手写 mesh collective。

remat 和 residual 的边界

NNX scan 项目里要避免把 prefetch 全部包进同一个 remat body。更稳的拆法是：

python

prefetched_next = prefetch(next_params_i)
hidden_next, state_next = checkpointed_layer_compute(prefetched_current, hidden, current_state)

这样 remat 只控制单层 compute 的 activation residual，prefetch 仍然作为 scan carry 的通信/搬运值存在。否则反向重算可能重新触发参数 gather，或者把本该 overlap 的通信塞进 checkpoint 重算路径。

风险与验证点

Carry shape 必须稳定：prefetched_current 和 prefetched_next 的 pytree、shape、dtype、sharding 必须每轮一致。
最后一层不能越界：next_params 需要 dummy、重复最后一层，或用 index guard；如果 dummy 会进入 HLO，需要确认不会产生额外无用 all-gather。
预取不能被 DCE：prefetched_next 必须进入下一轮 carry 并被下一轮 compute 使用。
prefetch 与 compute 是否真的 overlap 不能只看 Python 代码，要看 HLO / profiler。期望看到 all-gather 或 host/device copy 在 while body 内早于当前层主要 matmul，并且下一轮使用 carry 里的结果。
remat policy 需要重新验证：引入参数 carry 后，scan partial-eval 可能多出 residual 或改变 buffer lifetime；需要同时检查 peak memory、compile time、step time 和 backward HLO。

`scan` 与 `remat` 如何共同决定 residual

Decoder stack 里 residual 的保存不是只由 nn.remat 决定，也不是只由 nn.scan 决定，而是两层 partial-eval 规则叠加后的结果。可以把它拆成两步看：

remat 先在“单层 body”内部决定哪些中间值允许成为 residual。
scan 再决定这些 residual 在“多层 loop”里如何穿过前向和反向。

以 nn.scan(nn.remat(DecoderLayer)) 为例，JAX trace 后大致是：

python

def scanned_body(y, layer_vars):
  params_i, cache_i = layer_vars

  def rematted_layer(y):
    y0 = checkpoint_name(y, "decoder_layer_input")
    qkv = checkpoint_name(qkv_proj(params_i, y0), "qkv_proj")
    out = attention_and_mlp(params_i, qkv, y0)
    return out

  y_next = jax.checkpoint(rematted_layer, policy=policy)(y)
  return y_next, None

final_y, _ = jax.lax.scan(scanned_body, y0, scanned_layer_vars)

remat 看到的是单层 rematted_layer。它根据 policy 把该层内部值分成三类：

Saveable：保存为 residual，例如 minimal 策略下的部分投影输出。
Offloadable：保存为 residual，但前向放到 pinned_host，反向再搬回 device。
Recompute：不保存，反向重新执行对应计算。

但这些 residual 进入 scan 后，还会再被 scan 的规则分类。JAX 的 scan transpose/partial-eval 代码把 residual 分成两种：

intensive residual：不沿 scan 轴变化，可以作为反向 scan 的 const 传入，不需要按层 stack。JAX transpose 注释里把它写成 ires，即 intensive (not scanned over / const) residuals。
extensive residual：每个 scan step 都可能不同，必须按 scan 轴收集，形状会多出一个 leading layer/time 维度，并作为反向 scan 的 scanned input 传入。JAX transpose 注释里把它写成 eres，即 extensive residuals。

这不是按 residual 的名字、大小或 dtype 直接分类，而是 partial-eval 之后由“它是否能脱离当前 scan step 计算”决定。remat_policy 只回答“单层 body 内这个值要不要保存/是否 offload”，scan 进一步回答“如果保存，这个值是每层一份，还是 loop 外一份”。

1. JAX 如何识别 intensive 和 extensive residual

JAX 的普通 scan partial-eval 规则先把输入 tracer 按 consts、init carry、xs 分组，并用 t.pval.is_known() 得到 unknowns。然后它做一个 carry fixpoint：如果某个 carry 输入未知，或者 body 输出 carry 变成未知，就把对应 carry 标记为未知；循环直到 carry_uk 稳定。这个 fixpoint 很关键，因为 carry 是否未知会影响 body 内哪些值必须进入 staged/backward 部分。

这里的 unknown 不是“运行时值一定不可知”的泛化说法，而是 partial evaluation 中的一个二分：这个值能不能作为 known constant 留在 known jaxpr 里。满足下面条件时会被标记为 unknown：

输入 tracer 的 pval 本身不是 known。常见例子是 jit 的数组参数、scan 的初始 carry、scan 的 xs，以及反向/linearization 中需要 staged 的 tangent 或 cotangent。
某个 body 中间值依赖 unknown 输入。只要一个算子读了 unknown carry、unknown xs_i、动态 const 或其他 unknown 中间值，它的输出通常也会是 unknown。
某个 carry 输出依赖 unknown 值。_scan_partial_eval 会从 out_uk 里切出 carry_uk_out；如果 carry_uk_out[i] 是 true，对应 carry component 会在下一轮 fixpoint 中继续被当作 unknown 输入。
carry 的 unknown 状态是单调扩张的。代码用 carry_uk = carry_uk OR carry_uk_out 的方式迭代，所以一旦某个 carry component 被标成 unknown，就不会在同一个 fixpoint 中变回 known。

例如：

python

def body(carry, x):
  c1 = carry + x       # 依赖当前 scanned input，unknown
  c2 = jnp.ones_like(x)  # 依赖 x 的 aval/shape，但数值不依赖 x
  c3 = jnp.sin(c1)     # 依赖 unknown c1，unknown
  return (c1, c2), c3

如果 x 是 scan 的 scanned input，那么 c1 会是 unknown；c3 也会因为依赖 c1 进入 unknown/staged 路径。c2 是否能留在 known 侧取决于它是否只需要 aval/shape、是否被实例化为数组输出，以及后续是否被需要；如果它最终作为 residual 且只依赖 scan 外 known const，才有机会被 hoist 成 intensive residual。MaxText 的 hidden state carry 基本都会沿层变化，因此依赖 y_i 的 decoder_layer_input、qkv_proj(params_i, y_i) 等值通常会沿 unknown -> extensive residual 路径传播。

随后 pe.partial_eval_jaxpr_nounits_fwd 把 scan body 分成 known jaxpr 和 unknown jaxpr，并返回两类 residual 信息：

res_avals：known jaxpr 产生、unknown jaxpr 需要的非 forwarded residual。
in_fwd_res：可以从已有输入直接 forward 到 unknown jaxpr 的 residual。scan 只允许从 const 或 xs forward，不允许从 carry forward，也不允许从 Python np.ndarray forward。

在这一刻，所有非 forwarded residual 会先被当作 scan 的 extensive output：也就是 known scan 每一轮产出一份 residual，外层暂时按 [length, ...] 收集。之后 JAX 才尝试 hoist。最终 unknown jaxpr 的输入会被重排成：

text

[*int_res, *unknown_inputs, *ext_res]

其中 int_res 是 intensive residual，作为 scan const 进入 staged/transpose scan；ext_res 是 extensive residual，被移动到 scan input 的末尾，按照 scan calling convention 作为 scanned input 进入每一步。transpose 规则 _scan_transpose_fancy 也沿用这个调用约定：它先按 num_ires = len(consts_lin) - sum(consts_lin) 从 args 前缀切出 ires，再把末尾的 eres 和 cotangent 一起传入反向 scan。

所以识别逻辑可以简化成：

能从 non-carry const 直接 forward，或能被 hoist 成 const 的 residual，是 intensive。
需要由每个 scan step 产生，或 forward 来源是 scanned input xs 的 residual，是 extensive。
carry 相关 residual 不会被当作普通 input forward；如果它被 unknown/backward 需要，通常会落到 extensive 路径，因为每一层 carry 都不同。

2. `_scan_known_hoisting` 的判断依据

scan 的 partial-eval 规则有一个专门优化：先把非 forwarded residual 放在 known scan 的输出尾部，再调用 _scan_known_hoisting 尝试把 loop-invariant residual 提到 scan 外。它的判断方式不是做语法级依赖分析，而是重新 trace known jaxpr：

原 known const 输入被包装为 PartialVal.known(c)；如果 const 是 mutable ref，则保守地作为 unknown。
carry、xs 和其他非 const 输入被包装为 PartialVal.unknown(aval)。
known jaxpr 的普通输出要求 instantiate；residual 输出不要求 instantiate。
重新 trace 后，如果某个 residual 输出的 PartialVal 仍然是 known，说明它只依赖 known const，可以 hoist；否则说明它依赖当前 loop step 的输入，只能留在 loop 内。

因此 hoist 的充分条件是：这个 residual 可以只用 scan 外已知 const 计算出来。它不能依赖当前层参数切片、当前 carry、当前 xs、当前 dropout RNG 或当前 step 产生的激活。对于 custom partial-eval 规则，JAX 也做同类优化：把 known jaxpr 再分成 hoist 部分和 loop 部分，用 loop_dep_res 判断 residual 是否依赖 loop；loop_dep_res=False 的 residual 被放到 intensive_res，loop_dep_res=True 的 residual 被映射成带 leading scan 轴的 extensive_res。

一个 MaxText 例子：

python

def scanned_body(y, layer_vars):
  params_i, cache_i = layer_vars
  qkv = checkpoint_name(qkv_proj(params_i, y), "qkv_proj")
  mask_bias = checkpoint_name(make_mask_bias(mask), "mask_bias")
  return layer_out(qkv, mask_bias, y), None

如果 mask 是 scan 的 broadcast const，make_mask_bias(mask) 又不读 params_i 或 y，那么它有机会被 hoist 成 intensive residual：反向 scan 只需要拿一份 mask_bias。但 qkv 同时读当前层参数 params_i 和当前 carry y，每一层都不同，不能 hoist；如果 remat policy 要保存它，它会变成 extensive residual，最终形状等价于 [num_layers, batch, seq, ...]。

3. 是否依赖当前层或当前 carry 的影响

判断一个 residual 的 scan 代价，可以直接问它在 body 中依赖什么：

residual 依赖	JAX 分类倾向	结果
只依赖 scan 外 const / broadcast 输入	intensive	只保存一份，作为反向 scan const；不产生 layer 轴 residual buffer
依赖当前层参数切片 `params_i` / 当前 `xs_i`	extensive	每层一份，保存或 offload 时都会沿 scan 轴组织
依赖当前 carry `y_i`	extensive	每层 carry 不同，保存成本接近保存所有层边界 hidden states
依赖当前 dropout RNG / 当前 cache state / 当前中间激活	extensive	每 step 不同，不能 hoist
从 scan 的 const input 直接 forward	intensive	不需要由 known scan 每轮重新产出
从 scan 的 `xs` input 直接 forward	extensive	即使只是 forward，也带 scan 轴

这解释了 MaxText 中几个容易混淆的点：

对 MaxText 的 decoder stack 来说，常见结果如下：

值	remat policy 决策	scan 决策	结果
`qkv_proj`	`minimal` 可能保存	依赖每层参数和每层输入	extensive residual，按 layer 轴保存
`qkv_proj`	`qkv_proj_offloaded` offload	依赖每层参数和每层输入	extensive residual，但每层 residual offload 到 host
`decoder_layer_input`	custom 若保存	每层 carry 不同	extensive residual，保存会很贵
broadcast mask / positions 派生值	若被命名并保存	不依赖当前层或 carry 时可 hoist	intensive residual，只保存一份
静态配置 / mesh / mode	不作为数组 residual	loop-invariant	编译期或 const 处理

这个机制解释了几个配置上的现象：

minimal_with_quantization 在 scan_layers=True 下可能更慢：policy 允许保存更多命名值后，这些值大多是每层不同的 extensive residual，会额外形成按 layer 轴 stack 的 residual buffer。
decoder_layer_input 不适合随意加入自定义保存或 offload：它是每层 carry，scan 下通常是 extensive residual；保存它相当于保存所有层边界的 hidden states。
Offload 的粒度也是按 residual 决定的：如果某个 named residual 在 scan 中是 extensive，那么 offload 的不是一个单层值，而是沿 scan 轴组织的一组每层 residual。省 HBM 的同时，会增加前向 offload 和反向 reload 的数据流。

因此，判断一个 MaxText residual 的显存影响，需要同时问三个问题：它是否被 checkpoint_name 命名并被 remat_policy 选中；它在 scan body 中是否依赖当前层或当前 carry；如果被保存或 offload，它会是 intensive 还是 extensive。只有第三个问题回答完，才能估算它最终是“一份常量 residual”，还是“每层一份 residual buffer”。

4. 非 scan / Python `for` 场景下 remat 为什么可能“失效”

严格说，jax.remat 在非 scan 场景下不会在 JAX partial-eval 语义上失效。remat_partial_eval 仍然会调用 pe.partial_eval_jaxpr_custom，按 policy 产出 known jaxpr、staged jaxpr 和 residual；被判为 Recompute 的值不会作为 JAX residual 穿到 backward，而是在 staged/backward jaxpr 中重新计算。问题出在下一层：没有 scan lowering 出来的 WhileOp 边界时，Python for 在 jit 下通常被完全展开，前向层和反向重算会落在同一个大 HLO 图里。

JAX 对 remat 防止这种情况的机制是 prevent_cse。remat_p lowering 只有在 differentiated=True 且 prevent_cse=True 时，才会对相关输入插入 OptimizationBarrierOp。这个 barrier 的作用是阻止编译器把反向里“本来应该重算”的表达式和前向里已有的等价表达式做 common subexpression elimination。如果 prevent_cse=False，JAX jaxpr 仍然表达了“反向重算”，但 XLA 优化器可能把前向计算和反向重算合并成同一个 HLO 值。这样一来，前向产生的 buffer 会因为反向仍要引用而延长 live range，看起来就像 remat 没有释放前向激活。

scan 场景下 MaxText 通常关闭 prevent_cse，因为 rolled scan 会形成 WhileOp 边界，scan transpose/partial-eval 又会把需要跨前反向的值显式分成 ires 和 eres。编译器不能像完全展开的 for loop 那样在一个平铺图里随意把所有层的前向值和反向重算合并；真正需要保留的 extensive residual 会表现为 scan 输出/输入 buffer，不需要保留的值则在反向 scan body 内重算。

但这也有边界条件：如果 scan 长度为 1 或被优化 inline，while 边界可能消失，CSE 又有机会跨过前向和反向重算。MaxText 的 Ling3 MLA MoE prefix 就显式注释了这个问题：长度 1 scan 可能被 inline，导致 forward gather 和 backward reference 被 CSE 合并，从而 buffer 无法在 full remat 下释放；因此代码重新包一层 nn.remat(..., prevent_cse=True)，让 optimization barrier 保留下来。

所以“for loop 下 remat 失效”的准确机制是：Python for 被展开后缺少 WhileOp 边界，如果 prevent_cse 没有打开，JAX remat 产生的 backward recompute 可能被 XLA CSE 优化成复用 forward 结果，forward buffer 的 liveness 延伸到 backward。打开 prevent_cse 时，JAX 会插入 barrier，代价是减少编译器优化空间；使用 scan 时，loop 边界和 scan 自己的 residual 分类通常承担了这部分隔离作用。

HLO / sharding 要点

nn.scan / jax.lax.scan 在 HLO 中主要体现为一个 WhileOp，而不是 N 个复制的 layer body。cond region 判断迭代计数，body region 读取第 i 层参数和输入、执行 layer jaxpr、更新 carry 和输出 buffer。OpenXLA 自身的 HLO kScan 是另一条入口，通常由 ScanExpander 展开成同类 kWhile 结构。
scan 的 layer 参数通常多一个 layers 维度；param_scan_axis 决定这个维度放在哪里。为了让 scan/offload 更直接，JAX host-offloading 文档也强调参数 offload 与 scan 轴 0 的配合最自然；MaxText NNX 路径在 param_scan_axis != 0 时会先 moveaxis 到 0 再 scan，结束后再移回。
metadata_params={nn.PARTITION_NAME: "layers"} 让 Flax 变量树里的 scan 轴携带分区名；后续逻辑轴到物理 mesh 的映射可以把层轴用于 FSDP/参数切分，而不需要业务代码手写 device mesh 细节。
activation offload 不一定表现为显式 Python 层面的 host copy；在 JAX jaxpr 中它是 remat partial-eval 根据 Offloadable 插入的 device_put_p，在 StableHLO/compiled program 中则体现为带 memory kind 的 host/device 数据流。
parameter_memory_host_offload 是另一条路径：Linen set_remat_policy 会先用 nn.map_variables 在进入 remat layer 前把参数 device_put 到 device；NNX scan body 也会在每层切片后把当前参数放回 device。这和 activation remat/offload 可以叠加，但成本模型不同。

坑与权衡

scan body 必须同构：carry 的 shape/dtype 不能随层变化，参数树结构也要一致。异构 attention pattern、Dense/MoE 混排、prefix/remainder 层需要拆成多段 scan 或保留 unscanned 路径。
Python for 在 jit 下容易被展开，编译图随层数线性增长；scan 降低编译压力，但会引入 loop carry、动态切片/更新和更复杂的 transpose/backward 结构。
prevent_cse=True 能防止优化器把 remat 重算合并掉，但在 scan 内通常是额外开销。MaxText 默认在 scan_layers=True 时关闭它；Ling3 的长度 1 MLA MoE prefix 是例外，因为 while 可能被 inline，代码中重新包了一层 nn.remat(..., prevent_cse=True)。
offload 省 HBM，但引入 device ↔ pinned host 传输。适合体积大、反向前暂时不用、重算成本也高的 residual；如果张量很小或很快被反向读取，传输延迟可能抵消收益。
minimal_with_quantization 等策略在 scan layers 下有显式 warning：扫描层可能给被 checkpoint 的 quantization 值带来额外开销，某些配置下甚至比不保存更慢，需要按模型和硬件实测。
decoder_layer_input 在自定义策略里不能随意 remat：配置注释说明它是周期性 checkpoint 起点。如果把起点也要求重算，反向可能需要跨过过大的计算区间，显存和重算时间都会失控。

版本依据

MaxText：source snapshot submitter-v0.3-87-g7bcf1aaf。
Flax：0.12.7，source snapshot v0.12.7-1-gf9b51ef6。
JAX：0.10.1 development source，source snapshot jax-v0.10.0-715-g54e9c65432。
OpenXLA / XLA：source snapshot 3d312858c2。

Decoder 堆叠与重计算 / Offload ​

逐原语分析 ​

nn.scan：把同构 Decoder 层提升为单个 loop primitive ​

1. MaxText 是如何使用 nn.scan 的 ​

2. Flax 是如何翻译 Linen nn.scan 的 ​

3. Flax 新的 nnx.scan 的差异 ​

4. jax.lax.scan 是如何工作的 ​

5. XLA 是如何翻译 scan 的，以及核心 op 行为 ​

nn.remat：在 Module 边界上提升 jax.checkpoint ​

1. MaxText 是如何使用 nn.remat 的 ​

2. Flax 是如何翻译 Linen nn.remat 的 ​

3. Flax NNX remat 与 MaxText NNX 路径的差异 ​

4. jax.checkpoint / jax.remat 是如何工作的 ​

5. XLA / StableHLO 中 remat、CSE barrier 和 offload 的表现 ​

独立项目：NNX scan decoder stack ​

当前 MaxText NNX 形态 ​

nnx.scan 装饰器与手写 lax.scan ​

与 Linen nn.scan 的核心差异 ​

下一层权重预取的项目设计 ​

remat 和 residual 的边界 ​

风险与验证点 ​

scan 与 remat 如何共同决定 residual ​

1. JAX 如何识别 intensive 和 extensive residual ​

2. _scan_known_hoisting 的判断依据 ​

3. 是否依赖当前层或当前 carry 的影响 ​

4. 非 scan / Python for 场景下 remat 为什么可能“失效” ​

HLO / sharding 要点 ​

坑与权衡 ​

版本依据 ​

Decoder 堆叠与重计算 / Offload

逐原语分析

`nn.scan`：把同构 Decoder 层提升为单个 loop primitive

1. MaxText 是如何使用 `nn.scan` 的

2. Flax 是如何翻译 Linen `nn.scan` 的

3. Flax 新的 `nnx.scan` 的差异

4. `jax.lax.scan` 是如何工作的

5. XLA 是如何翻译 scan 的，以及核心 op 行为

`nn.remat`：在 Module 边界上提升 `jax.checkpoint`

1. MaxText 是如何使用 `nn.remat` 的

2. Flax 是如何翻译 Linen `nn.remat` 的

3. Flax NNX remat 与 MaxText NNX 路径的差异

4. `jax.checkpoint` / `jax.remat` 是如何工作的

5. XLA / StableHLO 中 remat、CSE barrier 和 offload 的表现

独立项目：NNX scan decoder stack

当前 MaxText NNX 形态

`nnx.scan` 装饰器与手写 `lax.scan`

与 Linen `nn.scan` 的核心差异

下一层权重预取的项目设计

remat 和 residual 的边界

风险与验证点

`scan` 与 `remat` 如何共同决定 residual

1. JAX 如何识别 intensive 和 extensive residual

2. `_scan_known_hoisting` 的判断依据

3. 是否依赖当前层或当前 carry 的影响

4. 非 scan / Python `for` 场景下 remat 为什么可能“失效”

HLO / sharding 要点

坑与权衡

版本依据