카테고리: 런타임: 프로그램 할당 실패
이 오류는 TPU 기기의 XLA 런타임이 컴파일된 XLA 프로그램 실행 파일을 TPU의 HBM에 로드하지 못했음을 나타냅니다.
샘플 오류 메시지:
XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).
XLA 백엔드: TPU
개요
이 오류는 일반적으로 다음 이유 중 하나로 인해 발생합니다.
- 프로그램 크기가 사용 가능한 HBM을 초과함: 컴파일된 XLA 프로그램(명령어, 정적 데이터, 삽입된 상수 포함)이 프로그램이 로드되는 특정 TPU 코어에서 현재 사용 가능한 총 여유 HBM보다 큽니다.
- HBM 조각화: 기기의 총 여유 HBM이 합계로는 충분할 수 있지만 전체 프로그램을 맞출 수 있을 만큼 큰 단일 연속 블록으로 사용할 수 없습니다.
TPU 런타임이 메모리에 우선순위를 지정하는 방식을 이해하는 것이 중요합니다. 버퍼 할당은 로드된 프로그램보다 우선합니다. 버퍼 할당이 실패하면 런타임은 이미 로드된 프로그램을 HBM에서 삭제하여 공간을 확보합니다. 이로 인해 이전에 성공적으로 로드된 프로그램이 이제 HBM이 더 많은 데이터 버퍼로 점유되어 OOM 오류로 실패하는 상황이 발생할 수 있습니다.
디버깅
- 버퍼 메모리 사용 공간 줄이기: 데이터 버퍼에서 사용하는 메모리를 확보하면 프로그램 자체에 더 많은 공간이 남습니다.
- 배치 크기 줄이기: 활성화에 사용되는 메모리 양을 줄이는 가장 효과적인 방법 중 하나입니다.
- 파라미터 샤딩: 매우 큰 모델의 경우 모델 병렬 처리 또는 샤딩 기법 (예: FSDP 또는 Megascale)을 사용하여 모델의 파라미터와 계산을 여러 TPU 코어 또는 호스트에 분산합니다.
- 시퀀스/컨텍스트 길이 단축: 순차적 데이터를 처리하는 모델(예: NLP 모델)의 경우 시퀀스 길이를 줄이면 메모리 사용량을 크게 줄일 수 있습니다.
- 버퍼 기부: 프레임워크 기능 (예:
jax.jit(..., donate_argnums=...))를 사용하여 XLA가 출력을 저장하기 위해 입력 버퍼의 메모리를 재사용하도록 하여 최대 메모리 사용량을 줄입니다.
- 임시 파일에 대한 프로그램의 메모리 요구사항을 줄입니다.
tpu_shared_memory_percent플래그를 사용하여 임시 파일의 프로그램 메모리 사용량을 줄입니다. 이로 인해 성능에 부정적인 영향을 줄 수 있습니다.
- 실행 전략 최적화/제공 부하 감소:
- 프로그램 로드 관리: 여러 함수를 JIT 컴파일하는 경우 각 함수로 인해 프로그램이 로드될 수 있습니다. 동시에 로드되는 프로그램 수를 최소화하도록 워크로드를 구조화하세요.
- 메모리 누수가 없는지 확인합니다.
jax.Array객체에 대한 참조가 의도한 것보다 오래 유지되지 않도록 합니다.jax.Array객체를 보유하면 프로그램 컴파일이 완료된 후에도 자동 할당 해제가 방지될 수 있습니다.
도구
- 할당자가 OOM이 발생할 때 모든 현재 할당에 관한 자세한 보고서를 덤프하는
tpu_log_allocations_on_oom플래그를 사용 설정합니다. 이는 디버깅에 매우 유용할 수 있습니다. - 프로그램 프로파일링: JAX 메모리 프로파일러 또는 TensorFlow 프로파일러를 사용하여 시간 경과에 따른 프로그램의 메모리 사용량을 자세히 확인합니다. 이를 통해 메모리 소비의 예기치 않은 급증을 식별할 수 있습니다.