카테고리: 컴파일 시간: 하드웨어에서 지원되지 않는 RHS 데이터 유형
이 오류는 행렬 곱셈 (예: jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul 또는 @ 연산자)에서 오른쪽 (RHS) 피연산자에 사용되는 데이터 유형이 사용 중인 특정 TPU 생성에서 기본적으로 지원되지 않는 경우에 발생합니다.
샘플 오류 메시지:
INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...
The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>
XLA 백엔드: TPU
개요
TPU의 MXU (행렬 곱셈 단위)는 모든 하드웨어 생성에서 Float32 작업을 기본적으로 지원합니다.
하지만 BFloat16 및 기타 양자화된 데이터 유형(예: Int4, Int8 또는 Float8)에 대한 기본 지원은 하드웨어 생성에 따라 다릅니다. 이 오류는 커널이 특정 TPU 생성에서 실행할 물리적 회로가 없는 데이터 유형을 사용하여 행렬 곱셈을 MXU에 매핑하려고 할 때 트리거됩니다.
이 오류는 일반적으로 지원되지 않는 유형을 지원되는 유형으로 자동 변환하려고 시도하는 컴파일러의 정규화 패스(예: 소프트웨어 에뮬레이션을 통해)가 유효한 변환 규칙을 찾지 못했거나 호환성 모드 가 사용 중지되어 변환 규칙을 찾지 못했음을 나타냅니다.
디버깅
이 오류를 해결하려면 데이터 유형을 하드웨어의 기능에 맞춰야 합니다. 다음과 같은 옵션이 있습니다.
1. 기본 유형으로 캐스팅
가장 안정적인 해결 방법은 matmul 작업 전에 커널 내에서 피연산자를 하드웨어 지원 데이터 유형 (예: TPU v4+의 Float32 또는 BFloat16)으로 수동으로 캐스팅하는 것입니다.
- 이유:
Float32는 모든 TPU 생성에서 MXU가 기본적으로 지원하는 범용 데이터 유형입니다. - 절충: 캐스팅을 수행하는 데 필요한 주기인 VPU (벡터 처리 장치) 비용이 발생하지만 커널이 현재 하드웨어에서 실행되도록 보장합니다.
2. 호환성 모드 확인
일반적으로 컴파일러는 기본적으로 사용 설정된 호환성 모드 에서 이러한 유형 불일치 문제를 자동으로 처리할 수 있습니다. XLA 구성을 다시 한번 확인하여 --xla_mosaic_compat_mode가 false로 설정되지 않았는지 확인합니다.
이는 하드웨어에서 기본적으로 지원하지 않는 작업에 소프트웨어 에뮬레이션 시퀀스를 삽입하는 'polyfill' 역할을 합니다.
호환성 모드에서 사용 설정되는 항목:
- 혼합 정밀도 MatMul: 캐스팅 작업 (예: matmul 전에 정수를
Float32로 확장)을 자동으로 삽입하여 정수 피연산자를 부동 소수점 누산기와 혼합할 수 있습니다. - 저정밀도 에뮬레이션: 특정 하드웨어 생성에서 실행 전에
BFloat16또는Float32와 같은 지원되는 유형으로 확장하여4-bit부동 소수점 (4E2M1FN) 또는8-bit부동 소수점 (8E4M3FN)과 같은 지원되지 않는 유형을 에뮬레이션합니다.
이 모드는 에뮬레이션에서 MXU가 작동하기 전에 데이터 형식을 변환하기 위한 추가 명령어가 필요하므로 최대 성능보다 호환성을 우선시합니다.
3. 하드웨어 업그레이드 또는 지원 요청
알고리즘에 캐스팅 또는 에뮬레이션 오버헤드 없이 Int4 또는 Float8과 같은 유형에 기본 성능이 엄격하게 필요한 경우 기본 지원이 있는 최신 TPU 생성에서 실행해야 합니다.
기능 요청: 하드웨어에서 이 작업을 지원한다고 생각하거나 호환성 모드에서도 컴파일러에 유효한 에뮬레이션 경로가 없는 경우 기능 요청을 제출하세요. 일반적으로 작업이 이전 버전과 호환되도록 보장합니다. 따라서 커널이 TPU 생성에서 실행되면 모든 향후 생성에서 실행되지만 이전 생성에 대한 에뮬레이션이 보장되지는 않습니다 (일부의 경우 캐스팅 비용이 매우 많이 듭니다).