Canzona 통합 비동기 로드밸런싱 매트릭스 옵티마이저 프레임워크

초록

본 논문은 대규모 언어 모델 학습에 사용되는 매트릭스 기반 옵티마이저(Shampoo, Muon, SOAP 등)의 전체 텐서 접근 요구와 Megatron‑style 분산 학습에서 발생하는 텐서 파편화 사이의 충돌을 해결한다. 저자는 원자성(Atomicity) 제약을 유지하면서도 부하 불균형을 최소화하는 α‑Balanced 정적 파티셔닝과 텐서 병렬에서 마이크로‑그룹 스케줄링을 결합한 Canzona 프레임워크를 제안한다. 256 GPU 환경에서 Qwen‑3 모델(1.7B‑32B) 실험 결과, 기존 방식 대비 전체 반복 시간이 1.57배 빨라지고 옵티마이저 단계 지연이 5.8배 감소하였다.

상세 분석

Canzona는 매트릭스 기반 옵티마이저가 요구하는 “전체 텐서 원자성”을 시스템 수준에서 강제하면서도 기존 ZeRO‑1 및 Tensor Parallel(TP) 구조와 호환되도록 설계되었다. 데이터 병렬(DP)에서는 기존의 “Equal Chunk” 방식이 파라미터 경계와 무관하게 버퍼를 균등하게 나누어 원자성을 깨뜨리는 반면, Canzona는 파라미터의 시작 인덱스를 기준으로 각 랭크에 전체 파라미터를 할당하는 정적 파티셔닝을 도입한다. 이때 파라미터는 절대로 분할되지 않으며, ZeRO‑1의 버킷 구조와 정확히 일치하도록 배치된다. 그러나 파라미터 크기가 크게 차이 나는 경우 단순 정적 할당은 특정 랭크에 과부하를 초래한다. 이를 해결하기 위해 저자는 부하 함수 W(p)를 파라미터 크기(또는 계산 복잡도)로 정의하고, 전체 버킷을 가중치 기준으로 내림차순 정렬한 뒤 LPT(Longest Processing Time) 알고리즘에 α‑Balanced 가중치를 혼합한 최적화 절차를 제시한다. α 파라미터는 부하 균형과 통신 균형 사이의 트레이드오프를 조정하며, 실제 파라미터 경계 내에서 가능한 가장 가까운 커트를 선택해 원자성 제약을 위배하지 않는다. 결과적으로 각 랭크는 거의 동일한 연산량을 갖게 되어 파이프라인 버블이 최소화되고, Reduce‑Scatter/All‑Gather와 같은 ZeRO‑1의 고효율 통신 패턴을 그대로 활용할 수 있다.

텐서 병렬(TP)에서는 각 텐서가 여러 랭크에 분산되어 있어 매트릭스 연산을 수행하려면 재구성이 필요하다. Canzona는 “Micro‑Group Scheduling”이라는 비동기 파이프라인을 도입한다. 먼저 동일한 연산 특성을 가진 텐서들을 마이크로‑그룹으로 묶고, 각 그룹을 담당할 호스트 랭크를 지정한다. 이후 All‑to‑All 재구성 연산을 그룹 단위로 배치하고, 재구성과 매트릭스 연산을 겹쳐 실행함으로써 통신 지연을 숨긴다. 또한, 그룹 스케줄링 단계에서 부하 균형을 고려해 무거운 텐서는 여러 랭크에 분산시키고, 가벼운 텐서는 하나의 랭크에 집중시켜 전체 실행 시간을 최소화한다. 이 설계는 기존의 동기식 재구성 방식이 초래하는 전역 대기 시간을 크게 감소시킨다.

실험에서는 Qwen‑3 시리즈(1.7B, 6B, 16B, 32B) 모델을 256 GPU 클러스터에 배포했으며, Shampoo, Muon, SOAP 등 세 종류의 매트릭스 옵티마이저를 모두 적용해 성능을 검증했다. 결과는 DP‑ASC(α‑Balanced 정적 파티셔닝)와 TP‑Micro‑Group 파이프라인을 결합했을 때, 기존 ZeRO‑1 기반 레이어‑와이즈 파티셔닝 대비 전체 반복 시간이 평균 1.57배 빨라졌고, 옵티마이저 단계만 따로 측정했을 때는 5.8배의 지연 감소를 보였다. 메모리 사용량은 기존 방식과 동일하거나 약간 감소했으며, 통신량도 ZeRO‑1 버킷 구조를 그대로 유지함으로써 추가적인 오버헤드가 발생하지 않았다.

이 논문은 매트릭스 기반 옵티마이저를 대규모 분산 학습에 적용하기 위한 시스템‑알고리즘 공동 설계의 좋은 사례를 제시한다. 원자성 제약을 만족하면서 부하 불균형을 정량적으로 모델링하고, α‑Balanced LPT와 마이크로‑그룹 스케줄링이라는 두 가지 독립적인 최적화 기법을 통해 DP와 TP 양쪽 모두에서 효율성을 확보했다. 향후 연구에서는 동적 워크로드 변화에 대응하는 온라인 스케줄링, 더 높은 차원의 텐서 파티셔닝, 그리고 다른 종류의 두 번째 차원(예: 전문가 모델)과의 결합 가능성을 탐색할 여지가 있다.