기술 요약: 저랭크 대규모 언어 모델을 위한 병목 최적화 확장성 훈련 프레임워크 BOOST

1. 문제 제기

대규모 언어 모델 (LLM) 의 사전 훈련은 점점 더 계산 및 통신 비용에 의해 제약받고 있습니다. 행렬 분해를 통해 파라미터 수, 메모리 발자국, 계산 요구 사항을 줄임으로써 유망한 알고리즘적 해결책을 제공하는 저랭크 병목 아키텍처 (예: CoLA, LORO, LaX) 는 대규모 분산 훈련에 적용될 때 상당한 확장성 장벽에 직면해 있습니다.

식별된 핵심 문제는 **기존의 3 차원 병렬화 전략 (데이터, 파이프라인, 텐서 병렬화)**이 풀랭크 모델을 위해 설계되었으나, 저랭크 병목 아키텍처에 단순하게 적용될 경우 성능이 매우 저하된다는 점입니다. 구체적으로 다음과 같습니다:

• 통신 폭발: 병목 아키텍처는 더 많은 동기화 지점을 가진 깊은 구조로 구성됩니다. 이러한 구조에 원본 텐서 병렬화 (TP) 를 적용하면 풀 차원 활성화에 대한 과도한 all-reduce 연산이 발생하여, 풀랭크 기준 대비 5~6.5 배 높은 통신 오버헤드를 초래합니다.
• 낮은 GPU 활용도: 원본 TP 는 저랭크 차원 ($r$) 을 따라 쉐딩되므로, 축소 차원이 작은 일반 행렬 곱셈 (GEMM) 커널이 생성됩니다. 이는 산술 강도를 낮추어 실행을 메모리 병목 영역으로 밀어넣고 GPU 계산 유닛을 포화시키지 못하게 하여, 이론적 알고리즘 효율성 향상을 무효화합니다.

2. 방법론: BOOST 프레임워크

저자들은 대규모 저랭크 병목 아키텍처에 맞춘 훈련 프레임워크인 BOOST를 제안합니다. 이 프레임워크는 종단 간 속도 향상을 달성하기 위해 새로운 병렬화 전략과 여러 시스템 수준의 최적화를 통합합니다.

2.1 병목 인식 텐서 병렬화 (BTP)

핵심 혁신은 병목 구조에 정렬되도록 TP 청크의 분할 방식을 재정의하는 BTP입니다.

• 전략 전환: 저랭크 레이어 쌍 ($d \times r$ 및 $r \times d$) 을 단일 청크로 취급하는 원본 TP 와 달리, BTP 는 청크 경계를 이동시킵니다. 이는 업 - 프로젝션 ($r \times d$) 을 컬럼 병렬로, 후속 다운 - 프로젝션 ($d \times r$) 을 행 병렬로 실행하며, 중간 연산은 쉐어드된 활성화에서 수행합니다.
• 통신 감소: 숨겨진 차원 ($d$) 이 아닌 좁은 병목 차원 ($r$) 에서 집단 통신 (all-reduce) 을 수행함으로써, BTP 는 페이로드 크기를 $[b, s, d]$에서 $[b, s, r]$로 줄입니다. 이는 원본 저랭크 TP 대비 통신 볼륨을 약 5.7 배 감소시킵니다.
• 계산 최적화: BTP 는 저랭크 차원 ($r$) 이 아닌 큰 숨겨진 차원 ($d$) 을 따라 쉐딩합니다. 이는 더 큰 GEMM 축소 차원을 유지하여 산술 강도를 높이고 커널이 계산 병목 상태로 유지되도록 하여 GPU 활용도를 개선합니다.

2.2 시스템 수준 최적화

효율성을 더욱 향상시키기 위해 BOOST 는 세 가지 구체적인 최적화를 통합합니다:

1. Online-RMSNorm: 표준 RMSNorm 은 전역 통계가 필요하므로 "쉐어드-안전하지 않음"입니다. BTP 는 RMSNorm 을 TP 청크 내부에 배치합니다. 비용이 많이 드는 동기화를 피하기 위해 저자들은 Online-RMSNorm을 도입하여 로컬 통계를 계산하고, 전역 정규화를 지연시키며, 통계 교환을 후속 TP all-reduce 와 융합합니다. 이는 수학적 동등성과 수치적 안정성을 유지하면서 별도의 작은 페이로드 콜렉티브를 제거합니다.
2. 선형 레이어 그룹화: 커널 시작 오버헤드와 불필요한 데이터 이동을 줄이기 위해 프레임워크는 병렬 선형 레이어를 그룹화합니다. 다운 - 프로젝션 (동일한 입력 공유) 의 경우 가중치를 연결하고, 업 - 프로젝션 (서로 다른 입력) 의 경우 배치화된 GEMM 을 사용합니다. 이는 유효 대역폭과 산술 강도를 증가시킵니다.
3. 통신 없는 저랭크 활성화 체크포인팅: TP 에서의 표준 활성화 체크포인팅은 재전달 (re-forward) 단계에서 추가 통신을 도입하는 경우가 많습니다. BTP 는 체크포인트 경계를 TP 청크 경계와 정렬하여 재전달 경로가 단일 청크 내부에 완전히 머무르도록 합니다. 이는 역전파 중 추가 동기화 지점을 제거하여 체크포인팅을 통신 없이 수행 가능하게 합니다.

