기술 요약: DASH – 배치 블록 프리컨디셔닝 및 효율적인 역제곱근 솔버를 통한 더 빠른 Shampoo

문제 정의

Shampoo는 낮은 활성화 이상치(activation outliers)와 개선된 압축성을 생성하는 것으로 알려진 선도적인 근사 2차 옵티마이저(approximate second-order optimizer)입니다. 그러나 이의 광범위한 채택은 주로 프리컨디셔너 행렬의 역행렬 제곱근($A^{-1/4}$)을 계산하는 데 드는 막대한 계산 오버헤드로 인해 저해되고 있습니다. Distributed Shampoo는 복잡도를 $O(n^3)$에서 $O(Bn^2)$로 줄이기 위해 프리컨디셔너를 $B \times B$ 블록으로 분할하여 이를 완화하려고 시도하지만, 기존 구현체들은 다음과 같은 결정적인 효율성 격차를 겪고 있습니다:

1. 순차적 처리: 표준 구현은 블록을 순차적으로 처리하여 현대적인 GPU 병렬성을 충분히 활용하지 못합니다.
2. 비효율적인 솔버: 기본값인 고유값 분해(EVD)에 대한 의존도는 GPU에서의 병렬화가 어렵습니다. Coupled-Newton(CN)과 같은 대안적인 반복법은 수치적 안정성과 수렴 속도에 대한 우려로 인해 기본적으로 비활성화되는 경우가 많습니다.
3. 최적화되지 않은 스케일링: 현재 방식들은 프로베니우스 노름(Frobenius norm)을 사용하여 행렬을 스케일링하는데, 이는 계산 비용은 저렴하지만 고유값을 0에 가깝게 밀어낼 수 있어 목표 정밀도에 도달하기 위해 훨씬 더 많은 반복 횟수를 요구하게 됩니다.

방법론

본 논문은 이론적 효율성과 하드웨어 역량 사이의 간극을 메우기 위해 설계된 고성능 구현체인 **DASH (Distributed Accelerated SHampoo)**를 제안합니다. 이 방법론은 두 가지 주요 기둥에 기반합니다:

1. 배치 블록 프리컨디셔닝 (시스템 레벨 최적화)

DASH는 프리컨디셔너 블록을 순차적으로 처리하는 대신, 독립적인 블록들을 3D 텐서로 쌓습니다.

• 3D 텐서 스태킹: 그래디언트 블록 $G_{ij}$를 $(N, B, B)$ 형상의 행렬으로 쌓으며, 여기서 $N$은 총 블록 수입니다.
• 병렬 실행: 이를 통해 옵티마이저가 배치 행렬 곱셈(bmm) 프리미티브를 사용하여 전체 블록 배치에 대해 역제곱근 절차(EVD, CN 또는 새로운 반복법)를 동시에 적용할 수 있게 합니다.
• GPU 활용도: 이러한 아키텍처 변경은 GPU 텐서 코어의 처리량을 극대화하며, 메모리 파편화를 줄이고 특정 솔버를 위한 반정밀도(FP16) 산술의 효율적인 사용을 가능하게 합니다.

2. 효율적인 역제곱근 솔버 (알고리즘 최적화)

본 논문은 표준 EVD 및 CN을 대체하거나 보완할 수 있는 고급 선형 대수적 접근 방식을 도입하고 분석합니다:

• Newton-Denman-Beavers (NDB) 반복법: 제곱근과 역제곱근을 모두 계산하는 새로운 Shampoo 방식입니다. 저자들은 중복된 행렬 곱셈을 제거하기 위해 첫 번째 반복 단계를 최적화하여, 오버헤드를 3번의 matmul(실제로는 이보다 적음)로 줄였습니다.
• 체비쇼프 다항식 근사 (Chebyshev Polynomial Approximations): 행렬 거듭제곱을 근사하기 위해 Clenshaw 알고리즘을 활용하는 대안적인 반복 솔버를 제공합니다.
• 스케일링을 위한 Multi-Power-Iteration: 수렴 문제를 해결하기 위해, 표준 프로베니우스 노름 스케일링을 Multi-Power-Iteration 방식으로 대체합니다.
  • 메커니즘: Power-Iteration을 통해 병렬로 처리되는 시작 벡터 풀(최대 32개)을 사용하여 스펙트럼 반지름($\lambda_{max}$)을 추정합니다.
  • 이점: 이는 프로베니우스 노로보다 더 정확한 스케일링 인자를 제공하여, 고유값이 반복 솔버가 빠르게 수렴할 수 있는 영역(1에 가까운 상태)에 머물도록 보장함으로써 필요한 반복 횟수를 줄이고 수치적 안정성을 향상시킵니다.

