PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Liangyu Wang, Siqi Zhang, Junjie Wang, Yiming Dong, Bo Zheng, Zihan Qiu, Shengkun Tang, Di Wang, Rui Men, Dayiheng Liu

arXiv:2602.06079v1cs.DCcs.LG

TL;DR

Matrix-based optimizers require holistic tensor updates, conflicting with the tensor fragmentation used in distributed LLM training. Canzona decouples logical optimizer assignment from physical parameter distribution through load-balanced data-parallel partitioning and asynchronous tensor-parallel micro-group scheduling, achieving a 1.57x end-to-end iteration speedup and 5.8x lower optimizer step latency than the baseline.

  • Problem

    Matrix-based optimizers require full tensor dimensions for holistic updates, while distributed frameworks fragment tensors across data- and tensor-parallel ranks, creating a system-algorithm conflict.

  • Method

    Canzona decouples logical optimizer assignment from physical parameter distribution using α-Balanced Static Partitioning for Data Parallelism and asynchronous Micro-Group Scheduling for Tensor Parallelism.

  • Results

    1.57x end-to-end iteration-time speedup and 5.8x lower optimizer step latency were achieved on Qwen3 models using 256 GPUs, compared with the baseline.

  • Takeaways & Limitations

    Canzona reconciles atomicity requirements with massive parallelism while preserving efficient communication overlap and reducing load imbalance in distributed matrix-based optimization.

  • Takeaways & Limitations

    The workload formulations in Sections 3.2 and 4.2 assume the simplified linear cost model W(p) = numel(p) for clarity, although the framework supports generalized nonlinear complexity.

Abstract

from arXiv · show

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.

1. Introduction

Canzona addresses the conflict between holistic matrix-based optimizer updates and fragmented distributed parameters by decoupling logical optimizer ownership from physical parameter placement. Its DP and TP strategies preserve atomicity while reducing redundant computation and balancing heterogeneous workloads.

  • Motivation: Matrix-based optimizers require full tensor dimensions for holistic updates, conflicting with ZeRO-1 and TP sharding that fragment tensors across ranks.This atomicity constraint creates a system–algorithm conflict in large-scale training.
  • Motivation: Existing synchronous approaches preserve atomicity through redundant computation, while layer-wise partitioning conflicts with ZeRO-1 geometric layout and efficient bucket communication.These alternatives compromise scalability or communication efficiency.
  • Data Parallelism: For DP, α-Balanced Static Partitioning assigns whole parameters to ranks, preserving atomicity and enabling zero-communication local optimizer updates.The layout also preserves ZeRO-style geometric alignment with bucketed parameters.
  • Tensor Parallelism: For TP, Asynchronous Micro-Group Scheduling batches fragmented updates and overlaps reconstruction communication with computation.Micro-groups support fused All-to-All communication and load-balanced host-rank scheduling.
  • Framework: Canzona decouples logical optimizer assignment from physical parameter distribution in a unified, optimizer-agnostic architecture.The framework assigns optimizer tasks to designated host ranks and executes them asynchronously across ranks.
  • Evaluation: 1.57x end-to-end iteration speedup and 5.8x lower optimizer step latency were achieved on Qwen3 models from 1.7B to 32B parameters using 256 GPUs.The evaluation also covered Muon, Shampoo, and SOAP.

2. Preliminary

Modern LLM training combines DP and TP to scale computation, using ZeRO-1 bucketed state sharding and TP weight-matrix partitioning. These mechanisms improve memory efficiency and distribute matrix operations across devices.

  • Data Parallelism: DDP replicates optimizer states, whereas ZeRO-1 partitions them into contiguous parameter and gradient-buffer buckets.Buckets pipeline Reduce-Scatter and All-Gather communication with computation.
  • Tensor Parallelism: Tensor Parallelism orthogonally splits weight matrices across devices to distribute model computation.

3. Load-Balanced Asynchronous Compute for Data Parallelism Matrix-based Optimizer

