양면 저랭크 통신으로 Adam 효율 극대화

초록

본 논문은 대규모 언어 모델 사전학습에서 발생하는 통신 병목을 해결하기 위해, Adam 계열 옵티마이저에 양면 저랭크(두‑사이드) 코어 동기화 방식을 도입한 TSR‑Adam을 제안한다. 매 스텝마다 행·열 두 개의 저랭크 기저 U, V를 유지하고, 핵심 행렬 C = UᵀGV (크기 r×r)만을 All‑Reduce로 전파함으로써 기존 O(mn) 혹은 O(r n) 통신량을 O(r²)로 감소시킨다. 또한 랜덤화된 SVD 기반의 기저 갱신을 사용해 피크 통신량을 억제하고, 임베딩 층에 특화된 저랭크 처리와 갱신 스케줄을 적용한다. 60M‑1B 규모 모델 사전학습과 GLUE 파인튜닝 실험에서 평균 통신량을 각각 13배·25배 절감하면서 성능 저하 없이 학습을 수행한다.

상세 분석

TSR‑Adam은 기존 저랭크 옵티마이저가 메모리 절감에 초점을 맞추어 한쪽(단일) 기저만을 전파하고, 전체 그라디언트를 재구성하기 위해 O(r n) 혹은 O(m r) 크기의 행렬을 교환하는 한계를 극복한다. 논문은 먼저 매 스텝마다 파라미터 블록 W ∈ ℝ^{m×n}에 대해 두 개의 직교 기저 U ∈ ℝ^{m×r}, V ∈ ℝ^{n×r}를 유지한다는 가정을 둔다. 각 워커는 로컬 그라디언트 G 을 받아 C = UᵀGV (크기 r×r)를 계산하고, All‑Reduce를 통해 평균 \bar{C} 을 얻는다. 이후 \bar{G}=U\bar{C}Vᵀ 를 복원해 AdamW의 1차·2차 모멘트를 r×r 코어 공간에 저장하고 업데이트한다. 이 과정에서 전송되는 데이터는 r² 스칼라에 불과하므로, m·n 규모의 파라미터에 비해 통신량이 급격히 감소한다.

핵심 기술은 (1) 양면 저랭크 코어 동기화와 (2) 랜덤화 SVD 기반 기저 갱신이다. 기저 갱신은 일정 주기 K 마다 수행되며, 전체 그라디언트를 직접 전송하지 않는다. 대신 각 워커가 로컬 그라디언트에 무작위 행렬 Ω 을 곱해 스케치 Y=GΩ 를 만든 뒤, QR 분해와 All‑Reduce를 통해 전역 기저 U, V 를 추정한다. 이때 오버샘플링 파라미터 p 을 추가해 정확도를 보강한다. 이렇게 하면 피크 통신량이 O(m r + n r)에서 O(r² + r·p) 수준으로 낮아진다.

또한 임베딩 층은 행렬 차원이 매우 크고, 기존 저랭크 방법이 이를 밀집 형태로 유지하는 문제점을 지적한다. TSR‑Adam은 임베딩 전용 저랭크 차원 r_emb 과 별도 갱신 주기 K_emb 을 설정해, 임베딩 그라디언트 역시 r_emb² 수준의 코어만 전파한다. 이는 전체 통신량을 추가로 절감하고, 메모리 사용량도 감소시킨다.

이론적으로는 동기화된 코어 \bar{C} 가 원래 그라디언트 \bar{G} 의 저랭크 근사임을 보이며, Adam 업데이트가 코어 공간에서 수행될 때 기대값이 원래 Adam과 동일함을 증명한다. 즉, 수렴성(정상점) 보장을 유지하면서 통신 효율을 높인다.

실험에서는 LLaMA‑60M,‑130M,‑350M,‑1B 모델을 대상으로 사전학습을 수행했으며, 평균 통신량이 13배 감소하고, GLUE 벤치마크 파인튜닝에서는 25배 절감했다. 손실 곡선은 동일하거나 약간 개선된 수준을 보였으며, 메모리 사용량도 저랭크 상태 저장 덕분에 크게 감소했다.

요약하면, TSR‑Adam은 (i) 양면 저랭크 코어 전파로 O(r²) 통신량 달성, (ii) 랜덤화 SVD로 피크 통신 억제, (iii) 임베딩 전용 저랭크 설계로 전체 시스템 효율을 극대화한다는 세 가지 핵심 기여를 제공한다.