Canzona : Un cadre unifié, asynchrone et à équilibrage de charge pour les optimiseurs distribués basés sur les matrices
Canzona: A Unified, Asynchronous, and Load-Balanced Framework for Distributed Matrix-based Optimizers
February 4, 2026
Auteurs: Liangyu Wang, Siqi Zhang, Junjie Wang, Yiming Dong, Bo Zheng, Zihan Qiu, Shengkun Tang, Di Wang, Rui Men, Dayiheng Liu
cs.AI
Résumé
L'essor des modèles de langage à grande échelle (LLM) suscite un intérêt croissant pour les optimiseurs matriciels (tels que Shampoo, Muon, SOAP) en raison de leur efficacité de convergence. Cependant, leur nécessité de mettre à jour les paramètres de manière holistique entre en conflit avec la fragmentation des tenseurs dans les cadres distribués comme Megatron. Les solutions existantes sont sous-optimales : les approches synchrones souffrent de redondance computationnelle, tandis que le partitionnement par couches échoue à résoudre ce conflit sans violer les contraintes géométriques des primitives de communication efficaces. Pour combler cette lacune, nous proposons Canzona, un cadre unifié, asynchrone et équilibré qui dissocie l'attribution logique de l'optimiseur de la distribution physique des paramètres. Pour le parallélisme de données, nous introduisons une stratégie de Partitionnement Statique Alpha-Équilibré qui respecte l'atomicité tout en neutralisant le déséquilibre de charge. Pour le parallélisme de tenseurs, nous concevons un pipeline de Calcul Asynchrone utilisant un Ordonnancement par Micro-Groupes pour regrouper les mises à jour fragmentées et masquer la surcharge de reconstruction. Des évaluations approfondies sur la famille de modèles Qwen3 (jusqu'à 32 milliards de paramètres) sur 256 GPU démontrent que notre approche préserve l'efficacité des architectures parallèles établies, obtenant une accélération de 1,57x du temps d'itération de bout en bout et réduisant la latence des étapes de l'optimiseur de 5,8x par rapport à la référence.
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.