TensorFlow XLA:TPU 编译报 Ran out of memory in memory space hbm 怎么排查【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow当 JAX/XLA:TPU 程序在编译阶段抛出类似下面的报错时说明程序需要的静态分配总量超过了 TPU 芯片物理 HBMHigh Bandwidth Memory容量这是 XLA 错误分类中的 **E1000: Compile Time: HBM OOM**后端TPURESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.以上是 E1000 文档 中的示例错误消息。XLA 在编译时会检查所有必要静态分配的总和是否能放进设备 HBM。编译器管理的 TPU HBM 分配包括六类程序输入输出训练批次、优化器状态等、TPU temporaries激活、梯度等中间计算、编译产物TensorCore 和 SparseCore 的机器码、系统开销XLA Runtime 预留空间、常量内嵌在 HLO IR 中的常量、编译器内部分配如 mesh 中节点的路由信息。当这六类的总和放不进 HBM 时就会报出这个错。排查路径在 XLA 仓库文档中是明确的E1000 文档 负责按错误消息分诊Debug OOM errors with XProf 负责用 XProf Memory Viewer 定位峰值内存XLA flags guidance 提供最后的内存 flag 调优项。下面按这条路径展开。第一步先看错误消息属于哪种形态error_1000.md 要求仔细分析错误消息和日志然后进入对应分支错误明确给出了 TC/SC 用量分解形如TC Hbm usage: X, SC Hbm usage Y示例TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G→ 进入TC/SC 失衡分支错误是Ran out of memory in memory space hbm且日志中列出了异常大的分配单个张量超过 HBM 上限的 50%→ 进入大分配分支错误是Ran out of memory in memory space hbm但日志中没有异常大的张量 → 进入累积压力分支需要用 XProf 可视化峰值内存。分支一错误显示了 TC/SC 用量分解此时是 TensorCoreTC SparseCoreSC的总用量超过了 HBM 上限对比两个数值找出瓶颈SparseCore 用量高时文档给出的检查项HBM stack 用量随feature_width、max_unique_nz_per_row和logical_replica_count增长。可以用--xla_sc_num_serialized_tables_to_optimize_hbmflag 把 table 的处理串行化以降低峰值 stack 用量代价是并行度下降检查 padding 开销SparseCore 会把 embedding table 对齐到 32B8 个 float。feature width 较小的表例如 8 个 float会产生显著 padding 浪费maximum_parallel_iterations取值过大会把更多输入数据预取进 HBM heap调低该值可以释放内存确认 embedding table 是否在所有 chip 之间正确做了 mod sharding。TensorCore 用量高转到大分配分支分支二继续排查。两者都不高但总和超限说明已经到了芯片容量上限需要同时降低两个组件的用量按分支二、三的建议综合处理。分支二日志里有异常大的分配 50% HBM 上限E1000 文档明确指出出现这种大分配时几乎从来不是硬件容量问题通常是配置错误。具体检查查看大分配的 XLA label如果存在label 里通常有指向 JAX 源码位置的提示移除调试残留在大规模运行里使用jax.debug.print()会强制编译器把完整张量实体化到 HBM 再传回 CPU破坏融合并抬高峰值内存。删掉遗留的jax.debug.print()修正 mesh shape 或 sharding 标注错误的 mesh shape 或缺失的 sharding 标注会让编译器退化为replication——把非常大的张量整个塞进单块芯片。检查大分配的 shape确认 sharding 被正确指定并被 XLA 传播。分支三没有单一大分配——用 XProf 定位峰值内存当总分配超限但没有明显大张量时需要先看峰值时刻到底是谁占着 HBM。oom_debugging.md 给出的完整流程给程序加 profiling trace。文档示例文档示例中一个触发 OOM 的 JAX 程序这样写import jax from jax import random import jax.numpy as jnp jax.profiler.trace(/tmp/xprof) jax.jit def oom(): a random.normal(random.PRNGKey(1), (327680, 327680), dtypejnp.bfloat16) return a a if __name__ __main__: oom()在你的程序中把jax.profiler.trace装饰到需要捕获的入口函数上第一个参数是 profile 存储目录。文档特别建议用jax.profiler.trace而不是jax.profiler.start_trace/stop_trace因为前者是上下文管理器在异常情况下也能安全结束 profiling。安装并启动 XProf指定 profile 目录和端口pip install xprof xprof --logdir/tmp/xprof/ --port6006打开 Memory Viewer。本地机器上访问http://localhost:6006在Tools下拉框选择Memory Viewer在Memory Types下拉框选择HBM通常默认已选中看 HLO Ops at Peak Memory Allocation 区块。该区块展示峰值内存使用点的 buffer 图buffer 包括Program Inputs and Outputs训练批次、优化器状态等TensorCore and SparseCore Temporaries中间计算激活、梯度等所需的动态内存。鼠标悬停在 buffer 图上可以看到该 Op 的 size、shape、allocation type 等细节用来识别占用高或生命周期长的 temporaries以及 padding 低效的大输入/中间/输出张量。下面按 Memory Viewer 里看到的主力归因选择文档给出的优化项。配置层面的调整往往最先有效减小 batch size中间激活和梯度的内存与 batch size 成正比。注意减小 batch size 可能需要同步重调学习率、动量或优化器超参以维持训练稳定性捐赠输入 buffer如果某个输入在计算后不再使用且其 shape 和元素类型与某个输出匹配可以通过jax.jit的donate_argnums参数把该输入 buffer 捐给输出内存减少量约为被捐 buffer 的大小对最大张量启用 bfloat16 或量化如模型架构和质量要求允许。这会改变数值行为需要谨慎评估Micro-batching可选当无法减小全局 batch size 或增加芯片数、且单芯片 batch size 已接近下限时把每个 batch 拆成n个 micro-batch逐个跑前向和反向最后累积梯度并整体更新权重——激活内存从M降到约M/n。文档标注的代价step 时间变长多次前向反向且模型与 micro-batch 尺寸差距过大会带来收敛问题。架构与 sharding 层面当配置调整不够时可能是模型拓扑对当前硬件过大换用更新一代的 TPU单芯片 HBM 更大在更大的芯片拓扑上运行把权重 shard 到更多芯片使用更高级的 data/tensor/pipeline 并行并为中间值和输出指定 sharding hint。注意把张量切到多芯片会带来网络通信开销Host offloading把大张量激活、优化器状态卸载到 host CPU 内存。数值上安全但文档明确警告会严重影响性能——系统要不断在 TPU HBM 和 CPU RAM 之间搬大张量属于最后手段。检查 tensor padding 与对齐TPU 上的低效形状是 OOM 的常见且隐蔽的成因。为了达到峰值性能XLA 会把 tensor 维度做 padding——minor-most 维度通常对齐到 128 的倍数第二小的维度对齐到 8 的倍数。padding 影响输入数组和中间 tensorHLO temporaries在小维度上可能显著放大内存用量。在 XProf Memory Viewer 中悬停 buffer 查看详情卡里的 padding 信息文档以 TPU v5 默认 layout 为例shape(129, 1024)可能被 pad 到(256, 1024)产生近 50% 的内存浪费文档示例改成(128, 1024)则不需要 padding。把大 tensor 的维度batch size、embedding 维度、hidden size调整为 128 的倍数。这会改变模型行为需要谨慎评估。Rematerialization 与手动 checkpointing模型接近能放进内存时可以用jax.checkpoint装饰器配合jax.grad手动控制哪些中间值在前向保留、哪些在反向重算——用算力换 HBM。也可以让XLA::Rematerializationpass 优先省内存代价是编译变慢。E1000 文档列出的 flag 及取舍Flag作用影响 / 取舍--xla_tpu_max_hbm_size_mib手动设置 Rematerialization pass 使用的 HBM 上限强迫编译器把程序塞进比物理 HBM 更小的限制--xla_tpu_rematerialization_algoPEAK_PRIORITY把优化集中到内存峰值点相比默认算法可能更高效地削减内存--xla_tpu_rematerialization_max_block_size_limit32控制一次可 rematerialize 的 block 内最大指令数调大可以省更多内存但显著增加编译时间--xla_tpu_rematerialization_block_effort_factor10.0定义搜索可 rematerialize block 的编译努力程度值越大搜索越彻底编译时间越长--xla_tpu_pre_fusion_remattrue在 fusion pass 之前增加一次 Rematerialization pass能发现更多内存节省但编译时间增加且可能影响数值稳定性文档明确提示修改 XLA flag 应作为最后手段可能损害性能。最后一档XLA Memory Flagsflags_guidance.md 的 Memory Flags 一节 说明这些 flag 就是为了解决编译期的 HBM OOM 提供的只在遇到 HBM out of memory 时调整其余场景保持默认值改动可能损害性能。其中的默认值与建议值Flag默认值文档建议值xla_latency_hiding_scheduler_rerun15每次 rerun 会逐步下调调度内存上限文档标注超过 10 次意义不大xla_tpu_rwb_fusiontruefalse关闭 reducebroadcast 融合可降内存xla_memory_schedulerkDefaultkBrkga更省内存的调度算法代价是编译更慢xla_tpu_enable_latency_hiding_schedulertruefalse以放弃异步 collective 的性能收益换内存xla_jf_spmd_threshold_for_windowed_einsum_mib-110Mb~1Gb提高阈值可省内存代价是失去 collective matmul 机会如何对照症状选手段E1000 文档末尾给了一张干预手段速查表下表为按原文整理典型症状列是文档给出的 telltale signs帮助你确认当前瓶颈是否对得上手段是否改变程序行为典型症状对上再动手高级 sharding 技术基本不改变数值正确性但增加网络通信开销Memory Viewer 中单个张量远大于其他如被复制到所有 TPUTensorBoard hooks 中数组显示未分片减小 batch size改变训练动态通常要重调学习率micro-batching 是不改变行为的替代梯度计算时 Temporaries 分配失败Op 名含 JVP内存 profile 中大量 batch 形状张量启用混合精度bfloat16有风险改变数值精度可能影响结果或导致不收敛Memory Viewer 确认最大张量目前是float32手动 checkpointingjax.checkpoint不改变行为用计算换内存反向传播时大量完全同尺寸的张量占满内存常伴随 JVP Op 名捐赠输入 bufferdonate_argnums不改变实验完整性用错会直接报清晰的错误无特定信号属于白捡的赢面值得先试修改模型维度改变模型行为可能直接破坏与数据集的兼容性Memory Viewer 显示大量 padding 浪费维度不是 128 的倍数等Host offloading数值安全但性能上是大坑通常只作为超大优化器状态或重预处理步骤的最后手段修复后如何确认这个错误发生在编译期检查阶段所以验证方式就是重新运行程序修改生效后XLA:TPU 编译应当通过不再出现RESOURCE_EXHAUSTED: ... Ran out of memory in memory space hbm。另外注意上文中会改变程序行为的几个手段减小 batch size、bfloat16/量化、修改模型维度、xla_tpu_pre_fusion_remat在文档中都有可能影响模型行为/收敛/数值稳定性的标注——对这类改动编译通过只是第一步还需要按文档提示重新核对训练指标如重调学习率、观察收敛而纯 sharding、donate_argnums、jax.checkpoint这类不改变数值正确性的手段编译通过即可视为问题解除。参考文档E1000 - Compile Time: HBM OOM、Debug OOM errors with XProf、XLA flags guidance、XLA 错误总览。【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 TensorFlow XLA:TPU 编译报 “Ran out of memory in memory space hbm“ 怎么排查？
TensorFlow XLA:TPU 编译报 Ran out of memory in memory space hbm 怎么排查【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow当 JAX/XLA:TPU 程序在编译阶段抛出类似下面的报错时说明程序需要的静态分配总量超过了 TPU 芯片物理 HBMHigh Bandwidth Memory容量这是 XLA 错误分类中的 **E1000: Compile Time: HBM OOM**后端TPURESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.以上是 E1000 文档 中的示例错误消息。XLA 在编译时会检查所有必要静态分配的总和是否能放进设备 HBM。编译器管理的 TPU HBM 分配包括六类程序输入输出训练批次、优化器状态等、TPU temporaries激活、梯度等中间计算、编译产物TensorCore 和 SparseCore 的机器码、系统开销XLA Runtime 预留空间、常量内嵌在 HLO IR 中的常量、编译器内部分配如 mesh 中节点的路由信息。当这六类的总和放不进 HBM 时就会报出这个错。排查路径在 XLA 仓库文档中是明确的E1000 文档 负责按错误消息分诊Debug OOM errors with XProf 负责用 XProf Memory Viewer 定位峰值内存XLA flags guidance 提供最后的内存 flag 调优项。下面按这条路径展开。第一步先看错误消息属于哪种形态error_1000.md 要求仔细分析错误消息和日志然后进入对应分支错误明确给出了 TC/SC 用量分解形如TC Hbm usage: X, SC Hbm usage Y示例TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G→ 进入TC/SC 失衡分支错误是Ran out of memory in memory space hbm且日志中列出了异常大的分配单个张量超过 HBM 上限的 50%→ 进入大分配分支错误是Ran out of memory in memory space hbm但日志中没有异常大的张量 → 进入累积压力分支需要用 XProf 可视化峰值内存。分支一错误显示了 TC/SC 用量分解此时是 TensorCoreTC SparseCoreSC的总用量超过了 HBM 上限对比两个数值找出瓶颈SparseCore 用量高时文档给出的检查项HBM stack 用量随feature_width、max_unique_nz_per_row和logical_replica_count增长。可以用--xla_sc_num_serialized_tables_to_optimize_hbmflag 把 table 的处理串行化以降低峰值 stack 用量代价是并行度下降检查 padding 开销SparseCore 会把 embedding table 对齐到 32B8 个 float。feature width 较小的表例如 8 个 float会产生显著 padding 浪费maximum_parallel_iterations取值过大会把更多输入数据预取进 HBM heap调低该值可以释放内存确认 embedding table 是否在所有 chip 之间正确做了 mod sharding。TensorCore 用量高转到大分配分支分支二继续排查。两者都不高但总和超限说明已经到了芯片容量上限需要同时降低两个组件的用量按分支二、三的建议综合处理。分支二日志里有异常大的分配 50% HBM 上限E1000 文档明确指出出现这种大分配时几乎从来不是硬件容量问题通常是配置错误。具体检查查看大分配的 XLA label如果存在label 里通常有指向 JAX 源码位置的提示移除调试残留在大规模运行里使用jax.debug.print()会强制编译器把完整张量实体化到 HBM 再传回 CPU破坏融合并抬高峰值内存。删掉遗留的jax.debug.print()修正 mesh shape 或 sharding 标注错误的 mesh shape 或缺失的 sharding 标注会让编译器退化为replication——把非常大的张量整个塞进单块芯片。检查大分配的 shape确认 sharding 被正确指定并被 XLA 传播。分支三没有单一大分配——用 XProf 定位峰值内存当总分配超限但没有明显大张量时需要先看峰值时刻到底是谁占着 HBM。oom_debugging.md 给出的完整流程给程序加 profiling trace。文档示例文档示例中一个触发 OOM 的 JAX 程序这样写import jax from jax import random import jax.numpy as jnp jax.profiler.trace(/tmp/xprof) jax.jit def oom(): a random.normal(random.PRNGKey(1), (327680, 327680), dtypejnp.bfloat16) return a a if __name__ __main__: oom()在你的程序中把jax.profiler.trace装饰到需要捕获的入口函数上第一个参数是 profile 存储目录。文档特别建议用jax.profiler.trace而不是jax.profiler.start_trace/stop_trace因为前者是上下文管理器在异常情况下也能安全结束 profiling。安装并启动 XProf指定 profile 目录和端口pip install xprof xprof --logdir/tmp/xprof/ --port6006打开 Memory Viewer。本地机器上访问http://localhost:6006在Tools下拉框选择Memory Viewer在Memory Types下拉框选择HBM通常默认已选中看 HLO Ops at Peak Memory Allocation 区块。该区块展示峰值内存使用点的 buffer 图buffer 包括Program Inputs and Outputs训练批次、优化器状态等TensorCore and SparseCore Temporaries中间计算激活、梯度等所需的动态内存。鼠标悬停在 buffer 图上可以看到该 Op 的 size、shape、allocation type 等细节用来识别占用高或生命周期长的 temporaries以及 padding 低效的大输入/中间/输出张量。下面按 Memory Viewer 里看到的主力归因选择文档给出的优化项。配置层面的调整往往最先有效减小 batch size中间激活和梯度的内存与 batch size 成正比。注意减小 batch size 可能需要同步重调学习率、动量或优化器超参以维持训练稳定性捐赠输入 buffer如果某个输入在计算后不再使用且其 shape 和元素类型与某个输出匹配可以通过jax.jit的donate_argnums参数把该输入 buffer 捐给输出内存减少量约为被捐 buffer 的大小对最大张量启用 bfloat16 或量化如模型架构和质量要求允许。这会改变数值行为需要谨慎评估Micro-batching可选当无法减小全局 batch size 或增加芯片数、且单芯片 batch size 已接近下限时把每个 batch 拆成n个 micro-batch逐个跑前向和反向最后累积梯度并整体更新权重——激活内存从M降到约M/n。文档标注的代价step 时间变长多次前向反向且模型与 micro-batch 尺寸差距过大会带来收敛问题。架构与 sharding 层面当配置调整不够时可能是模型拓扑对当前硬件过大换用更新一代的 TPU单芯片 HBM 更大在更大的芯片拓扑上运行把权重 shard 到更多芯片使用更高级的 data/tensor/pipeline 并行并为中间值和输出指定 sharding hint。注意把张量切到多芯片会带来网络通信开销Host offloading把大张量激活、优化器状态卸载到 host CPU 内存。数值上安全但文档明确警告会严重影响性能——系统要不断在 TPU HBM 和 CPU RAM 之间搬大张量属于最后手段。检查 tensor padding 与对齐TPU 上的低效形状是 OOM 的常见且隐蔽的成因。为了达到峰值性能XLA 会把 tensor 维度做 padding——minor-most 维度通常对齐到 128 的倍数第二小的维度对齐到 8 的倍数。padding 影响输入数组和中间 tensorHLO temporaries在小维度上可能显著放大内存用量。在 XProf Memory Viewer 中悬停 buffer 查看详情卡里的 padding 信息文档以 TPU v5 默认 layout 为例shape(129, 1024)可能被 pad 到(256, 1024)产生近 50% 的内存浪费文档示例改成(128, 1024)则不需要 padding。把大 tensor 的维度batch size、embedding 维度、hidden size调整为 128 的倍数。这会改变模型行为需要谨慎评估。Rematerialization 与手动 checkpointing模型接近能放进内存时可以用jax.checkpoint装饰器配合jax.grad手动控制哪些中间值在前向保留、哪些在反向重算——用算力换 HBM。也可以让XLA::Rematerializationpass 优先省内存代价是编译变慢。E1000 文档列出的 flag 及取舍Flag作用影响 / 取舍--xla_tpu_max_hbm_size_mib手动设置 Rematerialization pass 使用的 HBM 上限强迫编译器把程序塞进比物理 HBM 更小的限制--xla_tpu_rematerialization_algoPEAK_PRIORITY把优化集中到内存峰值点相比默认算法可能更高效地削减内存--xla_tpu_rematerialization_max_block_size_limit32控制一次可 rematerialize 的 block 内最大指令数调大可以省更多内存但显著增加编译时间--xla_tpu_rematerialization_block_effort_factor10.0定义搜索可 rematerialize block 的编译努力程度值越大搜索越彻底编译时间越长--xla_tpu_pre_fusion_remattrue在 fusion pass 之前增加一次 Rematerialization pass能发现更多内存节省但编译时间增加且可能影响数值稳定性文档明确提示修改 XLA flag 应作为最后手段可能损害性能。最后一档XLA Memory Flagsflags_guidance.md 的 Memory Flags 一节 说明这些 flag 就是为了解决编译期的 HBM OOM 提供的只在遇到 HBM out of memory 时调整其余场景保持默认值改动可能损害性能。其中的默认值与建议值Flag默认值文档建议值xla_latency_hiding_scheduler_rerun15每次 rerun 会逐步下调调度内存上限文档标注超过 10 次意义不大xla_tpu_rwb_fusiontruefalse关闭 reducebroadcast 融合可降内存xla_memory_schedulerkDefaultkBrkga更省内存的调度算法代价是编译更慢xla_tpu_enable_latency_hiding_schedulertruefalse以放弃异步 collective 的性能收益换内存xla_jf_spmd_threshold_for_windowed_einsum_mib-110Mb~1Gb提高阈值可省内存代价是失去 collective matmul 机会如何对照症状选手段E1000 文档末尾给了一张干预手段速查表下表为按原文整理典型症状列是文档给出的 telltale signs帮助你确认当前瓶颈是否对得上手段是否改变程序行为典型症状对上再动手高级 sharding 技术基本不改变数值正确性但增加网络通信开销Memory Viewer 中单个张量远大于其他如被复制到所有 TPUTensorBoard hooks 中数组显示未分片减小 batch size改变训练动态通常要重调学习率micro-batching 是不改变行为的替代梯度计算时 Temporaries 分配失败Op 名含 JVP内存 profile 中大量 batch 形状张量启用混合精度bfloat16有风险改变数值精度可能影响结果或导致不收敛Memory Viewer 确认最大张量目前是float32手动 checkpointingjax.checkpoint不改变行为用计算换内存反向传播时大量完全同尺寸的张量占满内存常伴随 JVP Op 名捐赠输入 bufferdonate_argnums不改变实验完整性用错会直接报清晰的错误无特定信号属于白捡的赢面值得先试修改模型维度改变模型行为可能直接破坏与数据集的兼容性Memory Viewer 显示大量 padding 浪费维度不是 128 的倍数等Host offloading数值安全但性能上是大坑通常只作为超大优化器状态或重预处理步骤的最后手段修复后如何确认这个错误发生在编译期检查阶段所以验证方式就是重新运行程序修改生效后XLA:TPU 编译应当通过不再出现RESOURCE_EXHAUSTED: ... Ran out of memory in memory space hbm。另外注意上文中会改变程序行为的几个手段减小 batch size、bfloat16/量化、修改模型维度、xla_tpu_pre_fusion_remat在文档中都有可能影响模型行为/收敛/数值稳定性的标注——对这类改动编译通过只是第一步还需要按文档提示重新核对训练指标如重调学习率、观察收敛而纯 sharding、donate_argnums、jax.checkpoint这类不改变数值正确性的手段编译通过即可视为问题解除。参考文档E1000 - Compile Time: HBM OOM、Debug OOM errors with XProf、XLA flags guidance、XLA 错误总览。【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考