카테고리: 런타임: 버퍼 할당 실패
이 오류는 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객체를 보유하면 프로그램 컴파일이 완료된 후에도 자동 할당 해제가 방지될 수 있습니다.
도구
할당자가 OOM이 발생할 때 모든 현재 할당의 세부 보고서를 덤프하는 tpu_log_allocations_on_oom 플래그를 사용 설정합니다. 이는 디버깅에 매우 유용할 수 있습니다.