从捕获文件中调试JAX OOM
概述
经过JIT编译的JAX步骤会将其所有暂存内存作为一个整体规划暂存块。因此，该步骤的内存就是XLA的规划，而.xplane.pb文件既包含该规划，也包含分配器的运行时事件。scripts/oom_report.py脚本可在2–5秒内同时读取这两部分信息，且仅依赖标准库：
- 规划： 每个模块的缓冲区分配信息，来自
/host:metadata。它给出了参数（以及有多少参数被捐赠了）、非别名的输出、暂存块以及堆峰值。对于峰值时刻存活的缓冲区，它还会提供大小、形状、角色和Python代码行。角色来自JAX名称栈：jvp(表示正向（保留用于反向传播），transpose(表示反向，s8自定义调用的输出是库的工作空间。 - 运行时： BFC分配器的事件。它们提供了容量、每次执行开始时的已用字节数、峰值、
AllocatorRetry等待次数、从未启动过内核的执行次数，以及最后两次执行之间的主机事件。它们还给出了最后一次执行开始时存活的内容，并按分配它们的主机事件进行分组：outputs of jit_step execution 3（步骤输出）、BufferFromHostBuffer(s32[…])（相当于device_put），或者“在捕获开始前已分配”。
报告最终会输出诊断行，每行都标记为“measured”（实测）或“inferred”（推断）。首先运行脚本：
python3 -P <此技能所在的基础目录>/scripts/oom_report.py <捕获目录或.xplane.pb文件> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <此技能所在的基础目录>/scripts/oom_report.py --self-test
默认情况下，它会分析最后一次执行从未启动内核的那个模块。如果不存在这样的模块，则选择运行过的最大模块。可以使用--module参数选择其他模块，例如eval步骤。
从诊断行到原因和修复方案
| 诊断行 | 原因 | 修复建议；节省的内存量来自报告的字节数 |
|---|---|---|
the step itself doesn't fit（步骤本身放不下） | 规划本身超出了内存池容量 | 按照峰值时刻的角色分布进行分析（见下文） |
↳ 主要角色是 forward, kept for the backward pass（正向，保留用于反向传播） | 激活值：它们随batch × seq增长；朴素注意力每层增加B·H·S·S | jax.checkpoint / 对此块进行重计算（能节省大部分正向保留的字节，代价大约是额外进行一次正向传播）；使用更小的micro-batch配合梯度累积；缩短序列长度（S²缓冲区在S减半时会缩小为原来的1/4）。融合注意力仅作为“确认B·H·S·S缓冲区已从规划中消失”的检查来提供 |
↳ 主要角色是 backward（反向） | 梯度及其临时缓冲 | 分片（FSDP）或累积；检查是否有bf16梯度的f32副本 |
↳ workspace（工作空间） | GEMM / conv库的临时存储 | 通常很小；在ROCm上，对于gfx950，每个hipBLASLt GEMM需要64 MiB |
grow by X MiB per execution（每次执行增长X MiB），= 1.00x the non-aliased outputs（等于非别名输出的1.00倍） | 程序保留了每一步的输出（跨步骤泄漏）：存活列表显示 outputs of … execution 1, 2, … | 在步骤之间的主机事件中找到持有者（例如 $<unknown> append 意味着 list.append）。然后删除引用、在步骤内部减少输出（例如只保留范数而非整个树结构），或者使用 jax.device_get。增大内存池只会延迟OOM的发生 |
no parameter is donated … can save up to X（没有参数被捐赠……最多可以节省X） | 旧参数和新参数/优化器状态同时存活 | 在状态参数上使用 donate_argnums |
other live arrays take the room; the largest … came from X（其他存活数组占据了空间；最大的……来自X） | 步骤之外的数组：放在GPU上的数据集或eval副本、预取、其他JIT输出 | 将其保留在主机内存中并流式传输，或者释放它，或者缩小其数据类型。“在捕获开始前已分配”意味着需要更早地做profiling来确定其来源 |
AllocatorRetry … longest N s（分配器重试……最长N秒），launched none（未启动任何内核） | OOM发生在捕获过程中 | 仅为证据。失败的请求通常该模块的暂存块或一个输出 |
XLA already rematerialized N instructions（XLA已经重算了N条指令） | XLA自身的重计算已在内存限制下运行过 | 预期修复方案应来自程序本身，而非XLA |
给出每个修复方案的估计节省量及其代价，并标记为“inferred”（推断）。确保估算值与内存池容量进行对比。“提高 XLA_PYTHON_CLIENT_MEM_FRACTION”只有在GPU确实有余量时，才能解决规划放不下的问题。容量远低于HBM意味着该分数被设置了；捕获文件不会记录环境变量。
JAX术语中的playbook内存结构
Ultra-Scale Playbook（“分析内存使用情况”）将训练内存分为权重、梯度、优化器状态和激活值。在报告中：
- 权重 + 优化器状态 = 步骤的参数。Adam在fp32下每个参数占用12字节（p, m, v）。被捐赠意味着没有第二份副本。
- 梯度 = 峰值时刻的
backward角色，以及步骤如果返回梯度时的输出。 - 激活值 =
forward, kept for the backward pass。使用Scan /nn.scan的模型将它们作为堆叠的[layers, …]缓冲区存放在while循环中。 - “步骤1能运行，步骤N失败”： 在PyTorch中这是由于优化器状态在第一步创建。在JAX中状态在步骤1之前就已存在，因此应寻找增长行。首次调用还会进行编译和自动调优，因此第一次执行前的峰值是编译期的。
规则
- 仅使用捕获文件。 不要重建模型或重新运行工作负载来测试修复方案：从规划的字节数进行估算，并说明如何在用户侧进行确认（
jit(...).lower(...).compile().memory_analysis()）。 - 不要向捕获目录写入内容。 XProf的转换器（
xspace_to_tool_names、memory_viewer、xprof命令行工具）会在.xplane.pb旁边写入*.hlo_proto.pb缓存文件。如果需要XProf视图，请先将.xplane.pb复制到一个临时目录中。 - 为每个数字标注类型： measured（实测）、read（读取）或inferred（推断）。
实测参考（单MI355X，小池子，2026-11-06）
| 捕获场景 | 报告显示的内容 |
|---|---|
| attention规划过大 | 需要3.72 GiB > 2.79 GiB；堆峰值81.6%为forward-kept；softmax行有5个f32[8,8,1024,1024]；暂存块3,495 MiB = 错误信息中的“3.41GiB”；AllocatorRetry等待10秒；89条重计算指令 |
| 每个步骤都保留了梯度 | 步骤开始时已使用内存316 → 1,252 MiB，每步增加104.0 MiB = 非别名输出的1.00倍；规划788 MiB可容纳；步骤之间有 $<unknown> append 操作 |
| 没有捐赠 | 需要3.88 GiB > 2.79 GiB；1,902 MiB参数 + 1,902 MiB输出，0捐赠；暂存块仅172 MiB |
| MaxText llama2-7b，8 GPU，FSDP（可容纳） | 暂存块61.9 GiB；堆峰值57%为forward-kept，扫描栈 bf16[32,4,4096,11008] ×2，通过 while.50 循环 |
局限性
- 每个文件仅一个主机： 多主机运行每个主机有一个
.xplane.pb文件，因此需要对每个文件分别运行。 - 仅HBM： 仅分解了内存空间0。主机卸载显示为
memory space 5。 - 连续性： “能容纳”比较的是字节数。BFC还需要一个连续的内存块，因此接近容量限制时仍可能因碎片化而失败。
- “other”角色： 由XLA创建的没有JAX名称栈、也没有一个附带名称栈的别名操作所定义的缓冲区。
- CUDA： 未经测试。它使用相同的BFC和XLA事件，但名称未经验证。
- TPU / eager模式： TPU捕获仅提供规划总量，没有堆峰值和执行信息。Eager（逐操作）模式程序会显示为许多小模块，因此请阅读运行时部分。