카테고리: 컴파일 시간: 범위 지정된 Vmem OOM
이 오류는 프로그램에 할당된 것보다 더 많은 범위 지정된 벡터 메모리 (Vmem)가 필요함을 나타냅니다.
샘플 오류 메시지:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
XLA 백엔드: TPU
개요
TPU에는 TensorCore (TC)에서만 사용하는 로컬 스크래치패드 메모리인 벡터 메모리 (VMEM)가 있습니다. 컴파일러는 다양한 유형의 할당에 대해 Vmem을 관리합니다.
- 명령어 범위 지정 할당: 단일 HLO 명령어를 실행하는 동안 Vmem의 임시 저장소입니다. 여기에는 피연산자 범위 버퍼 (예: 이중 버퍼링) 및 레지스터 스필이 포함됩니다.
- 프로그램 범위 지정 할당: 단일 HLO 명령어의 범위를 벗어나는 할당입니다. 일반적으로 HLO 명령어의 입력 또는 출력인 HLO 임시 및 중간 결과입니다.
컴파일 시간 범위 지정된 Vmem OOM은 명령어 범위 지정 할당이 해당 명령어의 할당 한도를 초과할 때 발생합니다. 이 한도는
- --xla_tpu_scoped_vmem_limit_kib 플래그를 통해 전체 프로그램에 대해 전역적으로
및
- vmem_limit_bytes 매개변수를 통해 커스텀 커널별로
이러한 오류는 일반적으로 내부 컴파일러 버그 또는 할당 한도를 초과하는 커스텀 커널로 인해 발생합니다.
디버깅
오류 메시지를 주의 깊게 분석하여 오류가 커스텀 커널에서 비롯되는지 아니면 표준 HLO에서 비롯되는지 확인합니다. 커스텀 커널로 인한 오류에는 다음과 같은 서명이 있어야 합니다.
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- 커스텀 커널 범위 지정된 Vmem OOM: 오류가 커스텀 커널을 가리키는 경우 → 커널 재조정으로 이동합니다.
- 커널이 아닌 Vmem 문제: 커스텀 커널이 아닌 작업으로 인해 Vmem OOM이 발생하는 경우 내부 컴파일러 버그일 가능성이 높습니다. HLO 덤프를 사용하여 XLA에 버그 신고를 제출해 주세요.
커널 재조정
오류가 커스텀 커널에서 비롯되는 경우 다음 기법을 사용하여 커널의 메모리 요구사항을 줄입니다.
- 블록 크기 조정: 커널 구성에서 블록 크기 (타일 크기)를 줄여 범위 지정된 Vmem 사용량을 줄입니다.
- 커널별 범위 지정된 Vmem 한도 설정: vmem_limit_bytes 매개변수를 사용하여 특정 커널에 필요한 메모리 양을 명시적으로 요청합니다.
- 메모리 색상 지정 수정: pallas.tpu.with_memory_space_constraint를 사용하여 커널의 입력/출력을 Vmem에 명시적으로 색상 지정/제한합니다. pallas.tpu.with_memory_space_constraint 전체 Vmem OOM이 발생할 수 있으므로 Vmem에 너무 많은 입력 출력을 색상 지정하지 않도록 주의하세요.
- Vmem 한도 조정: 커널별 재조정이 어렵거나 문제가 여러 커널에 영향을 미치는 경우 전역 Vmem 한도를 조정할 수 있습니다.--xla_tpu_scoped_vmem_limit_kib 플래그를 사용하여