주요 기여

1. DASH 구현: 순차적 루프를 병렬 3D 텐서 연산으로 대체하여, 잘 최적화된 Distributed Shampoo와 비교했을 때 최대 5.6배 빠른 옵티마이저 단계를 달성하는 GPU 효율적인 프레임워크를 구축했습니다.
2. 새로운 솔버: EVD 및 CN의 더 빠르고 유효한 대안으로서 NDB 반복법과 체비쇼프 다항식 근사를 도입했습니다. 본 논문은 NDB가 CN 및 표준 EVD보다 낮은 검증 퍼플렉시티(validation perplexity)를 달성함을 보여줍니다.
3. 수렴 분석: 표준 프로베니우스 노름 스케일링이 반복법에 있어 최적이 아니며, 종종 더 느린 수렴을 초래한다는 상세한 행동 분석을 제공합니다. 본 논문은 CN과 NDB의 뚜렷한 수렴 동작에 대한 직관적인 설명을 제공합니다.
4. Multi-Power-Iteration: 스펙트럼 반지름을 견고하게 추정하여 NDB 및 CN의 수렴 기준을 만족시키는 최적의 행렬 스케일링을 가능하게 하는 효율적인 반정밀도 기반의 Power-Iteration 구현을 제공합니다.

실험 결과

저자들은 C4 데이터셋으로 학습된 953M 파라미터(임베딩 크기 2048, 시퀀스 길이 1024)를 가진 Llama 모델에서 DASH를 평가했습니다.

• 속도:
  • DASH는 역제곱근 방법 및 정밀도 사용에 따라 Distributed Shampoo 대비 옵티마이저 단계 실행 시간을 4배에서 5.6배 단축합니다.
  • 가장 큰 속도 향상(5.6배)은 Distributed Shampoo의 EVD (FP32)를 DASH의 CN (FP16) 및 GPU 레벨 스태킹과 비교했을 때 관찰되었습니다.
  • 빈번한 프리컨디셔너 업데이트($f=1$) 상황에서도 DASH는 상당한 속도 우위를 유지합니다.
• 정확도 (검증 퍼플렉시티):
  • DASH는 거의 모든 안정적인 구성에서 Distributed Shampoo와 일치하는 성능을 보입니다.
  • NDB는 CN보다 일관되게 낮은 검증 퍼플렉시티를 나타내며, 종종 EVD보다 뛰어난 성능을 보입니다.
  • NDB에 Power-Iteration 스케일링을 사용했을 때(11.68)가 프로베니우스 스케일링(11.76)보다 낮은 검증 퍼플렉시티를 기록하여, 이론적 분석을 입증했습니다.
• 정밀도:
  • 구현체는 퍼플렉시티 저하 없이 CN 및 EVD를 위한 FP16을 지원하여 실행 시간을 더욱 단축합니다.
  • NDB를 FP16/BF16에서 사용할 경우 실험적으로 수치적 불안정성이 발견되었으나, 저자들은 이를 향후 연구 과제로 남겨두었습니다.

의의 및 주장

본 논문은 DASH가 2차 옵티마이저(Shampoo와 같은)를 사용하는 데 따르는 계산적 장벽을 크게 낮추어, 고품질의 프리컨디셔닝을 대규모 딥러닝에서 실용적으로 사용할 수 있게 만든다고 주장합니다.

• 이론과 하드웨어의 결려: 알고리즘 설계(반복 솔버, 행렬 스케일링)를 하드웨어 역량(텐서 코어, 배치 연산)과 정렬함으로써, DASH는 연구자들이 과도한 실행 시간 없이 옵티마이저의 특성을 조사할 수 있게 합니다.
• 수치적 안정성 vs 속도: 본 연구는 반복법이 본질적으로 EVD보다 불안정하거나 느릴 것이라는 가설에 도전합니다. 더 나은 스케일링(Multi-Power-Iteration)과 최적화된 반복(NDB)을 도입함으로써, 저자들은 반복법이 더 빠른 속도로 더 높은 수렴도와 더 낮은 퍼플렉시티를 달성할 수 있음을 입증했습니다.
• 향-미래 방향: 저자들은 NDB의 우수한 성능이 블록 조건수(block condition numbers)에 따른 동적 솔버 선택의 길을 열어주며, 확률적 또는 오류 수정 방법을 사용하여 저정밀도 형식에서 NDB를 안정화할 가능성을 시사한다고 언급했습니다.

결론적으로, Shampoo의 주요 병목 현상은 수학적 복잡성이 아니라 파편화된 실행과 수치적 불안정성이었으며, DASH는 이를 효과적으로 해결했습니다.