类别:编译时:SparseCore 分配失败
当 XLA:SparseCore 编译器无法在当前 SparseCore 程序所需的指定内存空间中分配连续的内存块时,会发生此错误。
示例错误消息:
INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>
XLA 后端:TPU
概览
SparseCore (SC) 是一种专门用于稀疏工作负载的处理器。它依赖于特定的内存层次结构来高效管理数据移动。XLA 编译器会尝试根据硬件限制和用户定义的形状静态调整缓冲区大小并分配缓冲区。此错误表示在此分配阶段出现了内存不足 (OOM) 情况。
错误消息通常会指定内存空间 ID。以下是常见的内存空间及其整数编码:
| 内存空间 ID | 名称 | 说明 |
|---|---|---|
| 0 | Smem | 本地标量内存。用于标量寄存器和控制流。 |
| 201 | TileSpmem | 特定于 Tile 的暂存内存。可供特定 SC tile 使用的快速本地 SRAM。 |
| 202 | Spmem | 共享暂存内存。用于机会性地暂存数据(输入、输出、中间数据),以隐藏 HBM 延迟。 |
| 203 | HBM | 高带宽内存。用于嵌入表、堆和堆栈的大型共享内存。 |
| 204 | 同步标志 | 用于协调的同步原语。 |
如需深入了解 SC 及其内存层次结构,请参阅 SparseCore 文档。
调试
解决方案取决于哪个内存空间分配失败。
场景 1。HBM 分配失败(内存空间 ID:203)
如果 SparseCore 程序请求的单个临时分配过大,无法放入可用的 HBM 中,则会发生此错误。在标准嵌入工作负载和 SC 分流集合中,过大的单核心分区或不正确的分片规范可能会迫使编译器请求大量缓冲区。
建议采取的措施:
- 检查分片:确保您的嵌入表和 SC 输入/输出张量已正确分区/分片。如果单个核心负责的数据过多,分配可能会失败。
- 调整限额:查看
max_ids_per_partition和max_unique_ids_per_partition。如果这些值设置得过高,编译器会预留比实际需要的更多内存。请参阅限制如何转化为表格。
情景 2。内部内存故障(内存空间 ID:0、201、202、204)
Smem、TileSpmem、Spmem 或同步标志中的分配失败通常是由于编译器 bug 或分配策略的限制所致,其中编译器未能考虑所有内存要求。
建议采取的措施:
- 隔离失败的 XLA 操作:为了确定导致失败的特定 SC HLO 或 Mosaic 内核,请生成中间编译器表示形式:
- 转储 SparseCore MLIR:设置标志
--xla_sc_dump_mlir_to=/path/to/dump。这会生成 SparseCore 程序的 MLIR,让您了解哪个分配大小与错误消息相符。 - 转储 Mosaic LLO:对于自定义内核,请使用
--xla_mosaic_dump_to=/path/to/dump检查 Mosaic 发出的所有低级优化器 (LLO) 程序。
- 转储 SparseCore MLIR:设置标志
- 减小暂存区大小(Pallas 用户):如果故障发生在 Mosaic 内核中,请检查您的
scratch_shapes配置。确保pltpu.SMEM请求符合特定 TPU 代系的硬件规范。 - 停用集合分流:如果错误源于 SC 分流的集合操作,请尝试停用 SC 分流功能:
--xla_tpu_enable_sparse_core_collective_offload_all_gather=false--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
- 提交 bug:如果上述步骤无法解决问题,则可能是编译器 bug。请提交 bug 报告。