카테고리: 런타임: 프로그램 입력 버퍼 불일치
이 오류는 XLA 런타임이 컴파일된 프로그램에서 예상하는 메모리 버퍼 크기와 실행 시간에 실제로 제공되는 버퍼 크기 간의 불일치를 감지할 때 발생합니다.
샘플 오류 메시지:
XlaRuntimeError: INVALID_ARGUMENT: Executable(jit_embedding_pipeline_step_fn) expected parameter 2482 of size 5242880 (bf16[16,1280,40]{2,1,0:T(8,128)(2,1)}) but got buffer with incompatible size 1638400 (bf16[16,1280,40]{1,2,0:T(8,128)(2,1)}): while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
XLA 백엔드: TPU
개요
오류 메시지에는 예상 크기와 실제 크기, 텐서 모양과 레이아웃이 모두 표시됩니다. 두 텐서의 모양이 동일하더라도 물리적 레이아웃 (하드웨어에서 데이터가 타일링되고 정렬되는 방식)이 다르면 메모리 크기가 다를 수 있으므로 이러한 오류가 발생할 수 있습니다.
이러한 오류는 주로 다음으로 인해 발생합니다.
- 체크포인트와 XLA 구성 불일치 - 모델이 학습되고 체크포인트가 저장됩니다. 해당 체크포인트의 가중치의 물리적 레이아웃은 당시의 정확한 XLA 버전과 구성 (예: XLA 플래그)에 따라 결정됩니다. 나중에 이 체크포인트는 구성이 변경된 다른 환경에서 로드됩니다. 새 플래그, 다른 기본값 또는 모델/XLA 코드의 변경으로 인해 런타임에서 가중치의 다른 실제 레이아웃을 예상할 수 있습니다. 체크포인트의 이전 버퍼가 새로 컴파일된 XLA 프로그램에 전달되면 런타임에서 오류가 발생합니다.
- 하드웨어/토폴로지별 레이아웃 - XLA 컴파일러는 텐서의 다양한 물리적 레이아웃을 선택하여 다양한 하드웨어에서 성능을 최적화할 수 있습니다. v4 TPU에 최적인 레이아웃은 v5 TPU와 다를 수 있으며, 동일한 칩의 포드 슬라이스 (예: 4x8). 이 오류는 모델이 하나의 토폴로지 레이아웃에 관한 가정으로 컴파일되었지만 런타임에 다른 토폴로지에서 예약되거나 특정 하드웨어의 컴파일러 레이아웃 로직에 버그가 있는 경우 발생합니다.
디버깅
- 체크포인트에서 모델 내보내기와 재실행 간의 구성 일관성을 보장합니다.
- 레이아웃에 영향을 미치는 변경사항이 없다고 확신하지 않는 한 새 코드에 이전 체크포인트를 사용하지 마세요.
- 저장된 모델 다시 내보내기: 체크포인트/구성 불일치가 의심되는 경우 가장 확실한 해결책은 추론 또는 미세 조정에 사용 중인 것과 정확히 동일한 (최신) 코드베이스와 구성을 사용하여 저장된 모델을 다시 내보내는 것입니다.
- 두 실행 간의 구성 변경사항 (예: XLA 플래그)을 확인합니다.
- 하드웨어/토폴로지별 레이아웃:
- 하드웨어나 토폴로지를 전환하는 경우 하드웨어 버전과 토폴로지가 일치하지 않는지 확인합니다.