Canzona: 분산 행렬 기반 최적화 프레임워크를 위한 통합적 비동기 부하 분산 아키텍처

초록

대규모 언어 모델(LLM)의 규모 확장으로 인해 수렴 효율성 때문에 행렬 기반 최적화 도구(예: Shampoo, Muon, SOAP)에 대한 관심이 높아지고 있습니다. 그러나 이들 최적화 도구가 요구하는 전체적 업데이트 방식은 Megatron과 같은 분산 프레임워크의 텐서 분할 방식과 상충됩니다. 기존 해결책은 최적이 아닙니다: 동기식 접근법은 계산 중복 문제가 있고, 계층별 분할 방식은 효율적인 통신 기본 요소의 기하학적 제약 조건을 위반하지 않고는 이 상충을 해결하지 못합니다. 이러한 격차를 해소하기 위해 우리는 논리적 최적화 도구 할당과 물리적 매개변수 분배를 분리하는 통합적이고 비동기적이며 부하 균형을 이룬 프레임워크인 Canzona를 제안합니다. 데이터 병렬화를 위해 원자성을 보장하면서 부하 불균형을 해소하는 알파 균형 정적 분할 전략을 도입합니다. 텐서 병렬화를 위해 분할된 업데이트를 일괄 처리하고 재구성 오버헤드를 숨기기 위해 마이크로 그룹 스케줄링을 활용한 비동기 컴퓨트 파이프라인을 설계합니다. 256개의 GPU에서 Qwen3 모델 패밀리(최대 320억 매개변수)에 대한 폭넓은 평가를 통해 우리의 접근 방식이 기존 병렬 아키텍처의 효율성을 유지하면서, 기준 대비 종단 간 반복 시간에서 1.57배의 속도 향상을 달성하고 최적화 도구 단계 지연 시간을 5.8배 줄이는 것을 입증했습니다.

English

The scaling of Large Language Models (LLMs) drives interest in matrix-based optimizers (e.g., Shampoo, Muon, SOAP) for their convergence efficiency; yet their requirement for holistic updates conflicts with the tensor fragmentation in distributed frameworks like Megatron. Existing solutions are suboptimal: synchronous approaches suffer from computational redundancy, while layer-wise partitioning fails to reconcile this conflict without violating the geometric constraints of efficient communication primitives. To bridge this gap, we propose Canzona, a Unified, Asynchronous, and Load-Balanced framework that decouples logical optimizer assignment from physical parameter distribution. For Data Parallelism, we introduce an alpha-Balanced Static Partitioning strategy that respects atomicity while neutralizing the load imbalance. For Tensor Parallelism, we design an Asynchronous Compute pipeline utilizing Micro-Group Scheduling to batch fragmented updates and hide reconstruction overhead. Extensive evaluations on the Qwen3 model family (up to 32B parameters) on 256 GPUs demonstrate that our approach preserves the efficiency of established parallel architectures, achieving a 1.57x speedup in end-to-end iteration time and reducing optimizer step latency by 5.8x compared to the baseline.

Canzona: 분산 행렬 기반 최적화 프레임워크를 위한 통합적 비동기 부하 분산 아키텍처

Canzona: A Unified, Asynchronous, and Load-Balanced Framework for Distributed Matrix-based Optimizers

초록

Support