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

arXiv: https://arxiv.org/abs/2602.06079 · HTML: https://mlsanity.com/paper/2602.06079/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 31 references that are in the map.
- First posted on arXiv 4 Feb 2026.

## Builds on (most cited first, 21 of 21)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Fantastic Pretraining Optimizers and Where to Find Them](https://mlsanity.com/paper/2509.02046/) (2025)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Preconditioned Stochastic Gradient Descent](https://mlsanity.com/paper/1512.04202/) (2015)
- [Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks](https://mlsanity.com/paper/1811.12019/) (2018)
- [MuonBP: Faster Muon via Block-Periodic Orthogonalization](https://mlsanity.com/paper/2510.16981/) (2025)
- [Controlled LLM Training on Spectral Sphere](https://mlsanity.com/paper/2601.08393/) (2026)
- [ROOT: Robust Orthogonalized Optimizer for Neural Network Training](https://mlsanity.com/paper/2511.20626/) (2025)
- [Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension](https://mlsanity.com/paper/2502.07752/) (2025)
- [Conda: Column-Normalized Adam for Training Large Language Models Faster](https://mlsanity.com/paper/2509.24218/) (2025)

## Built on by (most cited first, 4 of 4)

- [On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability](https://mlsanity.com/paper/2608.30320/) (2026)
- [SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales](https://mlsanity.com/paper/2607.20548/) (2026)
- [DMuon: Efficient Distributed Muon Training with Near-Adam Overhead](https://mlsanity.com/paper/2606.27153/) (2026)
- [FedSubMuon: Communication-Efficient Federated LLM Fine-Tuning via Structured Subspace Muon](https://mlsanity.com/paper/2609.06073/) (2026)

## Abstract (arXiv)

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.