3. 주요 기여

논문은 다음과 같이 기여 사항을 요약합니다:

• 이론적 분석: 병목 아키텍처의 분산 훈련에서 산술 강도와 통신 볼륨을 정량화하여 원본 설계의 확장성 한계를 규명했습니다.
• BTP 전략: 저랭크 활성화에서의 효율적인 통신을 용이하게 하고 GEMM 계산 강도를 높이기 위해 가중치 행렬 분할을 최적화하는 새로운 텐서 병렬화 전략입니다.
• BOOST 프레임워크: Online-RMSNorm, 레이어 그룹화, 저랭크 활성화 체크포인팅을 특징으로 하는 고성능 분산 훈련 구현체입니다.
• 실증적 검증: 다양한 저랭크 아키텍처와 모델 크기에서 속도 향상 및 효율성 향상을 입증했습니다.

4. 실험 결과

NERSC-Perlmutter 슈퍼컴퓨터에서 LLaMA-2 계열 모델 (1B~40B 파라미터) 과 다양한 저랭크 아키텍처 (CoLA, SVD, LaX) 를 사용하여 평가가 수행되었습니다.

• 속도 향상: BOOST 는 풀랭크 모델 기준 대비 1.46~1.91 배의 속도 향상을 달성했으며, 단순 통합 3 차원 병렬화 (Vanilla-TP) 를 사용하는 저랭크 모델 대비 1.87~2.27 배의 속도 향상을 보였습니다.
• 통신 효율성: BTP 는 통신 볼륨을 크게 줄여 Vanilla-TP 대비 최대 5.3 배 빠른 통신 시간을 달성했으며, 일부 시나리오에서는 풀랭크 TP 보다 더 나은 성능을 보였습니다.
• 하드웨어 활용도: 더 높은 산술 강도를 유지함으로써 BOOST 는 메모리 병목 실행으로 고통받는 Vanilla-TP 에 비해 훨씬 더 나은 GPU 활용도를 유지합니다.
• 확장성: 이 프레임워크는 1B 에서 40B 파라미터까지 강력한 약 확장성 (weak scaling) 을 보여주며, 감소된 메모리 발자국으로 인해 더 큰 마이크로 배치 크기를 지원합니다.
• 일반성: 이 접근법은 다양한 저랭크 분해 방법 (SVD, CoLA, LaX) 에서 검증되어 특정 병목 구조와 관계없이 일관된 개선을 보여줍니다.

5. 중요성과 주장

논문은 BOOST가 저랭크 알고리즘의 이론적 효율성과 대규모 분산 훈련의 실제 현실 사이의 간극을 해소한다고 주장합니다. 병목 아키텍처에 표준 텐서 병렬화를 적용할 때 발생하는 특정 비효율성을 해결함으로써, 이 프레임워크는 다음을 가능하게 합니다:

• 실용적 확장성: 이전에 채택을 제한했던 통신 및 활용도 패널티 없이 대규모 저랭크 기반 모델의 사전 훈련을 실현 가능하게 합니다.
• 자원 효율성: 모델 정확도를 유지하면서 훈련 시간과 메모리 사용량을 크게 줄입니다.
• 시스템 - 알고리즘 공동 설계: 알고리즘적 효율성 (저랭크 분해) 이 분산 환경에서 완전한 잠재력을 발휘하려면 시스템 수준의 공동 설계 (BTP, Online-RMSNorm) 와 결합되어야 함을 입증합니다.

저자들은 BOOST 를 저랭크 LLM 의 사전 훈련을 위한 실용적이고 확장 가능한 솔루션으로 위치시키며, 현재는 밀집 병목 아키텍처에 초점을 맞추고 있지만, 좁은 활성화로 TP 경계를 이동시키는 기본 원리는 향후 작업에서 전문가 혼합 (MoE) 과 같은 다른 아키텍처로 확장될 수 있다고 언급합니다.