Improved Scaling for Fast Mode of Ozaki Scheme II

기술 요약: Ozaki Scheme II의 Fast Mode를 위한 개선된 스케일링

문제 정의
Ozaki Scheme II는 중국인 나머지 정리(CRT)를 기반으로 저정밀도 정수 하드웨어(예: INT8 Tensor Cores)를 사용하여 고정밀도 행렬 곱셈(예: FP64/FP32)을 에뮬레이션하는 방법이다. 이 방식은 고정밀도 입력 행렬을 정수 행렬로 스케일링할 때, 이 정수들의 곱이 CRT를 통해 재구성될 때 특정 임계값($P/2$)을 초 exceed하지 않도록 하는 과정이 필요하다. 만약 이 임계값을 초과하면 CRT 복구가 실패하여 치명적인 정확도 손실이 발생한다.

이 방식은 스케일링 계수($\mu$와 $\nu$)를 결정하기 위해 두 가지 모드를 제공한다:

1. Accurate Mode (정밀 모드): 보조 INT8 행렬 곱셈(GEMM)을 사용하여 스케일링 계수를 추정하며, 높은 정확도를 보장하지만 상당한 계산 오버헤드가 발생한다.
2. Fast Mode (고속 모드): 보조 GEMM 없이 코시-슈바르츠 부등식(Cauchy–Schwarz inequality)을 사용하여 스케일링 계수를 추정하며, 더 높은 처리량을 제공하지만 역사적으로 정확도가 낮았다.

본 논문은 기존 Fast Mode 스케일링 공식의 결정적인 결함을 식별하였다: 즉, **스케일 불변성(scale invariance)**이 결여되어 있다는 점이다. 입력 행렬에 스칼라 상수 $\alpha$를 곱할 경우, 결과로 나오는 스케일링된 정수 행렬의 유효 비트 폭(effective bit width)이 예측 불가능하게 변한다.

• 큰 $\alpha$에 대해서는 유효 비트 폭이 감소하여 정확도 저하를 일으킨다.
• 작은 $\alpha$에 대해서는 스케일링된 정수 곱이 CRT 임계값($P/2$)을 초과하여 완전한 복구 실패를 초래할 수 있다.
  모든 원소가 1인 행렬 및 무작위 행렬에 대한 수치 실험 결과, 특정 스칼라 범위에서 Fast Mode는 올바른 결과를 복구하는 데 실패한 반면, Accurate Mode는 안정성을 유지함을 입示하였다.

방법론
저자들은 코시-슈바르츠 부등식을 통해 CRT 유일성 조건으로부터 직접 유도된, Fast Mode를 위한 수정된 스케일링 공식을 제안한다.

1. 이론적 유도:
   저자들은 스케일링 계수 $\mu_i$와 $\nu_j$에 대한 조건을 재정의한다. 기존의 휴리스틱 대신, 다음의 충분 조건을 정의한다:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   GPU 수학 API에서 제공하는 log2f 함수의 절대 오차 범위를 고려하고, base-2 로그를 사용하여 $\mu_i$와 $\nu_j$를 구함으로써 새로운 공식을 도출한다:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   여기서 $P'_{prop}$는 사전 계산된 상수이며, $\widehat{\log_2}$는 유계된 근사 로그 함수이다.

2. 스케일 불변성 증명:
   저자들은 새로운 공식이 스케일 불변성을 가짐을 이론적으로 입증한다. 입력 행렬 $A$가 $\alpha = 2^s$에 의해 스케일링되면, $\log_2(\alpha)$ 항이 스케일링된 정수 행렬의 유효 비트 폭 계산 과정에서 상쇄된다. 결과적으로, 입력 크기와 관계없이 비트 폭이 일정하게 유지되므로, 기존 Fast Mode에서 관찰되었던 정확도 저하 및 CRT 실패를 방지한다.

3. 구현:
   제안된 방법(OS II-prop)은 GEMMul8 라이브러리의 Fast Mode에 있는 스케일링 공식을 대체한다. 이는 기존 Fast Mode의 $O(mk)$및$O(kn)$ 계산 복잡도를 유지하며, 추가적인 오버헤드 없이 기존 Fast Mode와 동일한 속도를 유지하면서도 Accurate Mode(추가 GEMM 필요)보다는 빠르다.

주요 기여

• 한계점 식별: 본 논문은 기존 Fast Mode 스케일링 공식이 스케일 불변성이 부족하여, 특정 입력 스케일링 상황에서 정확도 저하 및 CRT 복구 실패를 초래함을 입증한다.
• 수정된 공식: 스케일 불변성을 보장하고 CRT 유일성 조건을 항상 만족하는 새로운 스케일링 공식을 제안한다.
• 제로 오버헤드: 새로운 공식은 기존 Fast Mode 대비 추가적인 계산 비용을 발생시키지 않는다.
• 포괄적 평가: 제안된 방법은 NVIDIA GH200 GPU에서 이중 정밀도(DGEMM) 및 단정밀도(SGEMM) 행렬 곱셈에 대해 평가되었다.

결과
실험은 다양한 원소 크기 편차($\phi$)를 가진 무작위 행렬과 행렬 차원($k \in \{1024, 16384\}$)을 사용하여 수행되었다.

• 정accuracy: 제안된 방법(OS II-prop)은 모든 $\phi$ 값에 대해 Accurate Mode(OS II-accu)와 대등한 정확도를 달달성한다. 반면, 기존 Fast Mode(OS II-fast)는 큰 $\phi$ 값(조건수가 나쁜 경우)에서 심각한 정확도 저하와 실패를 보였다.
• 처리량(Throughput): OS II-prop는 기존 Fast Mode와 대등한 처리량을 유지한다. OS II-prop는 모듈러스 개수가 적은 특정 시나리오(스케일링된 정수의 zero 요소가 적음)에서 기존 Fast Mode보다 약간 느릴 수 있으나, Accurate Mode보다는 훨씬 높은 처리량을 보인다.
• 트레이드오프: 정확도-처리량 트레이드오프 분석에서, OS II-prop는 OS II-accu를 지배하며(동일 정확도에서 더 높은 처리량), OS II-fast의 정확도 한계를 극복한다. 이는 광범위한 조건에서 cuBLAS보다 정확도와 처리량 모두에서 우수한 성능을 동시에 낼 수 있는 유일한 방법이다.

의의
본 논문은 제안된 방법이 성능상의 이점을 희생하지 않으면서 Ozaki Scheme II의 Fast Mode가 가진 근본적인 정확도 한계를 해결한다고 주장한다. 스케일 불변성을 보장함으로써, 이 방법은 입력 스케일링에 관계없이 CRT 복구가 성공하도록 보장하며, 더 나은 정확도와 성능의 균형을 제공한다. 이를 통해 사용자는 Fast Mode의 속도로 고정밀도 결과를 얻으면서도 Accurate Mode의 신뢰성을 확보할 수 있으며, 결과적으로 Accurate Mode의 처리량 제약과 기존 Fast Mode의 정확도 제약을 효과적으로 극복할 수 있다.

저자들은 현재의 연구가 DGEMM 및 SGEMM을 위한 INT8 행렬 엔진에 집중되어 있지만, 이 접근 방식이 FP8 엔진 및 복소수 행렬 곱셈(ZGEMM/CGEMM)에도 적용 가능함을 언급하며, 이는 향후 연구 과제로 남겨두었다.