Canzona: un Framework Unificato, Asincrono e a Bilanciamento del Carico per Ottimizzatori Distribuiti Basati su Matrici

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

February 4, 2026

Autori: Liangyu Wang, Siqi Zhang, Junjie Wang, Yiming Dong, Bo Zheng, Zihan Qiu, Shengkun Tang, Di Wang, Rui Men, Dayiheng Liu

cs.AI

Abstract

Il ridimensionamento dei Large Language Model (LLM) stimola l'interesse per ottimizzatori basati su matrici (ad es. Shampoo, Muon, SOAP) per la loro efficienza di convergenza; tuttavia, la loro necessità di aggiornamenti olistici entra in conflitto con la frammentazione dei tensori nei framework distribuiti come Megatron. Le soluzioni esistenti sono subottimali: gli approcci sincroni soffrono di ridondanza computazionale, mentre il partizionamento per layer non riesce a conciliare questo conflitto senza violare i vincoli geometrici delle primitive di comunicazione efficiente. Per colmare questa lacuna, proponiamo Canzona, un framework Unificato, Asincrono e Bilanciato nel Carico che disaccoppia l'assegnazione logica dell'ottimizzatore dalla distribuzione fisica dei parametri. Per il Parallelismo dei Dati, introduciamo una strategia di Partizionamento Statico Bilanciato-alfa che rispetta l'atomicità neutralizzando al contempo lo squilibrio di carico. Per il Parallelismo dei Tensori, progettiamo una pipeline di Calcolo Asincrono che utilizza lo Scheduling a Micro-Gruppi per raggruppare gli aggiornamenti frammentati e mascherare l'overhead di ricostruzione. Valutazioni estese sulla famiglia di modelli Qwen3 (fino a 32 miliardi di parametri) su 256 GPU dimostrano che il nostro approccio preserva l'efficienza delle architetture parallele consolidate, ottenendo un miglioramento di 1,57x nel tempo di iterazione end-to-end e riducendo la latenza del passo di ottimizzazione di 5,8x rispetto al baseline.

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.