카테고리: 컴파일 시간: 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 | 타일별 스크래치패드 메모리입니다. 특정 SC 타일에서 사용할 수 있는 빠른 로컬 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 또는 동기화 플래그 의 할당 실패는 일반적으로 컴파일러가 모든 메모리 요구사항을 고려하지 못하는 컴파일러 버그 또는 할당 전략의 제한사항으로 인해 발생합니다.
권장 조치:
- 실패한 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
- 버그 신고: 위 단계를 수행해도 문제가 해결되지 않으면 컴파일러 버그일 가능성이 높습니다. 버그 신고를 제출해 주세요.