Canzona’s DP design assigns complete parameters within ZeRO-1-compatible layouts, then balances their heterogeneous costs through static scheduling. This preserves atomicity, avoids optimizer-step collectives, and reduces stragglers while retaining bucket communication.

  • Design Paradigms: ZeRO-1 equal chunking can split individual tensors across ranks, making local holistic matrix operations such as SVD impossible without reconstruction.
  • Design Paradigms: Synchronous compute replicates optimizer states and matrix operations across ranks, preserving correctness but limiting scalability through redundant work.
  • Design Paradigms: Layer-wise partitioning misaligns logical load balancing with ZeRO-1’s physical bucket geometry, sacrificing efficient bucket-based communication.
  • Static Partitioning: Static Partitioning assigns each parameter’s optimizer states exclusively to one rank according to its flattened-buffer start position.This preserves parameter atomicity and sequential ZeRO-1 ordering without physically reordering parameters.
  • Load Imbalance: Naive atomic assignment creates stragglers and pipeline bubbles because matrix-based optimizer costs vary non-linearly across parameters.The resulting waits can also increase peak memory usage.
  • Load Balancing: Boundary shifting balances bucket workloads without reordering parameters, retaining coalesced Reduce-Scatter primitives.
  • Load Balancing: α-Balanced Greedy LPT treats partitioning as a discrete load-balancing problem and trades global workload balance against bucket communication balance.The heuristic processes heavy blocks early and projects assignments onto valid atomic cuts.

4. Tensor Parallelism with Load-Balanced Asynchronous Compute

Canzona’s TP strategy treats each split-parameter update as an asynchronous task hosted on a designated rank. It batches gradients into capacity-bounded micro-groups, balances their computation, and executes a static runtime plan.

  • TP Compute Strategies: Standard TP synchronously All-Gathers full tensors and redundantly computes identical updates on every TP rank.This limits scalability and efficiency.
  • Asynchronous Compute: TP-ASC batches multiple tensor gradients into fused Micro-Groups and assigns them to host ranks using offline load balancing.The approach is viable within high-bandwidth intra-node TP communication.
  • Task Abstraction: Each TP-split parameter becomes an atomic Compute Task assigned to a specific host rank, with optimizer states initialized locally and retained there.
  • Workload Scheduling: Micro-Group Construction uses global LPT sorting, greedy packing, and rollback when a group would exceed the capacity constraint Cmax.
  • Execution Workflow: Fused All-to-All communication gathers gradients for designated host ranks while avoiding small-message overhead and excessive memory peaks.Host ranks then execute matrix operations using gathered gradients and locally resident optimizer states.
  • Workload Scheduling: The scheduling objective first minimizes within-group computational makespan, then maximizes group saturation without violating Cmax.
  • Runtime Integration: The precomputed micro-group sequence serves as a static execution plan that the runtime iterates during training.
  • Unified Architecture: DP and TP share a unified, optimizer-agnostic architecture that treats tensor updates as generic cost-defined tasks.Neither scheduler requires changes to the optimizer’s internal mathematics.

5. Experiment

Experiments evaluate Canzona’s LB-ASC strategy on Qwen3 models across distributed configurations, measuring load balance, iteration efficiency, and convergence preservation. LB-ASC reduces workload imbalance, accelerates training relative to NV-layerwise, and matches the synchronous baseline’s loss trajectory.

  • Experimental Setup: Experiments evaluate Qwen3 models from 1.7B to 32B parameters using Muon, with verification on Shampoo and SOAP.The main cluster configuration uses 256 GPUs with DP size 32 and TP size 8.
  • Load Balancing: α-Balanced Partitioning reduces DP FLOPs and memory imbalance ratios from 3.24× and 2.46× to 1.43× and 1.11×.The strategy flattens workload distribution relative to static partitioning without DP load balancing.
  • Load Balancing: Micro-Group Scheduling improves TP FLOPs and memory ratios to 2.46× and 1.16×, while LB-ASC achieves the lowest maximum step time at 1.05 TFLOPS equivalent.The results connect improved balance with reduced computation bubbles.
  • End-to-End Efficiency: 5.8× speedup: optimizer step time falls from 0.383s to 0.066s, while forward-backward time decreases from 0.998s to 0.811s.The method uses zero-communication DP updates and asynchronous TP processing, while preserving efficient communication structure.
  • Precision Verification: LB-ASC has an indistinguishable training-loss trajectory from the synchronous baseline on Qwen3-1.7B trained for 400B tokens.This verifies convergence behavior under DP=8 and TP=4 with Muon.

