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

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.

💡 Research Summary

The paper addresses a fundamental systems‑algorithm mismatch that prevents the practical use of matrix‑based optimizers (such as Shampoo, Muon, and SOAP) in modern large‑scale language‑model training pipelines. These optimizers require access to an entire weight tensor to compute second‑order statistics (e.g., SVD, Newton‑Schulz updates), but popular distributed training frameworks (Megatron‑style) shard both parameters and optimizer states across data‑parallel (DP) and tensor‑parallel (TP) ranks for memory efficiency. The resulting “Atomicity Constraint” (full‑tensor access) conflicts with the “Zero‑1” sharding strategy, leading existing solutions to either duplicate work synchronously (SC) or to partition at the layer level, which breaks the geometric alignment required for efficient bucket‑based Reduce‑Scatter communication.

Canzona proposes a unified framework that decouples logical optimizer ownership from the physical placement of parameters, satisfying atomicity while preserving the communication patterns of ZeRO‑1 and TP. For DP, the authors introduce a static partitioning scheme that assigns whole parameters to a single rank based on the start index in the flattened buffer. This guarantees that each optimizer state resides entirely on one device, enabling local matrix updates without extra collectives. Because parameter sizes vary dramatically, naïve static assignment creates severe load imbalance. To solve this, the paper formulates a load‑balancing optimization problem and presents an α‑Balanced Greedy LPT (Longest Processing Time) algorithm. The algorithm first sorts buckets by total load, then distributes them across ranks using a blended target that mixes an even‑share vector with a deficit‑fill vector controlled by α. Cuts are snapped to actual parameter boundaries to respect atomicity, resulting in near‑equal computational load and minimal pipeline bubbles.

For TP, the challenge is that each weight matrix is split across ranks, so holistic matrix operations still require reconstruction. Canzona’s solution is a Micro‑Group Scheduling pipeline. Tensors with similar computational characteristics are grouped into micro‑batches; each group is assigned a host rank that drives an All‑to‑All reconstruction followed by the matrix‑based update. By overlapping reconstruction with computation and by balancing the groups’ workloads, the pipeline hides communication latency and eliminates global synchronization points. The design remains compatible with the existing bucketed communication used in forward and backward passes, preserving the high‑throughput Reduce‑Scatter/All‑Gather pattern.

The implementation modifies the Megatron‑LM codebase minimally: DP‑ASC adds metadata for the α‑balanced cuts and reorders the optimizer‑state buffer; TP‑Micro‑Group introduces asynchronous CUDA streams and a scheduler that launches grouped All‑to‑All and matrix kernels. Experiments were conducted on a 256‑GPU cluster (8×A100 80 GB) training the Qwen‑3 family (1.7 B, 6 B, 16 B, 32 B parameters). Three matrix‑based optimizers (Shampoo, Muon, SOAP) were evaluated. Compared with the baseline ZeRO‑1 + layer‑wise partitioning, Canzona achieved an average 1.57× speed‑up in total iteration time (forward + backward + optimizer) and a 5.8× reduction in optimizer‑step latency alone. Memory consumption remained on par with the baseline, and communication volume did not increase because the bucketed Reduce‑Scatter pattern was retained.

The contributions are fourfold: (1) a unified, atomicity‑preserving framework that works with both DP and TP; (2) a static load‑balancing algorithm (α‑Balanced Greedy LPT) that eliminates stragglers in DP; (3) a micro‑group asynchronous pipeline for TP that overlaps reconstruction and computation; (4) extensive large‑scale validation across multiple optimizers and model sizes, demonstrating practical gains. The authors suggest future work on online dynamic scheduling, extension to higher‑dimensional expert or multimodal tensors, tighter integration with hardware accelerators, and automatic tuning of the α parameter based on runtime profiling.

In summary, Canzona bridges the gap between the algorithmic requirements of matrix‑based optimizers and the engineering constraints of state‑of‑the‑art distributed training, delivering significant speed‑ups without sacrificing memory efficiency or communication performance. This work paves the way for broader adoption of second‑order optimizers in training ever larger language models.