错误代码：0100

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

此错误表示 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

概览

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

  • 通过 jax.device_put 或
  • 用户调度程序输出分配的失败。

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

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

TPU 运行时内置了多种机制来重试分配失败，包括：

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

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

调试

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

工具

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