6. Related Work

Related work frames distributed matrix-based optimization as a conflict between strict atomicity requirements and standard sharding and communication primitives. Existing approaches address parts of this conflict through approximation or layerwise assignment but incur communication or geometric costs.

  • System-Algorithm Conflict: Distributed matrix-based optimizers require strict atomicity that conflicts with ZeRO’s parameter sharding in modern distributed training.The conflict arises because exact optimizers such as Muon and Shampoo need holistic matrix access.
  • Existing Approaches: Layerwise optimizer assignment disregards ZeRO geometric constraints, forcing costly All-Reduce communication with 2× volume instead of efficient bucket-based Reduce-Scatter.The geometric incompatibility also requires additional Broadcast or All-Gather synchronization during optimizer updates.
  • Existing Approaches: Distributed K-FAC mitigates communication costs through approximate curvature, whereas modern exact optimizers retain stricter atomicity requirements.This distinguishes approximation-based distributed optimization from the exact matrix-based setting addressed by Canzona.

7. Conclusion

Canzona reconciles matrix-optimizer atomicity with large-scale parallel training through balanced DP partitioning and asynchronous TP scheduling. Experiments on Qwen3 demonstrate a 1.57× end-to-end speedup while reducing communication overhead and load imbalance.

  • Contributions: α-Balanced Static Partitioning for DP and Asynchronous Micro-Group Scheduling for TP reconcile strict atomicity with massive parallelism.The two mechanisms address data-parallel ownership and tensor-parallel fragmented updates respectively.
  • Results: 1.57× end-to-end speedup: Qwen3 experiments show that Canzona reduces communication overhead and neutralizes load imbalance in matrix-based optimization.The conclusion presents this as the principal reported performance outcome.

B.1. Data Parallelism and ZeRO

Distributed training frameworks partition model states and tensors for memory and communication efficiency, but matrix-based optimizers require complete tensors, creating an atomicity conflict with arbitrary sharding.

  • ZeRO-1 partitioning: ZeRO-1 partitions optimizer states across data-parallel ranks, reducing each rank’s responsibility to approximately 1/R of the parameters.Megatron implements this through a unified flattened parameter and gradient buffer.
  • ZeRO-1 partitioning: Megatron flattens parameters and gradients into contiguous buffers divided into sequential buckets for communication overlap.Each bucket is rigidly split into R equal contiguous segments, independent of tensor boundaries.
  • ZeRO-1 workflow: Reduce-Scatter shards bucket gradients during backward, ranks update assigned ranges, and All-Gather reconstructs updated parameter shards before forward.These operations are overlapped with neighboring computation where supported by the bucketed workflow.
  • Optimizer compatibility: Element-wise optimizers tolerate arbitrary tensor cuts because each scalar update depends only on its own gradient and historical states.Matrix-based optimizers instead operate at tensor granularity using operations such as SVD or matrix multiplication.
  • Optimizer compatibility: Matrix-based optimizers impose an Atomicity Constraint requiring complete tensors, so arbitrary Megatron slicing can make local matrix updates impossible without reconstruction communication.Partition strategies must therefore respect parameter boundaries.

C.1. Full Experiment Setup

