提出Canzona框架,解决分布式优化器的负载不均与通信瓶颈问题。
Canzona: A Unified, Asynchronous, and Load-Balanced Framework for Distributed Matrix-based Optimizers
- 通过逻辑与物理分离设计,实现异步更新与负载均衡。
- 在256张GPU上训练320亿参数模型,迭代速度提升1.57倍,优化步延迟降低5.8倍。
- 适合大规模语言模型训练,尤其适用于数据并行与张量并行混合场景。
大语言模型的扩展推动了基于矩阵的优化器(如Shampoo、Muon、SOAP)的研究,因其收敛效率高而受到关注;然而,这些优化器要求整体更新,与Megatron等分布式框架中的张量碎片化存在冲突。现有方案不理想:同步方法存在计算冗余,层级划分又无法在不违反高效通信原语几何约束的前提下调和此矛盾。为此,我们提出Canzona,一个统一、异步且负载均衡的分布式矩阵优化框架。针对数据并行,引入alpha平衡静态分区策略,在保持原子性的同时消除负载不均;针对张量并行,设计异步计算流水线,利用微组调度批量处理碎片化更新,隐藏重构开销。在256张GPU上对Qwen3系列模型(最大320亿参数)的大量实验表明,该方法保持了现有并行架构的效率,实现端到端迭代时间1.57倍加速,优化器步延迟降低5.8倍。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。