错误代码：E0100

类别：运行时：缓冲区分配失败

此错误表示 XLA:TPU 运行时的内存分配器无法在加速器的 HBM 上找到适合所请求分配的内存块。

示例出错提示：

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

XLA 后端：TPU

概览

在以下情况下会抛出此错误：

或

  • 用户安排的程序输出分配失败。

这些失败通常是由以下几个原因造成的：

  • 内存不足 (OOM)：用户尝试分配的内存块大于 TPU HBM 上可用的总空闲内存量。
  • 内存碎片：分配失败，因为内存空间中没有单个连续的空闲块足够大，无法满足请求的大小。可用内存总量足以进行分配，但这些内存分散在内存空间中，以小块不连续的块形式存在。

TPU 运行时具有多种机制来重试分配失败，包括：

  • 如果存在已排队的取消分配，运行时会重试失败的分配。
  • 对于由碎片化导致的 OOM，运行时可以自动触发碎片整理和重试。
  • TPU 运行时优先考虑缓冲区分配，而不是保持程序加载状态。如果由于 HBM 不足而导致缓冲区分配失败，系统将逐出已加载的 TPU 程序，直到有足够的内存可用于缓冲区为止。

因此，在采取上述缓解措施后遇到的错误通常需要用户采取行动。

调试

  • 减少模型的内存占用空间
    • 减小批次大小：减小批次大小可直接降低内存用量。
    • 参数分片：对于非常大的模型，请使用模型并行或分片等技术将参数分布到多个 TPU 核心或主机的 HBM 中。
    • 缩短序列/上下文长度：对于处理序列的模型（例如语言模型），缩短输入序列长度可以显著减少内存占用量。
    • 缓冲区捐赠：利用框架功能（例如：jax.jit(..., donate_argnums=...)）向 XLA 发出信号，表明某些输入缓冲区可以被覆盖并重新用于输出。如需了解详情，请参阅缓冲区捐赠。
    • 优化检查点策略：考虑仅保存模型权重或使用分片检查点策略，而不是一次性保存整个模型状态。
  • 寻址内存布局和填充
    • TPU 内存以块为单位进行分配，填充可能会增加张量的实际大小。
  • 确保没有内存泄漏
    • 确保对 jax.Array 对象的引用不会保留过长时间。保留 jax.Array 对象可能会阻止自动取消分配，即使在程序编译完成后也是如此。

另请参阅错误代码：E1000，了解可用于减少每个程序使用的内存量的其他策略。

工具

启用 tpu_log_allocations_on_oom 标志后，当发生内存不足 (OOM) 时，分配器将转储所有当前分配的详细报告，这对于调试来说非常宝贵。