Canzona: Un Marco Unificado, Asincrónico y de Equilibrio de Carga para Optimizadores Distribuidos Basados en Matrices
Canzona: A Unified, Asynchronous, and Load-Balanced Framework for Distributed Matrix-based Optimizers
February 4, 2026
Autores: Liangyu Wang, Siqi Zhang, Junjie Wang, Yiming Dong, Bo Zheng, Zihan Qiu, Shengkun Tang, Di Wang, Rui Men, Dayiheng Liu
cs.AI
Resumen
La escalada de los Modelos de Lenguaje Grandes (LLM) impulsa el interés en optimizadores basados en matrices (por ejemplo, Shampoo, Muon, SOAP) por su eficiencia de convergencia; sin embargo, su requisito de actualizaciones holísticas entra en conflicto con la fragmentación de tensores en frameworks distribuidos como Megatron. Las soluciones existentes son subóptimas: los enfoques síncronos adolecen de redundancia computacional, mientras que la partición por capas no logra reconciliar este conflicto sin violar las restricciones geométricas de las primitivas de comunicación eficientes. Para salvar esta brecha, proponemos Canzona, un framework Unificado, Asíncrono y con Equilibrio de Carga que desacopla la asignación lógica del optimizador de la distribución física de parámetros. Para el Paralelismo de Datos, introducimos una estrategia de Partición Estática alfa-Equilibrada que respeta la atomicidad mientras neutraliza el desequilibrio de carga. Para el Paralelismo de Tensores, diseñamos un pipeline de Cómputo Asíncrono que utiliza Micro-Grupos de Planificación para agrupar actualizaciones fragmentadas y ocultar la sobrecarga de reconstrucción. Evaluaciones exhaustivas en la familia de modelos Qwen3 (hasta 32B parámetros) en 256 GPUs demuestran que nuestro enfoque preserva la eficiencia de las arquitecturas paralelas establecidas, logrando una aceleración de 1.57x en el tiempo de iteración integral y reduciendo la latencia del paso del optimizador en 5.8x comparado con la línea base.
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.