The evaluation tests Canzona across Qwen3 scales, matrix-based optimizers, distributed baselines, and runtime metrics, using repeated measurements on configurations reaching 256 GPUs.

  • Models and optimizers: Experiments scale across Qwen3 models from 1.7B to 32B parameters and include Muon, Shampoo, and SOAP.The detailed performance analysis primarily uses Qwen3-32B.
  • Baselines and methods: The comparison includes synchronous compute, NVIDIA’s layerwise optimizer, asynchronous compute, and Canzona’s load-balanced asynchronous strategy.Canzona combines α-Balanced Static Partitioning for DP with Greedy LPT scheduling for TP.
  • Evaluation metrics: Runtime breakdowns report Forward-Backward time, standard AdamW optimizer time, and Target Optimizer time as the primary comparison metric.Timings are averaged over 10 runs and 10 steps within each run.

C.2. Full Performance Comparison with layerwise optimizer

Canzona consistently outperforms NVIDIA’s layerwise optimizer across tested Qwen3 configurations, with advantages driven mainly by lower optimizer-step cost while preserving efficient forward-backward communication.

  • End-to-end comparison: Approximately 8.3x lower specific optimizer latency is achieved for Qwen3-32B with DP16-TP8 versus NV-layerwise.The framework consistently outperforms NV-layerwise in every tested configuration.
  • End-to-end comparison: The performance gap widens with model size as matrix-based optimizer computation becomes dominant and NV-layerwise exposes collective communication overhead.The contrast is attributed primarily to reduced optimizer time in the proposed framework.
  • Parallelism robustness: The advantage remains robust across different parallelism splits, including DP16-TP8 and DP32-TP4.Static partitioning and micro-group scheduling adapt to changing topology and parallelism dimensions.
  • Communication efficiency: Canzona’s forward-backward latency closely tracks AdamW Reduce-Scatter, whereas NV-layerwise matches AdamW All-Reduce.Variable-sized chunks cause a marginal discrepancy from ideal Reduce-Scatter, but the reported impact is negligible relative to All-Reduce overhead.

C.3. Scaling Analysis

Across parallelism and model-size scaling, Canzona’s load-balancing strategies maintain near-optimal workload balance and optimizer latency, while preserving correctness and improving efficiency. The framework also reduces communication and scheduling costs through fused operations and offline planning.

  • Parallelism Scaling: As DP scales from 16 to 128, DP LB-ASC maintains a load-balance ratio near 1.0 and stable optimizer time, unlike ASC’s worsening imbalance.The baseline’s imbalance increases with parameter-size variance and produces straggler-driven optimizer slowdown.
  • Parallelism Scaling: As TP increases, Micro-Group Scheduling keeps the FLOPs load-balance ratio substantially below the baseline and preserves low optimizer latency.The strategy mitigates the fragmentation-induced imbalance that worsens with larger TP groups.
  • Model Size Scaling: Across model sizes from 1.7B to 32B, LB-ASC maintains a consistent DP load-balance profile despite increasing tensor-shape heterogeneity.For TP, greedy scheduling sustains near-optimal packing across architecture-dependent imbalance patterns.
  • Precision: The framework validates exact training behavior: Shampoo and SOAP loss curves overlap the standard synchronous baseline, with zero mathematical precision loss reported.The optimization is system-level and does not approximate curvature or skip updates.
  • Load Balance: Muon, Shampoo, and SOAP evaluations show that LB-ASC flattens workload variance, reducing the FLOPs load-balance ratio from > 2.0 to ≈1.05 for Shampoo and SOAP.The result indicates effective hardware utilization across ranks despite differing optimizer computational characteristics.
  • Communication Efficiency: Fused communication lowers optimizer latency from ≈0.11s to ≈0.073s, with performance stabilizing around Cmax = 512MB.The no-fuse baseline suffers from many small kernel launches and poor bandwidth utilization.
  • Initialization Overhead: Offline scheduling completes in milliseconds on Qwen3-32B with 256 GPUs, making its initialization overhead effectively indistinguishable from system noise.The reported work omits dedicated scheduling-latency experiments because the overhead is imperceptible.

D.4. Discussion on the Choice of Training Framework

