카테고리: 런타임: 프로그램 입력 버퍼 불일치
이 오류는 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와 다를 수 있으며, 동일한 칩의 포드 슬라이스 (예: 4x4x4와 4x8)에서도 다를 수 있습니다. 이 오류는 모델이 하나의 토폴로지 레이아웃에 관한 가정을 사용하여 컴파일되었지만 런타임에 다른 토폴로지에서 예약되거나 특정 하드웨어의 컴파일러 레이아웃 로직에 버그가 있는 경우 발생합니다.
디버깅
- 체크포인트에서 모델 내보내기 및 재실행 간 구성 일관성 보장
- 레이아웃에 영향을 미치는 변경사항이 없다고 확신하지 않는 한 새 코드와 함께 이전 체크포인트를 사용하지 마세요.
- 체크포인트/구성 불일치가 의심되는 경우 가장 확실한 해결책은 추론 또는 미세 조정에 사용하는 것과 정확히 동일한 (현재) 코드베이스와 구성을 사용하여 저장된 모델을 다시 내보내는 것입니다.
- 두 실행 간의 구성 변경사항 (예: XLA 플래그)을 확인합니다.
- 하드웨어/토폴로지별 레이아웃
- 하드웨어나 토폴로지를 전환하는 경우 하드웨어 버전과 토폴로지가 일치하지 않는지 확인합니다.