카테고리: 컴파일 시간: 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 또는 모자이크 커널을 식별하려면 중간 컴파일러 표현식을 생성하세요.
- Dump SparseCore MLIR:
--xla_sc_dump_mlir_to=/path/to/dump플래그를 설정합니다. 이렇게 하면 SparseCore 프로그램의 MLIR이 생성되므로 오류 메시지와 일치하는 할당 크기를 확인할 수 있습니다. - 덤프 모자이크 LLO: 맞춤 커널의 경우
--xla_mosaic_dump_to=/path/to/dump를 사용하여 모자이크에서 내보낸 모든 하위 수준 옵티마이저(LLO) 프로그램을 검사합니다.
- Dump SparseCore MLIR:
- 스크래치 크기 줄이기 (Pallas 사용자): 모자이크 커널 내에서 오류가 발생하는 경우
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
- 버그 신고: 위 단계로 문제가 해결되지 않으면 컴파일러 버그일 수 있습니다. 버그 신고를 제출해 주세요.