Canzona is implemented primarily in Megatron because its tensor parallelism, pipelining, and geometric partitioning expose the targeted communication constraints. The framework’s abstractions can extend to FSDP, but hybrid sharding introduces distinct communication and engineering limitations.

  • Framework Choice: Megatron provides the primary evaluation framework because its optimized tensor parallelism, pipelining, and strict geometric partitioning expose Canzona’s target constraints.The framework is designed to preserve ZeRO-1 geometric alignment and coalesced communication overlap.
  • FSDP Generalization: LB-ASC is described as generalizable to FSDP because both FSDP and TP fragment parameters and reconstruct full tensors before matrix-based optimizer updates.Micro-Group Scheduling and asynchronous compute are presented as transferable to FSDP shards treated as logical units.
  • FSDP Generalization: FSDP adaptation depends on the ZeRO stage: FSDP-ZeRO1 can skip the TP pipeline’s All-to-All, whereas FSDP-ZeRO2/3 requires fused All-to-All reconstruction.The distinction follows whether gradients are replicated or sharded across devices.
  • Limitations: FSDP-ZeRO2/3 combined with TP requires a coupled 2D communication mesh, imposing slower inter-node DP synchronization and substantial topology-management complexity.This hybrid setting is a practical boundary for the framework’s otherwise transferable TP abstraction.
  • Cost Metric: Numel preserves optimizer-agnostic infrastructure across Muon, Shampoo, SOAP, and future algorithms without requiring optimizer-specific complexity models.The choice reflects Canzona’s separation of system architecture from optimizer internals.
  • Cost Metric: The numel proxy relies on strong shape–cost correlation in standard Transformer architectures and reduces engineering fragility from maintaining separate optimizer cost models.These assumptions support the simplified metric primarily for architectures such as Qwen.
  • Cost Metric: Using numel instead of exact FLOPs changes latency by approximately 10−4s, with 0.0718s for Numel versus 0.0717s for FLOPs.The experiment uses Qwen3-32B with DP=16 and TP=8, supporting numel as a practical proxy under the tested setting.
  • Positioning: Canzona is positioned as reconciling mathematical exactness and system efficiency, unlike layerwise partitioning, which preserves exactness but incurs communication penalties.The comparison distinguishes system-level compromises, algorithmic approximations, and Canzona’s proposed approach.

E.3. Algorithmic Attempts: Efficiency at the Cost of Fidelity

Algorithmic approximations reduce communication or memory costs, but they may sacrifice fidelity, convergence quality, or generality. Canzona instead claims exact optimizer behavior through system-level scheduling and a unified, optimizer-agnostic design.

  • Block-Diagonal Approximations: Block-diagonal methods approximate full preconditioners by ignoring off-diagonal correlations, which may degrade convergence speed or solution quality for LLMs.The limitation is tied to settings where dense correlations matter.
  • Shard-Local Orthogonalization: Shard-local orthogonalization replaces global updates with local approximations, creating directional drift that can cause instability or require synchronization.MuonBP is cited as an example of this approach.
  • Low-Rank and Subspace Approximations: Low-rank and subspace methods improve efficiency but assume updates lie in a low-rank subspace, creating a fidelity trade-off when full-rank curvature matters.These methods alter the optimization path rather than preserving the exact holistic update.
  • Overall Limitation: Algorithmic approximations may require further verification at scale and often use optimizer-specific modifications that restrict extensibility to other matrix-based optimizers.The paper contrasts this with a unified system-level design.
  • Canzona’s System-Level Alternative: Canzona claims zero-fidelity-loss by preserving the original optimizer definition while handling atomicity through memory layout and scheduling.The claimed convergence behavior is identical to single-device baselines.
  • Canzona’s System-Level Alternative: Canzona combines balanced DP partitioning, efficient ZeRO-1 communication overlap, and optimizer-agnostic tensor-task scheduling instead of altering optimizer mathematics.The framework is described as supporting Muon, Shampoo, SOAP, and future matrix-based algorithms without algorithm-specific hacks.
Loading 2602.06079v1…