ArXiv: 2602.06079

🎯 Pitch

Matrix-based optimizers like Muon demand atomic, whole-tensor access for SVD and Newton-Schulz updates, yet the tensor fragmentation in standard distributed frameworks like Megatron blocks their deployment at scale. Canzona eliminates this conflict without sacrificing speed—achieving a 5.8× optimizer-step speedup by asynchronously hiding all reconstruction overhead using a novel micro-group scheduling pipeline.

1. Executive Summary

This paper introduces Canzona, a unified distributed framework that resolves the fundamental conflict between matrix-based optimizers—which require holistic tensor access for operations like SVD or Newton-Schulz iterations—and the tensor fragmentation imposed by standard ZeRO-1 and Tensor Parallelism sharding strategies. Evaluated on the Qwen3 model family (up to 32B parameters) on 256 GPUs, Canzona decouples logical optimizer assignment from physical parameter distribution through two complementary mechanisms: an α-Balanced Static Partitioning strategy for Data Parallelism (which enforces parameter atomicity by assigning whole tensors to ranks while using a blended allocation policy to neutralize load imbalance) and an Asynchronous Compute pipeline with Micro-Group Scheduling for Tensor Parallelism (which batches fragmented gradient updates into fused All-to-All collectives and schedules them onto load-balanced host ranks to hide reconstruction overhead). The framework achieves a 1.57× speedup in end-to-end iteration time and reduces optimizer step latency by 5.8× compared to NVIDIA's layerwise_optimizer baseline, while preserving the efficient Reduce-Scatter communication primitives that layer-wise partitioning structurally abandons—establishing that matrix-based optimizers can be deployed at scale without sacrificing either mathematical exactness or system throughput.

2. Context and Motivation

The Core Problem: Matrix-Based Optimizers Break Distributed Training Assumptions

The fundamental problem this paper addresses is a system-algorithm conflict that arises when trying to deploy matrix-based optimizers (Shampoo, Muon, SOAP) in modern distributed training frameworks like Megatron. These optimizers offer superior convergence efficiency compared to element-wise alternatives like AdamW by leveraging second-order information or structural properties of weight matrices—for instance, Muon applies Newton-Schulz iterations to orthogonalize weight matrices, while Shampoo constructs Kronecker-factored preconditioners. However, they impose a mathematical requirement that standard distributed training strategies were never designed to satisfy: the Atomicity Constraint.

The constraint is deceptively simple. Matrix-based optimizers operate at tensor granularity. For a weight matrix W∈Rdin×doutW \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}, the update rule involves holistic operations like SVD or matrix multiplication that require access to the complete tensor dimensions:

Wt+1=Wt−η⋅MatrixOp(Wt,∇Wt)W_{t+1} = W_t - \eta \cdot \text{MatrixOp}(W_t, \nabla W_t)

In contrast, element-wise optimizers like AdamW update each parameter independently, meaning the ii-th element's update depends only on its own gradient and historical states—never on the jj-th element. This makes element-wise optimizers trivially compatible with arbitrary tensor slicing: you can fragment a weight matrix however you want, and every fragment can be updated locally without communication.

The problem emerges because modern distributed training frameworks aggressively fragment tensors to minimize memory. Section 2 and Appendix B detail the two fragmentation mechanisms at play:

  • ZeRO-1 sharding (Data Parallelism): To reduce per-GPU memory consumption (optimizer states alone can consume 2–3× the model size in mixed-precision training), ZeRO-1 partitions optimizer states across data-parallel ranks. Megatron's implementation flattens all parameters into a contiguous param_and_grad_buffer, logically divides it into "buckets," and then rigidly slices each bucket into RR equal contiguous segments—entirely agnostic to individual tensor boundaries. Rank rr owns the rr-th geometric slice. For AdamW, this is perfectly fine. For Muon, a weight matrix that straddles the cut between rank 2 and rank 3 is now inoperable locally: neither rank has the full tensor, and no rank can perform the required orthogonalization without first reconstructing the tensor via communication.

  • Tensor Parallelism (TP): In Megatron's TP scheme, individual weight matrices are split across devices. Column-parallel linear layers split the output dimension; row-parallel layers split the input dimension. Each rank holds only a shard of the full weight. When the optimizer step requires holistic access (e.g., to compute preconditioners from the full gradient statistics), every rank must reconstruct the complete tensor—again requiring communication.

This conflict is not a minor implementation inconvenience. It is structural: the memory-saving mechanism that made training trillion-parameter models feasible (sharding) is the exact same mechanism that makes matrix-based optimizers locally inexecutable. Resolving this without sacrificing either memory efficiency or mathematical correctness is the central challenge.

Why This Problem Matters: The Convergence-Efficiency Tradeoff

The significance of this problem extends beyond a niche systems concern. Three factors make it urgent and practically consequential:

1. Matrix-based optimizers are not a curiosity—they are becoming mainstream. Recent work (Liu et al., 2025; Wen et al., 2025) has demonstrated that optimizers like Muon can substantially accelerate LLM pretraining convergence compared to AdamW, reducing the total number of training steps required to reach a given loss. In the era of multi-million-dollar training runs, algorithmic efficiency gains translate directly to cost savings and faster research iteration. However, these gains remain inaccessible to practitioners using standard distributed training stacks if the system overhead of deployment outweighs the algorithmic benefit. The paper's own experiments (Section 5.2, Figure 4) quantify this: NVIDIA's layerwise_optimizer, the current best-practice for preserving atomicity, spends 0.383 seconds per optimizer step for Qwen3-32B with Muon on 256 GPUs—a latency that would dominate end-to-end training time and potentially negate any convergence advantages.

2. The gap between algorithmic innovation and system support is widening. The paper identifies a growing divergence: algorithmic researchers continue to propose more sophisticated optimizers (Shampoo, SOAP, Muon, Conda, ROOT, PSGD, Sophia, and their variants), while distributed training frameworks remain optimized exclusively for element-wise algorithms. Without a general, optimizer-agnostic system solution, each new optimizer requires bespoke engineering to deploy at scale—a pattern the paper explicitly critiques in Appendix E.3, noting that existing approaches like MuonBP resort to "ad-hoc, algorithm-specific modifications" that "lack a Unified design philosophy." Canzona positions itself as a unified framework precisely to bridge this gap: it treats tensor updates as generic computational tasks defined by cost metrics, making it compatible with any current or future matrix-based optimizer without requiring per-algorithm system engineering.

3. There is a genuine tension between mathematical fidelity and system throughput that prior solutions force practitioners to trade off. The paper identifies two categories of prior approaches (Appendix E), each sacrificing one axis:

  • System-level approaches (like layerwise_optimizer) preserve mathematical exactness—they compute the true, full-tensor optimizer update—but incur severe communication overheads. As detailed in Section 3.1 and Appendix D.2, these approaches create a ZeRO Geometric Incompatibility: because layer-wise assignment disregards the physical parameter layout, the system cannot use the efficient bucket-based Reduce-Scatter primitive for gradient synchronization. It must fall back to All-Reduce, which incurs 2× the communication volume, and often requires additional Broadcast or All-Gather operations during the optimizer step to redistribute updated parameters. Appendix C.2 (Figure 7) directly quantifies this: the Fwd-Bwd latency of NV-layerwise matches the AdamW All-Reduce baseline, confirming it is bottlenecked by this 2× communication penalty.

  • Algorithmic approaches preserve system efficiency by modifying the optimizer itself—using block-diagonal approximations (Distributed Shampoo, K-FAC), shard-local orthogonalization (MuonBP), or low-rank subspace projections (Dion). These avoid communication by relaxing the atomicity constraint, but at the cost of mathematical fidelity: the "Local Newton-Schulz" update in MuonBP differs from the "Global Newton-Schulz" update, introducing directional drift between the true gradient geometry and the applied update. As the paper notes, this "may degrade convergence speed or solution quality for Large Language Models where dense correlations matter."

This tradeoff is unacceptable for large-scale training, where both convergence quality and throughput are critical. The paper's core contribution—the third category in Appendix E.4—is demonstrating that this tradeoff is false: with careful decoupling of logical assignment from physical distribution, it is possible to achieve both mathematical exactness and high throughput simultaneously.

Where Prior Approaches Fall Short

The paper provides a detailed taxonomy of prior approaches in Appendix E, but the key deficiencies can be understood through three lenses:

The Layer-wise Partitioning Trap (System-Level Exactness, System-Level Inefficiency). NVIDIA's layerwise_optimizer (an open-source PR to Megatron-LM) represents the most direct attempt to deploy matrix-based optimizers while preserving exactness. It assigns optimizer states at the granularity of whole layers rather than individual parameters, ensuring that no single tensor is split across ranks. Section 3.1 and Appendix D.2 provide a careful analysis of why this approach is fundamentally limited. The core issue is the ZeRO Geometric Incompatibility visualized in Figure 15. Megatron's communication primitives rely on a strict geometric partition: for a bucket of size BB across RR ranks, the rr-th contiguous segment is always sent to Rank rr. The destination is determined solely by physical position. Layer-wise assignment, however, allocates parameters to ranks based on computational load, creating a "data-task mismatch" where parameter P3P_3 physically located in what is geometrically Rank 2's region might be assigned to Rank 1. This interleaving of ownership breaks bucket coalescing: the system cannot launch a single monolithic Reduce-Scatter kernel and is forced into the "lose-lose dilemma" described in Appendix D.2—either fall back to All-Reduce (2× bandwidth) or dismantle buckets into per-parameter communication kernels (high latency).

Beyond the gradient synchronization penalty, this geometric violation compounds during parameter update. Under standard ZeRO-1, updated parameters are gathered via bucket-based All-Gather overlapped with the forward pass. Because layer-wise ownership misaligns with geometric shards, a coalesced All-Gather is impossible. Implementations must either perform explicit All-Gather/Broadcast within the optimizer step itself (adding exposed latency, as Section 5.2 quantifies at 0.383 seconds for the optimizer step alone) or resort to inefficient per-parameter operations that destroy overlap efficiency. This is the structural reason why NV-layerwise shows a 1.23× slowdown in Fwd-Bwd time compared to Canzona (Figure 4).

Algorithmic Approximations as Band-Aids (System-Level Efficiency, Mathematical Fidelity Loss). The paper surveys several lines of work that modify optimizers to fit system constraints:

  • Block-diagonal approximations (Distributed Shampoo from Anil et al., 2020; Distributed K-FAC from Osawa et al., 2019) replace the full preconditioner with a block-diagonal matrix, effectively ignoring off-diagonal correlations. This allows parallel computation on each block. The limitation: for LLMs where transformer weight matrices can be 4096×4096 or larger, dense correlations across the full matrix may be important, and the small-block assumption can degrade convergence.

  • Shard-local orthogonalization (MuonBP from Khaled et al., 2025) performs Newton-Schulz iterations strictly on the local shard of each weight matrix, avoiding global communication entirely. The limitation is mathematical: the local update approximates—but does not equal—the global update. The paper characterizes this as "directional drift" between the true gradient geometry and the applied update, which can require frequent correction or lead to training instability.

  • Low-rank and subspace projections (Dion from Ahn & Xu, 2025; Gong et al., 2025) constrain updates to lower-dimensional manifolds. While efficient, they "impose strong structural assumptions" that "may degrade convergence quality in regimes where full-rank curvature information is critical."

The paper's critique of these approaches is not that they are ineffective—they demonstrably improve efficiency—but that they represent an architectural defeat: accepting that matrix-based optimizers cannot be deployed faithfully at scale, and compromising the algorithm to fit the system. Canzona rejects this premise.

The Lack of a Unified, Optimizer-Agnostic Framework. Beyond the specific deficiencies of individual approaches, the paper identifies a broader gap: there is no general system infrastructure for matrix-based optimizers. Each prior solution is optimizer-specific (MuonBP targets Muon; Distributed Shampoo targets Shampoo; layerwise_optimizer requires per-optimizer adaptation of the layer-to-rank mapping). This means each new optimizer proposal requires its own distributed-system implementation, creating a barrier to experimentation and deployment. Canzona's design goal is explicitly to fill this gap: by abstracting optimizer tasks as generic operations defined by cost metrics, the framework supports Muon, Shampoo, SOAP, and future optimizers without modification (Section 4.3, Appendix E.4). The experiments in Appendix C.4 validate this claim, showing that Canzona achieves comparable efficiency gains across all three optimizers and preserves exact convergence with the synchronous baseline.

How This Paper Positions Itself

Canzona positions itself at the intersection of two research communities—distributed systems and optimization algorithms—arguing that the conflict between them is resolvable through a decoupling principle. The key intellectual moves are:

1. Decoupling logical assignment from physical distribution. The paper frames all existing approaches as coupling these two concerns: layerwise_optimizer ties logical ownership to physical placement (layers), while ZeRO-1 ties physical placement to geometric constraints (position in the buffer). Canzona's insight is that these can be separated. For Data Parallelism (Section 3), the Static Partitioning strategy assigns whole parameters to ranks (respecting atomicity) by defining ownership intervals [si,r−1,si,r)[s_{i,r-1}, s_{i,r}) for each bucket BiB_i and rank rr, but instead of using equal chunks, these intervals are optimized to equalize load while staying within the sequential physical ordering. For Tensor Parallelism (Section 4), each TP-split parameter's update becomes an atomic "Compute Task" assigned to a Host Rank that owns the full optimizer states locally, with All-to-All communication reconstructing only the necessary gradients. Neither approach requires modifying the underlying communication primitives or the optimizer mathematics.

2. Formulating load balancing as a static optimization problem. The decoupling creates a new challenge: if you assign whole parameters to ranks rather than equal-sized shards, the computational load becomes severely imbalanced. A rank that happens to own a large embedding matrix may have 3.24× the FLOPs of the average (Figure 3c). The paper formulates this as two discrete optimization problems—a makespan-minimization with atomicity constraints for DP (Equation 2) and a hierarchical bin-packing with embedded multiprocessor scheduling for TP (Section 4.2)—and proposes efficient offline heuristics (the α-Balanced Greedy LPT algorithm and the Micro-Group Greedy Rollback algorithm). Crucially, these are one-time offline computations that complete in milliseconds (Appendix D.1), not runtime overhead.

3. Preserving—rather than abandoning—ZeRO-1 geometric primitives. The paper's sharpest critique of layerwise_optimizer is that it "structurally forces" a fallback to suboptimal communication (All-Reduce instead of Reduce-Scatter). The geometric alignment is what enables the Forward-Backward pass to efficiently overlap communication with computation. Canzona's α-Balanced Partitioning preserves this alignment explicitly: by anchoring ownership to Start_Index (Equation 1), the physical parameter ordering remains monotonic and sequential, meaning the standard bucket-based Reduce-Scatter and All-Gather primitives remain usable. The paper validates this claim in Appendix C.2 (Figure 7), showing that Canzona's Fwd-Bwd latency tracks the AdamW Reduce-Scatter baseline rather than the All-Reduce baseline.

4. Establishing a unified design philosophy across parallelism dimensions. The paper explicitly connects the DP and TP strategies (Section 4.3): both adopt a fully asynchronous execution model where expensive TensorOps execute in parallel across ranks, both rely on load-balanced static planning to minimize stragglers, and both preserve the optimizer-agnostic abstraction that treats updates as generic tasks. This unified design is what distinguishes Canzona from point solutions—it is a framework, not a collection of optimizations.

5. Emphasizing zero fidelity loss. Unlike algorithmic approximations that modify the optimizer's update rule, Canzona's system-level approach preserves the exact mathematical definition of the optimizer. The paper validates this through precision experiments (Section 5.3, Appendix C.4): training loss trajectories for Muon, Shampoo, and SOAP are indistinguishable between the synchronous baseline and Canzona's LB-ASC strategy across 400B training tokens. This is a crucial selling point: practitioners do not need to trust that an approximation is "good enough" or re-validate convergence behavior—they get the same optimization trajectory with better throughput.

In summary, the paper fills a gap that was created by the simultaneous success of two research directions: increasingly sophisticated matrix-based optimizers that demand holistic tensor access, and increasingly aggressive tensor fragmentation strategies that make training large models memory-feasible. By demonstrating that these demands can be reconciled through careful system design rather than algorithmic compromises, Canzona opens the door for matrix-based optimizers to become the default—rather than the exception—in large-scale LLM training.

3. Technical Approach

This is primarily a systems paper whose core contribution is a unified distributed execution framework that decouples the logical assignment of optimizer update tasks from the physical sharding layout imposed by Data Parallelism (ZeRO-1) and Tensor Parallelism, enabling matrix-based optimizers (Muon, Shampoo, SOAP) to execute at scale without sacrificing either mathematical exactness or communication efficiency.

3.1 Reader Orientation

Canzona is a distributed execution framework that plugs into existing training stacks (Megatron) and reorganizes how optimizer update work is partitioned and scheduled across GPUs so that matrix-based optimizers—which mathematically require access to complete tensors rather than arbitrary fragments—can run efficiently without redundant computation or expensive communication. The core problem it solves is load imbalance: when you stop slicing tensors into equal-sized shards (which element-wise optimizers like AdamW tolerate but matrix-based optimizers cannot) and instead assign whole tensors to individual ranks to preserve atomicity, the resulting workload distribution is severely skewed because matrix operations have non-linear costs that depend on tensor shape. The "shape" of the solution is a two-level static planning stage—one for Data Parallelism and one for Tensor Parallelism—that runs once during initialization (in milliseconds) and produces a fixed, optimized assignment of optimizer tasks to ranks, which the runtime then executes asynchronously without further synchronization.

3.2 Big-Picture Architecture (Diagram in Words)

Canzona has five major components, grouped by when they operate (offline planning versus runtime execution) and which parallelism dimension they address:

1. α-Balanced Greedy LPT Partitioner (Offline, Data Parallelism): Takes the model's full parameter list (organized into ZeRO-1 buckets) and a configurable balance factor $\alpha$, and computes per-bucket slicing vectors $\mathbf{s}_i$ that define which contiguous intervals of each bucket each DP rank owns. These intervals respect parameter boundaries (atomicity) and physical ordering (ZeRO-1 geometric constraints). The algorithm processes buckets in descending order of total load (Longest Processing Time rule), blends two allocation objectives—filling accumulated deficits across ranks versus uniform communication-size partitioning—via $\alpha$, and discretizes the resulting continuous target onto valid parameter boundaries. Output: a Global Partition Map $\Pi$ that overrides the standard equal-chunk shard registration.

2. Micro-Group Construction with Greedy Rollback (Offline, Tensor Parallelism): Takes the full set of TP-split parameters, a cost model $\mathcal{W}(p)$, a maximum group capacity $C_{max}$, and the number of TP ranks $R$, and outputs a sequence of Micro Groups $\mathbb{M} = \{M_1, \dots, M_K\}$. Each group $M_k$ encodes which tensors should be fused into a single All-to-All communication operation and which host rank each tensor is assigned to for computation. The algorithm globally sorts tensors by cost (LPT), iteratively packs them into candidate groups, simulates the exact makespan using a Min-Heap scheduling solver, and triggers a rollback (finalizing the current group) when adding a tensor would violate the capacity constraint. Output: a static execution plan consumed by the runtime.

3. Runtime Static-Layout Enforcement (Data Parallelism): At runtime, the partition map $\Pi$ governs three variable-size collective operations—Reduce-Scatter (backward), optimizer step (local), and All-Gather (forward)—all using the non-uniform shard sizes $S_{i,r}$ defined by $\Pi$. Critically, the standard Megatron overlap pattern is preserved: Reduce-Scatter for bucket $B_i$ is overlapped with backward computation of $B_{i-1}$, and All-Gather for bucket $B_i$ is overlapped with forward computation of $B_{i-1}$. The optimizer step itself involves zero communication because every parameter's full optimizer states and gradients are locally resident on its assigned owner rank.

4. Asynchronous Compute Unit Lifecycle (Tensor Parallelism): For each Micro Group $M_k$ in the pre-computed sequence, the runtime executes a four-stage pipeline: (1) a fused All-to-All gathers gradient shards to their designated Host Ranks (which already store the full optimizer states locally), (2) Host Ranks asynchronously execute the matrix-based update computation using the gathered gradients and local states, (3) a fused All-to-All scatters the computed update shards $\Delta W$ back to the original parameter owners, and (4) each rank applies its received update shards to its local parameter shards. All-to-All fusion ensures the communication volume per group is large enough to saturate NVLink bandwidth.

5. Unified Framework Abstraction (Cross-Cutting): Both dimensions converge on a common design: optimizer updates are treated as generic computational tasks defined by cost metrics ($\mathcal{W}(p) = \text{numel}(p)$ by default), the assignment of tasks to ranks is determined by static offline optimization, and execution is fully asynchronous across ranks—enabling compute-compute overlap that reduces the optimizer-step makespan.

3.3 Roadmap for the Deep Dive

  • First, the Atomicity Constraint and the Design Paradigm Analysis (Section 3.1): Understanding why existing strategies fail requires a precise characterization of what matrix-based optimizers demand, how standard ZeRO-1 violates it, and why layer-wise partitioning creates a Geometric Incompatibility. This section justifies why Static Partitioning is the only viable paradigm and introduces the load-imbalance challenge that motivates the optimization problem.

  • Second, the DP Load-Balance Optimization (Section 3.2): This is the mathematical core—formulating the static partitioning as a discrete optimization over two competing objectives (global DP balance and per-bucket communication balance), bounding the solution by the ZeRO-1 Geometric Constraints, and presenting the α-Balanced Greedy LPT heuristic. We walk through the algorithm step-by-step, explaining what each variable physically represents and how α controls the blend between compute-balance and communication-uniformity.

  • Third, the DP Runtime Workflow (Section 3.3): Once the partition map is computed, we explain how it overrides the standard Megatron shard registration (offline planning) and how the variable-size Reduce-Scatter, zero-communication optimizer step, and variable-size All-Gather execute at runtime with overlapping.

  • Fourth, the TP Task Abstraction and Asynchronous Compute Unit (Section 4.1): We move to Tensor Parallelism and explain the fundamental difference from DP (intra-node bandwidth permits communication during the optimizer step), the lifecycle of a single Micro Group (All-to-All gather, async compute, All-to-All scatter, local update), and why fusing tensors into groups is necessary to saturate bandwidth while avoiding memory peaks.

  • Fifth, the TP Hierarchical Partitioning and Scheduling (Section 4.2): This section mirrors the DP optimization but with a different structure: bin-packing tensors into groups while simultaneously solving a multiprocessor scheduling problem within each group. We walk through the dual lexicographic objectives (minimize imbalance within groups, maximize group saturation), the Global LPT sort, the rollback mechanism, and the MinHeapSolver subroutine.

  • Sixth, the Unified Framework Synthesis (Section 4.3): We tie the DP and TP strategies together, showing how they share the same design philosophy (decoupled assignment, asynchronous execution, load-balanced static planning, optimizer-agnostic cost abstraction) and explain how Canzona achieves all three axes of the "trilemma" simultaneously.

3.4 Detailed, Sentence-Based Technical Breakdown

The Atomicity Constraint and Why Standard Approaches Fail

Before presenting Canzona's mechanisms, the paper carefully defines the constraint that creates the problem and analyzes three design paradigms for addressing it, explaining why only one is viable.

What matrix-based optimizers demand versus what element-wise optimizers tolerate. The paper formalizes the difference in Appendix B.3. Element-wise optimizers like AdamW update each parameter scalar independently:

wt+1(i)=Update(wt(i),gt(i),mt(i),vt(i))w_{t+1}^{(i)} = \text{Update}(w_t^{(i)}, g_t^{(i)}, m_t^{(i)}, v_t^{(i)})

where $w_{t+1}^{(i)}$ is the updated $i$-th parameter scalar at step $t+1$, $g_t^{(i)}$ is its gradient, and $m_t^{(i)}, v_t^{(i)}$ are its historical moment estimates.

What this equation means operationally: the computation for element $i$ only reads and writes element $i$'s values. It never touches element $j$. This means that if you slice a 4096×4096 weight matrix arbitrarily at row 2048, the first half can be updated by rank 1 and the second half by rank 2 without any communication between them—each element's update is self-contained.

In contrast, matrix-based optimizers operate holistically:

Wt+1=Wt−η⋅MatrixOp(Wt,∇Wt)W_{t+1} = W_t - \eta \cdot \text{MatrixOp}(W_t, \nabla W_t)

where $W_t \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$ is the full weight matrix, $\nabla W_t$ is its full gradient matrix, $\eta$ is the learning rate, and $\text{MatrixOp}$ represents an operation like SVD, matrix square root via Newton-Schulz iterations (Muon), or Kronecker-factored preconditioner construction (Shampoo/SOAP) that requires simultaneous access to all entries of $W_t$ and $\nabla W_t$.

What this equation means operationally: the $\text{MatrixOp}$ function internally performs operations where every output element depends on many input elements. For Muon's Newton-Schulz iteration, computing the orthogonalized matrix involves repeated matrix multiplications that mix all rows and columns of the gradient. You cannot compute it from shards independently—the partial results from two shards do not compose simply. If rank 1 holds the top half of the gradient and rank 2 holds the bottom half, neither can locally compute $\text{Muon}_{\text{step}}$ without first reconstructing the full gradient.

This creates the Atomicity Constraint: the optimizer must access the complete tensor dimensions to execute a valid update step. Any sharding strategy that fragments a single parameter tensor across multiple ranks violates this constraint, making local computation impossible without reconstruction communication.

Standard ZeRO-1 violates atomicity by design. Section 2 and Appendix B.1 describe how Megatron's ZeRO-1 implementation flattens all parameters into a contiguous param_and_grad_buffer, partitions this buffer into logical "buckets" for communication overlapping, and then slices each bucket of size $|B|$ into $R$ equal contiguous segments of size $|B|/R$. The slicing is purely position-based and entirely agnostic to where individual parameter tensors begin and end. As illustrated in Figure 1 (the gray dashed lines), a single parameter $P_2$ might be start in the middle of rank 0's geometric slice and end in rank 1's slice. Neither rank has the full tensor, and neither can perform the matrix-based update locally.

Paradigm 1: Synchronous/Redundant Compute (the DDP baseline). The most straightforward way to preserve atomicity is to abandon sharding entirely—replicate all optimizer states on every DP rank, as standard Distributed Data Parallelism (DDP) does. Every rank performs identical matrix-based operations on the full tensors. This is mathematically correct but "severely limits scalability" because it eliminates the 1/R memory reduction that makes ZeRO-1 essential for large models. In mixed-precision training, optimizer states alone consume 2–3× the model size (FP32 master weights plus momentum buffers), and replicating this across ranks is typically infeasible for billion-parameter models. The paper treats this as the SC (Synchronous Compute) baseline for comparison.

Paradigm 2: Layer-wise Partitioning (geometric incompatibility). A seemingly natural compromise is to assign optimizer states at the granularity of whole layers rather than individual parameters, ensuring atomicity while still partitioning work. This is the approach taken by NVIDIA's layerwise_optimizer. A Global Load Balancing strategy (Global LPT) assigns each layer's optimizer states to a specific DP rank based on computational cost, placing heavy layers on underloaded ranks.

The paper's critical analysis (Section 3.1, expanded in Appendix D.2) is that this approach introduces a ZeRO Geometric Incompatibility that structurally prevents using the most efficient communication primitives. The problem is a mismatch between two assignment logics:

  • ZeRO-1 communication primitives are position-based. For a contiguous bucket of size $B$ across $R$ ranks, the Reduce-Scatter primitive rigidly sends the $r$-th geometric slice (the interval $[\frac{r-1}{R}B, \frac{r}{R}B)$) to Rank $r$. The destination is determined solely by the data's physical position in the buffer. This geometric alignment allows launching a single coalesced kernel for the entire bucket, achieving high bandwidth utilization.

  • Layer-wise assignment is weight-based. The task allocator assigns parameter $P_3$ to DP Rank 1 because Rank 1 is currently underloaded, even though $P_3$'s physical position in the buffer lies within a region that geometrically belongs to DP Rank 2. As visualized in Figure 15 (Right), this creates an interleaved ownership pattern: $P_1$ (DP1), $P_2$ (DP2), $P_3$ (DP1)—the assignment sequence no longer follows the monotonic physical ordering.

The consequence (the "Lose-Lose Dilemma" in Appendix D.2): Because the Reduce-Scatter primitive assumes destination rank is determined by geometric position, and the layer-wise assignment dictates destinations based on load, the system cannot combine both constraints into a single coalesced bucket-based communication. It faces two bad options. Option A (Forced All-Reduce): Abandon Reduce-Scatter and use All-Reduce, which broadcasts full gradients to all ranks, ensuring everyone receives the needed data regardless of assignment. This incurs 2× the communication volume and is the default choice for current layerwise_optimizer implementations. Option B (Forced Per-Parameter Communication): Try to enforce Reduce-Scatter but dismantle the bucket into discrete per-parameter kernels, each routing a specific tensor to its assigned rank. This avoids the 2× volume penalty but introduces massive kernel launch overhead and prevents saturating the interconnect bandwidth, shifting the bottleneck from bandwidth to latency.

This geometric violation also compromises the optimizer step itself. Under standard ZeRO-1, updated parameters are gathered via bucket-based All-Gather overlapped with the forward pass. Because layer-wise ownership misaligns with geometric shards, a coalesced All-Gather is impossible. Implementations must either perform explicit Broadcasts within the optimizer step (adding exposed latency) or resort to inefficient per-parameter operations that destroy overlap efficiency. The paper quantifies this overhead in Figure 4: NV-layerwise spends 0.383 seconds in the optimizer step versus Canzona's 0.066 seconds—a 5.8× difference—and additionally suffers 1.23× slower Fwd-Bwd time (0.998s vs. 0.811s) due to the forced All-Reduce.

Paradigm 3: Static Partitioning (Canzona's approach). The paper's key design decision is to enforce an Atomic Ownership Rule under ZeRO-1 Geometric Constraints. Instead of equal chunks, each rank's ownership interval within a bucket is defined by parameter boundaries that fall within its geometric slice. Formally, given a "stride" size $S = |B|/R$ (total buffer size divided by number of ranks, the width of each geometric slice), a parameter $p$ is assigned exclusively to rank $r$ based on its starting position in the flattened buffer:

(r−1)⋅S≤Start_Index(p)<r⋅S(r-1) \cdot S \leq \text{Start\_Index}(p) < r \cdot S

where $\text{Start\_Index}(p)$ is the byte offset of parameter $p$'s first element in the flattened param_and_grad_buffer, and $r \in \{1, \dots, R\}$ is the rank index.

What this rule computes: it assigns each parameter to the rank whose geometric slice contains the parameter's starting position. If a parameter starts at byte offset 1,048,577 and the stride $S$ is 1,048,576 bytes, the parameter falls in rank 1's slice (since $1 \times S \leq 1,048,577 < 2 \times S$). Crucially, the entire parameter—including any bytes that extend beyond $2 \times S$—is assigned to rank 1. The rule guarantees that every parameter is owned by exactly one rank (no fragmentation) and that ownership is determined by physical position (preserving geometric alignment).

Why this form matters: by anchoring ownership to $\text{Start\_Index}$, the assignment respects the monotonic sequential ordering of parameters in the buffer. Within each bucket, the sequence of parameter ownerships follows rank order—parameters starting in rank 0's slice are owned by rank 0, then parameters starting in rank 1's slice by rank 1, and so on. This means the system can still launch a single coalesced Reduce-Scatter for the entire bucket (the receivers are geometrically aligned), but the shard sizes are non-uniform: rank $r$ receives $S_{i,r}$ bytes rather than exactly $S$ bytes, where $S_{i,r}$ is the total size of all parameters whose start positions fall within rank $r$'s geometric slice. The paper refers to this as "boundary-shifting"—adjusting the cut points within each bucket to align with parameter boundaries while preserving sequential ordering.

The immediate problem this creates: load imbalance. While static partitioning solves the geometric compatibility problem, it introduces severe load heterogeneity. Figure 3c quantifies this: naive static partitioning (without load balancing) produces a FLOPs imbalance ratio (Max/Avg) of 3.24× and a Memory imbalance ratio of 2.46×. A rank assigned a large embedding matrix or a wide attention projection has dramatically more work than a rank holding only small bias terms and layer norms. This creates computational stragglers—all ranks must wait for the slowest rank to finish its optimizer step—and pipeline bubbles (highlighted by the dashed box in Figure 1). The rest of Section 3 is dedicated to solving this imbalance while preserving the geometric alignment.

The DP Load-Balance Optimization Problem

Section 3.2 formalizes the static partitioning as a discrete optimization problem over slicing vectors, with two competing objectives and a crucial geometric constraint.

Problem setup and notation. We have $R$ ranks and a param_and_grad_buffer divided into $N$ logical buckets $\mathcal{B} = \{B_1, \dots, B_N\}$. Each bucket $B_i$ consists of an ordered sequence of non-splittable parameters $P_i$ (i.e., parameters are atomic—you cannot assign half a parameter to one rank and half to another). We define a load function $\mathcal{W}(p)$ for each parameter, which in the paper's practical implementation is $\mathcal{W}(p) = \text{numel}(p)$ (number of elements), serving as a proxy for both memory footprint and computational cost. The paper notes in Appendix D.5 that while the framework mathematically supports arbitrary non-linear cost functions to capture the cubic complexity of matrix operations, numel is used because (1) it preserves optimizer-agnostic universality (the system doesn't need to know whether an optimizer is $\mathcal{O}(N)$ or $\mathcal{O}(N^3)$), (2) for standard transformer architectures, shape correlates strongly with actual computational cost, and (3) a controlled experiment (Figure 16) shows the performance difference between numel-based and exact FLOPs-based scheduling is negligible ($\approx 10^{-4}$ seconds).

Let $\mathcal{W}^i = \sum_{p \in B_i} \mathcal{W}(p)$ be the total load of bucket $i$, and let $\Phi_i(u) = \sum_{p \in B_i[0:u]} \mathcal{W}(p)$ be the cumulative load up to cut point $u$ within bucket $i$ (where $u$ indexes the sequence of parameters in $B_i$).

The goal is to determine slicing vectors $\mathbf{s}_i = (s_{i,0}, s_{i,1}, \dots, s_{i,R})$ for each bucket $B_i$, where $s_{i,0} = 0$, $s_{i,R} = \text{Size}(B_i)$, and $s_{i,r-1} < s_{i,r}$ for all $r$. The interval $[s_{i,r-1}, s_{i,r})$ defines the parameters owned by rank $r$ within bucket $i$. The cumulative load assigned to rank $r$ from bucket $B_i$ is:

Li,r=∑p∈range(si,r−1,si,r)W(p)L_{i,r} = \sum_{p \in \text{range}(s_{i,r-1}, s_{i,r})} \mathcal{W}(p)

where the range includes all parameters between the cut points $s_{i,r-1}$ and $s_{i,r}$ in the bucket's sequential ordering.

Objective 1: Global DP Balance (minimize stragglers). The primary optimization target is to minimize the maximum deviation from the ideal mean load across all ranks:

JDP=max⁡r∣∑i=1NLi,r−μload∣\mathcal{J}_{\text{DP}} = \max_{r} \left| \sum_{i=1}^{N} L_{i,r} - \mu_{\text{load}} \right|

where $\mu_{\text{load}} = \frac{1}{R} \sum_{i} \mathcal{W}^i$ is the ideal average load per rank (total work divided by number of ranks), and $\sum_{i=1}^{N} L_{i,r}$ is the total load (across all buckets) assigned to rank $r$.

What this equation computes: for each rank $r$, it computes the absolute difference between that rank's total assigned load and the ideal average load, then takes the maximum over all ranks. This is the worst-case deviation from perfect balance—the amount by which the slowest rank exceeds the ideal. Minimizing $\mathcal{J}_{\text{DP}}$ directly attacks the straggler problem: if every rank has load exactly $\mu_{\text{load}}$, then $\mathcal{J}_{\text{DP}} = 0$ and no rank waits for any other.

Objective 2: Fwd-Bwd Bucket Communication Balance (minimize idle time). A secondary objective ensures that within each individual bucket, the data volume is distributed as evenly as possible:

JComm=∑i=1N∑r=1R∣Si,r−∣Bi∣R∣\mathcal{J}_{\text{Comm}} = \sum_{i=1}^{N} \sum_{r=1}^{R} \left| S_{i,r} - \frac{|B_i|}{R} \right|

where $S_{i,r}$ is the actual data volume (in bytes) of the shard assigned to rank $r$ in bucket $i$, and $|B_i|/R$ is the ideal uniform shard size for that bucket.

What this equation computes: it sums the absolute deviations from uniform shard sizes across all buckets and all ranks. When shard sizes are highly non-uniform, the Reduce-Scatter and All-Gather collectives experience imbalance—some ranks finish their communication phase earlier than others, creating idle time. This idle time is partially hidden by overlapping with forward and backward computation, but excessive imbalance can outlast the overlap window and become exposed. The paper's ablation study (Appendix C.5, Figure 13) reveals that this objective is secondary: as $\alpha$ increases from 0 (pure communication balance) to 1 (pure DP compute balance), the optimizer step time decreases monotonically, while Fwd-Bwd time remains relatively stable, confirming that communication imbalance is largely absorbed by the overlap window.

The ZeRO-1 Geometric Constraint. Both objectives are optimized subject to a hard constraint: the slicing vectors $\mathbf{s}_i$ must respect the sequential, monotonic parameter ordering within each bucket. Parameters cannot be physically reordered—you can only shift the cut boundaries $s_{i,r}$ to align with parameter boundaries. This constraint is what distinguishes Canzona from layer-wise partitioning: by keeping the physical ordering intact and only adjusting where the cuts fall, the system preserves the ability to launch coalesced Reduce-Scatter and All-Gather primitives. The "boundary-shifting" mechanism is the mathematical bridge between load balancing and communication efficiency.

The paper notes that this discrete partitioning problem under strict atomicity constraints is NP-hard (it reduces to a variant of multiprocessor scheduling with contiguous allocation constraints), motivating the heuristic approach in Algorithm 1.

The α-Balanced Greedy LPT Algorithm (Algorithm 1)

Algorithm 1 processes buckets one at a time in descending order of total load (Longest Processing Time rule), constructing a blended target allocation for the current bucket and then discretizing it onto valid parameter boundaries.

Why LPT ordering matters. The algorithm sorts buckets virtually (without physically reordering them—the sort is only for decision order) such that $\mathcal{W}^{\pi(1)} \geq \mathcal{W}^{\pi(2)} \geq \dots \geq \mathcal{W}^{\pi(N)}$, where $\pi$ is a permutation. This means the heaviest bucket (the one containing the largest parameters) is processed first, when the "deficit space" is largest and most flexible. If you process heavy buckets last, you may find that all ranks are already near capacity and the heavy bucket cannot be balanced—creating a straggler that no later adjustment can fix. This is the standard LPT insight from scheduling theory (Graham, 1969).

Step 1: Deficit calculation. At the start of processing bucket $k = \pi(t)$ (the $t$-th bucket in LPT order), the algorithm has an accumulated load vector $\mathbf{L} = (L_1, \dots, L_R)$ where $L_r$ is the total load assigned to rank $r$ from all previously processed buckets. The target mean is $\mu = (\sum_i \mathcal{W}^i) / R$. For each rank, the deficit $d_r$ is how far that rank is below the mean:

dr=max⁡(0,μ−Lr)d_r = \max(0, \mu - L_r)

What this computes: if rank $r$ has already accumulated load equal to or above $\mu$, its deficit is 0 (it doesn't need more work). If it has accumulated $L_r < \mu$, the deficit $d_r = \mu - L_r$ is how much additional load it would ideally receive to reach the mean. The total deficit $D_{\text{total}} = \sum_r d_r$ is the total amount of "catch-up" work to be allocated. Ranks with $d_r > 0$ are currently underloaded; ranks with $d_r = 0$ are at or above the mean.

Step 2–3: Blended target allocation. The algorithm computes two basis allocation vectors:

  • The even basis $\mathbf{v}_{\text{even}} = [1/R, \dots, 1/R]$ — every rank gets an equal fraction of the current bucket's load. This corresponds to standard ZeRO-1 behavior and prioritizes within-bucket communication balance.
  • The fill basis $\mathbf{v}_{\text{fill}} = \mathbf{d} / D_{\text{total}}$ (when $D_{\text{total}} > 0$; otherwise fall back to $\mathbf{v}_{\text{even}}$) — each rank gets a fraction proportional to its deficit. Underloaded ranks get more of the current bucket to help them catch up.

These are blended via the parameter $\alpha \in [0, 1]$:

v∗=(1−α)veven+αvfill\mathbf{v}^* = (1 - \alpha) \mathbf{v}_{\text{even}} + \alpha \mathbf{v}_{\text{fill}}

What this computes: $\mathbf{v}^*$ is a probability vector (entries sum to 1) that defines the ideal fraction of the current bucket's load that should go to each rank. When $\alpha = 0$, $\mathbf{v}^* = \mathbf{v}_{\text{even}}$ — the algorithm ignores accumulated deficits and partitions the bucket uniformly (prioritizing communication balance). When $\alpha = 1$, $\mathbf{v}^* = \mathbf{v}_{\text{fill}}$ — the algorithm aggressively fills deficits, giving more of the current bucket to ranks that are behind (prioritizing DP compute balance). At intermediate $\alpha$, the algorithm strikes a compromise.

The target allocation in absolute load units is $\mathbf{target\_alloc} = \mathcal{W}^k \cdot \mathbf{v}^*$, where $\mathcal{W}^k$ is the total load of the current bucket.

Why this blending is necessary. Without the fill basis ($\alpha = 0$), every bucket is partitioned uniformly regardless of history, and the global load imbalance from past buckets never gets corrected. Without the even basis ($\alpha = 1$), the algorithm might make individual bucket shards extremely non-uniform to compensate for past imbalances, potentially creating communication bottlenecks that outlast the overlap window. The $\alpha$ parameter provides a tunable tradeoff. Appendix C.5 (Figure 13) empirically evaluates $\alpha = 0.0$ through $\alpha = 1.0$ and finds that $\alpha = 1.0$ yields the best total step time because computational stragglers are the dominant bottleneck and communication imbalance is effectively hidden by overlap.

Step 4: Atomic discretization. The continuous target allocation $\mathbf{target\_alloc}$ specifies how much load (in numel units) each rank should ideally receive from the current bucket, but the actual assignment must respect parameter boundaries—you cannot split a parameter to hit a target exactly. The algorithm iterates through ranks $r = 1, \dots, R-1$, maintaining a cumulative target $C$. For rank $r$, it accumulates $C \leftarrow C + \mathbf{target\_alloc}[r]$, then selects the cut point $s_{k,r}$ as the parameter boundary where the cumulative load function $\Phi_k(u)$ is closest to the target $C$:

sk,r=arg⁡min⁡u∈Uk∣Φk(u)−C∣s_{k,r} = \arg\min_{u \in \mathcal{U}_k} |\Phi_k(u) - C|

where $\mathcal{U}_k$ is the set of valid parameter boundaries in bucket $k$ (positions where one parameter ends and the next begins).

What this computes: it walks through the bucket's parameter sequence in order, accumulates the target load for each rank, and at each step chooses the parameter boundary that makes the actual assigned load as close as possible to the target. The $\arg\min$ over valid boundaries ensures atomicity—the cut is always placed between two parameters, never inside one. After selecting $s_{k,r}$, it updates the global load vector: $L_r \leftarrow L_r + (\Phi_k(s_{k,r}) - \Phi_k(s_{k,r-1}))$, where the difference $\Phi_k(s_{k,r}) - \Phi_k(s_{k,r-1})$ is the actual load of the slice assigned to rank $r$.

The final rank $R$ receives the remainder: $s_{k,R} = \text{Size}(B_k)$, and $L_R \leftarrow L_R + (\Phi_k(s_{k,R}) - \Phi_k(s_{k,R-1}))$.

Output and offline nature. The algorithm returns the set of slicing vectors $\{\mathbf{s}_i\}$ for all buckets, sorted back to their original physical indices (undoing the virtual LPT reordering). This constitutes the Global Partition Map $\Pi$. Appendix D.1 emphasizes that this entire computation is a one-time offline step during model initialization, completing in milliseconds (time complexity $\mathcal{O}(N \log N)$ where $N$ is the number of buckets, typically in the thousands). It introduces negligible overhead compared to CUDA context creation, memory allocation, and NCCL communicator setup.

Runtime Static-Layout Enforcement (DP Workflow)

Section 3.3 describes how the partition map $\Pi$ governs the runtime execution. The workflow has two stages: offline setup and per-iteration execution.

Offline planning phase. Before training begins, the α-Balanced Greedy LPT algorithm computes $\Pi$. Based on $\Pi$, the standard Megatron shard registration is overridden: instead of allocating uniform chunks of size $|B_i|/R$, each rank $r$ allocates memory proportional to its assigned shard sizes $S_{i,r}$ across all buckets. This ensures the physical buffer layout matches the logical partition. Importantly, the paper notes that the variable chunk sizes of each bucket introduce Reduce-Scatter and All-Gather communication imbalance, "but this is a secondary concern as it can be effectively hidden by overlapping communication with forward and backward computation" (Section 3.1, end of The Challenge subsection).

Per-iteration execution (three phases). The runtime preserves Megatron's standard overlapping pattern while handling variable-size shards:

  1. Backward pass (Variable-Size Reduce-Scatter): As gradients are computed and accumulated into bucket $B_i$, a non-uniform Reduce-Scatter is triggered. Unlike the standard implementation that assumes equal split sizes, this primitive handles the variable shard sizes $S_{i,r}$ defined by $\Pi$. The operation aggregates gradients across all DP ranks and scatters the result such that rank $r$ receives exactly the $S_{i,r}$ bytes corresponding to the parameters it owns. Crucially, this communication is overlapped with the backward computation of the previous bucket $B_{i-1}$, following Megatron's standard pipelining pattern. The paper notes that "the slight latency variance caused by the non-uniform data sizes is effectively hidden" by this overlap.

  2. Optimizer step (load-balanced asynchronous computation, zero communication): Rank $r$ updates the parameters strictly within its assigned intervals $[s_{i,r-1}, s_{i,r})$ for each bucket $B_i$. Because the partition map respects atomic boundaries, every parameter $p$ (with its full gradient and optimizer state) is fully available locally on its owner rank. The matrix-based optimizer executes without any additional communication—no All-Gather to reconstruct tensors, no Broadcast to distribute updated weights. This is the critical advantage over layer-wise partitioning, which requires explicit communication during the optimizer step to redistribute parameters. The "asynchronous" aspect means that different ranks can execute their updates in parallel without synchronization barriers, and the load-balanced assignment ensures they finish at approximately the same time (minimizing idle waiting).

  3. Forward pass (Variable-Size All-Gather): Before the forward computation for bucket $B_i$, the updated parameter shards must be gathered from all ranks. A non-uniform All-Gather reconstructs the full bucket parameters, with each rank contributing its $S_{i,r}$ bytes. Similar to the backward pass, this operation is overlapped with the forward computation of the previous bucket $B_{i-1}$. The paper's experiments (Figure 7) confirm that this overlapping is effective: Canzona's Fwd-Bwd latency closely tracks the ideal AdamW Reduce-Scatter baseline, while NV-layerwise tracks the All-Reduce baseline (reflecting its 2× communication volume penalty).

The TP Task Abstraction (from Reconstruction to Asynchronous Compute)

Sections 4.1 and 4.2 move to Tensor Parallelism, which presents a fundamentally different communication environment from Data Parallelism. While DP spans inter-node connections where any optimizer-step communication is prohibitively expensive (motivating the strict zero-communication strategy of Section 3), TP typically operates within the high-bandwidth intra-node domain (e.g., NVLink). This permits communication during the optimizer step, which Canzona exploits to eliminate the redundant computation of synchronous TP approaches.

The TP-SC (Synchronous Compute) baseline and its redundancy. In standard TP (Figure 2, Left), every rank holds only a shard of each weight matrix—a column-parallel shard for the first linear layer in a transformer block, a row-parallel shard for the second. When the optimizer requires holistic tensor access, the synchronous approach (as implemented in NVIDIA's layerwise_optimizer for TP) performs an All-Gather to reconstruct the full tensor on every rank, then every rank independently computes the identical matrix-based update, and finally each rank extracts its local shard from the result. This is the Redundant Compute highlighted in Figure 2—all ranks perform the same expensive operation, wasting $R-1$ ranks' worth of computation. For cubic-complexity operations like SVD or eigendecomposition, this redundancy is particularly costly.

Canzona's TP-ASC (Asynchronous Compute) approach. The paper's insight is that for TP, you can eliminate redundancy by designating a specific "Host Rank" for each tensor's update, having only that rank reconstruct the full tensor and perform the computation, and then scattering the update shards back. However, this requires careful orchestration: (1) the full optimizer states must be resident on the Host Rank so they don't need transmission, (2) only gradients need to be gathered (since the Host Rank already has the states), (3) the assignment of tensors to Host Ranks must be load-balanced so no single rank becomes a bottleneck, and (4) communication must be batched to saturate NVLink bandwidth.

Task abstraction as the organizing concept. Section 4.1 introduces the central abstraction: the update of each TP-split parameter is treated as an atomic "Compute Task" assigned to a specific Host Rank $r$. The assignment is static—determined offline during initialization—which allows a crucial optimization: the full optimizer states (e.g., Shampoo preconditioners, Muon momentum buffers) are initialized directly on their designated Host Ranks and never transmitted. This means the communication during the optimizer step only needs to move gradients (to the Host Rank) and update shards (back from the Host Rank), never the typically larger optimizer states.

The four-phase lifecycle of a Micro Group. The "Micro Gradient Group j" panel in Figure 2 (Right) illustrates the execution pipeline for a batch of tensors fused into a single communication group. Let the group $M_k$ contain tensors $\{G_1, G_2, G_3\}$ with their designated Host Ranks pre-assigned:

  1. All-to-All for gathering: Since the full optimizer states are already resident on the Host Ranks, the system only needs to aggregate the gradient shards. The asynchronous gather operations for all tensors in the group are fused into a single All-to-All collective. Each rank contributes its local gradient shards for every tensor in the group; the All-to-All routes each shard to the appropriate Host Rank. This fusion is essential: launching individual point-to-point transfers for each tensor would incur prohibitive kernel launch overhead and fail to saturate NVLink bandwidth. The paper's ablation in Appendix C.6 (Figure 14) quantifies this: the "No-Fuse" baseline (per-tensor communication) suffers from ~0.11 seconds optimizer time; enabling fusion immediately drops this to ~0.073 seconds, with further marginal improvements as the group size increases until bandwidth saturation around 512 MB.

  2. Asynchronous computation: Each Host Rank now has the full gradients for its assigned tensors (e.g., Host Rank 1 has the complete $\nabla G_2$) and the locally resident optimizer states. It executes the matrix-based update independently: $\Delta W = \text{Muon}_{\text{step}}(G, \text{States}_{\text{local}})$. Different Host Ranks compute different tensors in parallel—Host Rank 1 works on $G_2$ while Host Rank 2 works on $G_1$ and $G_3$. The load-balanced assignment ensures that the heaviest computation (e.g., $G_2$, represented by a longer block in Figure 2) is assigned to a dedicated rank that can handle it without stalling other tensors.

  3. All-to-All for scattering: Once the update tensors $\Delta W$ are computed, the Host Ranks slice them back into shards corresponding to the original TP partition. These update shards are scattered back to the parameter owners via a second fused All-to-All. Each rank receives the $\Delta W$ contributions for all parameters it owns shards of.

  4. Local update: Each rank applies the received update shards to its local parameter shards: $P_{i,\text{local}} \leftarrow P_{i,\text{local}} + \Delta W_{i,\text{local}}$. This is a purely element-wise operation requiring no further communication.

Why this sequence is communication-minimal. The key property is that only gradients and update shards are transmitted—both are the same size as the parameter shards themselves. The optimizer states, which for algorithms like Shampoo can be significantly larger than the parameters (preconditioner matrices scale quadratically with certain dimensions), never leave their Host Ranks. This is possible precisely because the assignment is static: the system knows at initialization time which rank will own each tensor's states, so it can place them once and leave them. This contrasts with dynamic scheduling, which would require transmitting states along with gradients.

The TP Hierarchical Partitioning and Scheduling (Micro-Group Construction)

While Figure 2 illustrates a single Micro Group executing efficiently, the system-level challenge is partitioning the entire model's TP-split parameters into a sequence of such groups and assigning each tensor to a Host Rank such that the overall execution time is minimized. Section 4.2 formulates this as a hierarchical optimization problem and presents Algorithm 2.

The dual nesting of optimization problems. The TP scheduling problem is naturally hierarchical: at the outer level, we must partition a set of tensors into groups (bin packing—which tensors go together in one All-to-All); at the inner level, within each group, we must assign tensors to Host Ranks to minimize the makespan (multiprocessor scheduling). These two levels interact: the total latency of processing one group is determined by the slowest Host Rank in that group, and a poor assignment within a group can make the entire group take longer, wasting communication slots that could have been used for other tensors.

Lexicographic objectives. The paper defines two priorities:

  • Priority 1: Minimize computational imbalance within each group. For a micro group $M_k$ with $R$ ranks, let $\mathcal{L}_{k,r}$ be the total computational load on rank $r$ from all tensors in $M_k$ assigned to that rank. The objective is:

min⁡Φ1(Mk)=max⁡r(Lk,r)−min⁡r(Lk,r)\min \Phi_1(M_k) = \max_r (\mathcal{L}_{k,r}) - \min_r (\mathcal{L}_{k,r})

where $\mathcal{L}_{k,r}$ is the sum of $\mathcal{W}(p)$ for all parameters $p \in M_k$ assigned to rank $r$.

What this computes: the difference between the maximum and minimum load across ranks within the group. A perfectly balanced group has $\Phi_1 = 0$. This is the makespan within the group—all ranks must finish before the next communication phase can begin, so the slowest rank determines the group's latency.

  • Priority 2: Maximize group saturation. Subject to the balance constraint, pack as much total load as possible into each group to minimize the number of communication rounds (each round incurs kernel launch and synchronization overhead):

max⁡Φ2(Mk)=∑r=1RLk,r,s.t. max⁡r(Lk,r)≤Cmax\max \Phi_2(M_k) = \sum_{r=1}^{R} \mathcal{L}_{k,r}, \quad \text{s.t. } \max_r (\mathcal{L}_{k,r}) \leq C_{max}

where $C_{max}$ is a capacity constraint representing either the memory buffer limit for the fused All-to-All or the maximum acceptable makespan for a single group.

What this computes: the total load in the group, subject to the constraint that no single rank's load exceeds $C_{max}$. This captures the tradeoff between communication efficiency (fewer, larger groups) and load balance (groups must stay within capacity).

The LPT with Greedy Rollback algorithm (Algorithm 2, detailed in Algorithm 3 in Appendix D.3). The algorithm comprises two nested phases:

Phase 1: Deterministic Global LPT Sort. All parameters $\mathcal{P}$ are sorted in descending order of their cost $\mathcal{W}(p)$ (with parameter ID as a secondary key for determinism across ranks):

Psorted←Sort(P,key=(W(p),IDp),descending)\mathcal{P}_{\text{sorted}} \leftarrow \text{Sort}(\mathcal{P}, \text{key}=(\mathcal{W}(p), \text{ID}_p), \text{descending})

Why global sorting matters. Processing the heaviest tensors first prevents them from accumulating at the end where they would be impossible to pack into balanced groups. This is the same LPT principle as in the DP algorithm, but now applied to individual tensors rather than buckets. The secondary sort by ID ensures all ranks independently compute the identical schedule without communication (crucial for a static, offline plan).

Phase 2: Greedy packing with rollback. The algorithm maintains a current candidate group $M_{\text{curr}}$ and iterates through $\mathcal{P}_{\text{sorted}}$. For each tensor $p$, it tentatively adds $p$ to $M_{\text{curr}}$ to form $M_{\text{test}} = M_{\text{curr}} \cup \{p\}$, then invokes a MinHeapSolver (Algorithm 4 in Appendix D.3) to simulate the optimal assignment of all tensors in $M_{\text{test}}$ to the $R$ Host Ranks.

The MinHeapSolver (inner multiprocessor scheduling). Algorithm 4 implements the standard LPT scheduling heuristic: it sorts the items in $M_{\text{test}}$ descending by cost (a local sort, separate from the global sort), then iteratively assigns each item to the currently least-loaded rank using a min-heap priority queue. Specifically, it initializes a min-heap $PQ$ with $(0, r)$ for each rank $r$ (load 0, rank index). For each tensor $(c, p, s)$ (cost, parameter, shape) in descending cost order, it pops the least-loaded rank $r_{\text{best}}$, assigns the tensor to that rank, and pushes back $(\text{load} + c, r_{\text{best}})$. The maximum load across all ranks after all assignments is $L_{\text{max}}$.

The rollback decision. If the solver reports $L_{\text{max}} \leq C_{\text{max}}$ (the makespan of the simulated assignment is within the capacity constraint), the tensor is accepted: $M_{\text{curr}} \leftarrow M_{\text{test}}$. If $L_{\text{max}} > C_{\text{max}}$, the constraint is violated. The algorithm then triggers a rollback: the current group $M_{\text{curr}}$ (without the violating tensor) is finalized and appended to $\mathbb{M}$, and a new group is started with $p$ as its first member ($M_{\text{curr}} \leftarrow \{p\}$). The iteration then continues from $p$ (the loop index is not incremented, so $p$ is retried as the seed of the new group).

Why the rollback mechanism is necessary. A simpler approach would be to estimate group load as $\sum \text{Cost} / R$ and accept a tensor if this average stays below $C_{\text{max}}$. However, this ignores fragmentation—depending on the specific cost distribution, a group with a modest total load might be impossible to balance across ranks (e.g., one very large tensor plus many small tensors where the large tensor dominates any one rank's load). The rollback mechanism uses the actual scheduling simulation as an oracle: it accepts a tensor only if a feasible assignment exists, and if not, it finalizes the group at the last feasible point. This ensures that every finalized group satisfies the balance constraint by construction.

Handling remaining tensors. After the loop, any tensors remaining in $M_{\text{curr}}$ are finalized as the last group.

Algorithmic complexity and offline nature. Although simulating the MinHeapSolver at every step appears computationally intensive, the solver operates in $\mathcal{O}(K \log R)$ for $K$ items in the current group. Since $K$ is typically small (on the order of tens of tensors per Micro Group), and the total number of tensors is in the thousands, the offline planning completes in milliseconds (Appendix D.1). The resulting sequence $\mathbb{M}$ serves as the static execution plan consumed by the runtime.

The Unified Framework Design Philosophy (Section 4.3)

Section 4.3 synthesizes the DP and TP strategies into a cohesive design philosophy that spans both parallelism dimensions. The paper identifies four unifying principles:

1. Unified abstraction. Neither the DP partitioner nor the TP scheduler requires any modification to the optimizer's internal mathematics. Tensor updates are treated as generic computational tasks defined solely by cost metrics ($\mathcal{W}(p) = \text{numel}(p)$). This makes the framework optimizer-agnostic—it supports Muon, Shampoo, SOAP, and any future matrix-based optimizer without requiring per-algorithm system engineering. The experiments in Appendix C.4 validate this: the identical framework codebase achieves comparable efficiency gains across all three optimizers, and training loss curves (Figures 5, 10b, 11b) match the synchronous baseline exactly, confirming zero fidelity loss.

2. Asynchronous execution across both dimensions. In DP, the optimizer step involves zero communication—each rank independently executes updates on its assigned parameters, and different ranks naturally proceed in parallel. In TP, the Micro Group pipeline executes the four-phase lifecycle (gather, compute, scatter, update) asynchronously: different Host Ranks compute different tensors concurrently, and the All-to-All communication phases are non-blocking where the hardware permits. The result is compute-compute overlap that substantially reduces the optimizer-step makespan compared to synchronous approaches where all ranks wait for the slowest operation.

3. Load-balanced static planning. Both dimensions rely on offline optimization (the α-Balanced Greedy LPT algorithm for DP, the Micro-Group Greedy Rollback algorithm for TP) to pre-compute assignment schedules that minimize stragglers. These schedules are computed once during initialization and impose negligible runtime overhead (the paper estimates milliseconds for Qwen3-32B on 256 GPUs, compared to minutes for CUDA context creation and NCCL setup).

4. Preservation of communication primitives. Critically, Canzona's design respects—rather than abandons—the efficient communication infrastructure of the underlying training framework. The DP Static Partitioning preserves the geometric alignment required for bucket-based Reduce-Scatter and All-Gather (validated in Figure 7). The TP Micro-Group batching saturates All-to-All bandwidth (validated in Figure 14). This distinguishes Canzona from both layer-wise partitioning (which structurally abandons ZeRO-1 geometric primitives) and algorithmic approximations (which modify the optimizer itself to avoid communication).

The paper frames these four principles as collectively resolving the "trilemma" between mathematical exactness (preserving the optimizer's original update rule), system throughput (minimizing end-to-end iteration time), and optimizer generality (supporting any matrix-based optimizer without modification). Prior approaches could achieve at most two of these simultaneously; Canzona is the first to achieve all three.

4. Key Insights and Innovations

Innovation 1: Decoupling as a System Design Principle, Not Merely an Optimization

The paper's most intellectually distinctive contribution is not any specific partitioning algorithm but the architectural diagnosis that the core conflict between matrix-based optimizers and distributed training is a coupling problem, and the corresponding insight that it can be resolved by decoupling logical optimizer-task ownership from physical parameter placement. This conceptual move redefines what a solution even looks like.

What the field assumed before this paper. The dominant approaches—both system-level (layer-wise partitioning) and algorithmic (block approximations, shard-local orthogonalization)—implicitly accepted that atomicity preservation and communication efficiency were in tension, forcing practitioners to trade one for the other. Layer-wise approaches said "preserve exactness, sacrifice communication geometry." Algorithmic approximations said "preserve efficiency, compromise mathematical fidelity." The underlying assumption was that the requirements of matrix-based optimizers (holistic tensor access) and the mechanisms of efficient distributed training (position-based geometric sharding) were fundamentally incompatible—that you could satisfy one constraint or the other, but not both simultaneously.

What this paper realized. The conflict is artificial. It arises because existing systems coupled the decision of which rank owns a parameter's optimizer states to two independent questions: (1) how the physical memory buffer is geometrically sliced for communication, and (2) how the computational work of updating parameters is distributed across ranks. In ZeRO-1, these are coupled by construction: rank $r$ owns the optimizer states for the $r$-th geometric slice, period. In layer-wise partitioning, they are coupled by a different logic: rank ownership is determined by load balancing, which then dictates what data must be communicated where. Canzona's key insight is that the ownership question and the communication question are separable: you can assign whole parameters to ranks based on load-balancing considerations (enforcing atomicity, eliminating redundant compute), while independently preserving the geometric slicing that enables efficient coalesced communication primitives—as long as you enforce a monotonic physical ordering constraint.

This decoupling is what enables the "boundary-shifting" mechanism (Section 3.2): by adjusting slice boundaries within each bucket to align with parameter edges while keeping parameters in their original physical order, Canzona simultaneously satisfies the Atomicity Constraint (whole parameters owned by single ranks) and the ZeRO-1 Geometric Constraint (the sequence of ownerships is monotonic, so coalesced Reduce-Scatter and All-Gather remain valid). The conceptual clarity of this separation—and the recognition that it resolves what appeared to be an irreconcilable tension—is the paper's deepest contribution. It transforms the problem from "which compromise do we accept?" to "how do we optimize the decoupled assignment?"

Evidence anchoring this insight. The paper's architecture diagrams (Figures 1 and 2) are not merely illustrations—they are arguments. Figure 1 contrasts three paradigms not as performance comparisons but as different coupling strategies: the gray dashed lines (Equal Chunk) show coupling-by-geometry, the layer-wise approach (discussed in text, visualized in Figure 15) shows coupling-by-load, and the orange/green arrows show decoupling (assign parameters atomically, but optimize boundaries for balance while maintaining geometric order). Figure 7 provides the empirical proof that this decoupling works: Canzona's Fwd-Bwd latency tracks the AdamW Reduce-Scatter baseline (the theoretical lower bound for communication-efficient ZeRO-1), while NV-layerwise tracks the All-Reduce baseline (confirming the 2× communication penalty from failing to decouple). The 5.8× optimizer-step speedup (Figure 4) is the downstream consequence of decoupling—but the insight that decoupling is the mechanism is more fundamental than the speedup itself.

Significance beyond performance. This is a fundamental conceptual reframing, not an incremental optimization. It provides a design template for future distributed optimizers: separate assignment logic from communication geometry, enforce a monotonic ordering constraint, and solve the resulting load-balancing problem offline. The paper explicitly acknowledges this generality (Section 4.3, Appendix D.4), noting that the same principle applies to FSDP and other sharding strategies—not just Megatron's ZeRO-1. The decoupling insight is what makes Canzona a framework rather than a collection of point optimizations.

Innovation 2: The ZeRO Geometric Incompatibility as a Diagnosed Structural Bottleneck

The paper introduces a new diagnostic concept—the ZeRO Geometric Incompatibility—that identifies why layer-wise partitioning approaches (like NVIDIA's layerwise_optimizer) are structurally incapable of matching the communication efficiency of standard ZeRO-1, regardless of implementation quality. This is not merely an observation about a specific baseline; it is a general negative result that explains the failure mode of an entire class of approaches and provides a criterion for evaluating future proposals.

What the field missed. Prior to this paper, layer-wise partitioning was the state-of-the-art for deploying matrix-based optimizers in Megatron while preserving mathematical exactness. The approach seemed natural: if you must assign whole parameters to ranks, assign them at layer granularity to keep the mapping manageable. The performance penalty of layer-wise approaches (higher iteration time) was known but attributed to general communication overhead—the implicit assumption being that any approach preserving atomicity would incur similar costs. What the paper's analysis reveals (Section 3.1, Appendix D.2) is that the penalty is not inherent to atomicity preservation at all—it is specific to the geometric misalignment created by layer-wise assignment logic.

The diagnostic move. The paper identifies the root cause as a mismatch between two assignment logics: ZeRO-1 communication primitives are position-based (the $r$-th geometric slice of a contiguous buffer goes to Rank $r$), while layer-wise allocation is weight-based (parameters are assigned to ranks based on computational cost, regardless of physical position). This mismatch creates a "data-task mismatch" where a parameter physically located in what is geometrically Rank 2's region might be assigned to Rank 1, interleaving ownership in a way that breaks bucket coalescing. The system is then forced into the "lose-lose dilemma" (Appendix D.2): either fall back to All-Reduce (2× communication volume) or dismantle buckets into per-parameter communication kernels (high latency). Neither option is a performance bug; both are structural consequences of the geometric incompatibility.

Why this is a conceptual advance, not just a performance critique. The Geometric Incompatibility concept provides a falsifiable criterion for evaluating future proposals: any approach that scrambles the monotonic rank-to-position mapping within ZeRO-1 buckets will necessarily incur this penalty—it cannot be engineered around with better overlapping or faster networks because the primitive itself (coalesced Reduce-Scatter) requires geometric alignment. This transforms the design space from "try various assignment heuristics and measure performance" to "verify that your assignment preserves monotonic physical ordering, or accept the 2× communication penalty." The paper uses this criterion to explain not just why NV-layerwise underperforms but also why Canzona avoids the penalty (its Start_Index-anchored ownership rule preserves monotonicity by construction).

Evidence. The Geometric Incompatibility is not merely asserted; it is empirically verified. Figure 7 shows that NV-layerwise's Fwd-Bwd latency precisely matches the AdamW All-Reduce baseline, while Canzona's matches the AdamW Reduce-Scatter baseline. This controlled comparison eliminates confounding factors—both systems are running the same model, same optimizer, same hardware, differing only in whether the assignment logic geometrically aligns with the buffer layout. The result is a clean demonstration that the 2× communication volume is structural, not incidental. Figure 15 visualizes the interleaving that creates the incompatibility, making the abstract concept concrete. The paper also notes (Appendix D.2) that certain buckets might coincidentally align (where layer-wise assignment happens to produce monotonic ownership), but this is luck, not a property of the approach—and the overall system must be designed for the worst case.

Significance. This is a fundamental diagnostic contribution that changes how system designers should think about distributed optimizer architectures. It establishes that the geometric constraint is not an implementation detail of Megatron but a first-class design requirement that any atomicity-preserving system must satisfy to achieve efficient communication. The concept will likely generalize to other frameworks (FSDP, DeepSpeed) that use position-based sharding of contiguous buffers.

Innovation 3: Load Balance as a Static Optimization Problem with Explicit Geometric Constraints

While load balancing in distributed systems is well-studied, the paper's formulation of it as a discrete optimization over slicing vectors subject to a hard geometric constraint (monotonic sequential ordering) and a tunable blend between two competing objectives (global DP balance vs. per-bucket communication balance) is novel in its specificity and its integration with the decoupling principle. The key intellectual move is recognizing that the load-imbalance problem created by atomic assignment is not an unfortunate side effect to be tolerated but a tractable offline optimization that can be solved once during initialization with negligible overhead.

What prior work did. Existing approaches to load balancing in distributed training—including layer-wise LPT assignment and general-purpose scheduling heuristics—typically treated load as an aggregate quantity to be equalized without considering the geometric constraints of communication primitives. Layer-wise approaches performed Global LPT to balance total load across ranks but ignored the physical parameter layout, creating the Geometric Incompatibility discussed above. Standard ZeRO-1 enforced geometric uniformity ($|B|/R$ equal chunks) but ignored load, trusting that the model architecture would naturally produce roughly balanced work. Neither approach recognized that load balancing and geometric alignment could be jointly optimized through boundary-shifting within a monotonic constraint.

The paper's formulation. Section 3.2 defines the problem over slicing vectors $\mathbf{s}_i$ for each bucket, with two objectives ($\mathcal{J}_{\text{DP}}$ for global balance, $\mathcal{J}_{\text{Comm}}$ for per-bucket communication balance) and a hard constraint that parameters remain in sequential physical order. The α-Balanced Greedy LPT algorithm (Algorithm 1) then provides a heuristic solution that interpolates between the two objectives via the $\alpha$ parameter. This formulation is distinctive in three ways:

  1. It encodes the geometric constraint directly into the optimization variable structure. The decision variables are cut points $s_{i,r}$, and the valid solution space is restricted to parameter boundaries. This means every solution the optimizer considers is guaranteed to be geometrically valid—the constraint is enforced by construction, not verified post-hoc.

  2. It separates objectives that are naturally in tension and provides a tunable tradeoff. The blend parameter $\alpha$ is not arbitrary—it corresponds to a physically meaningful choice: $\alpha \to 0$ recovers standard ZeRO-1 behavior (prioritize communication uniformity), while $\alpha \to 1$ aggressively corrects computational imbalance. The ablation in Appendix C.5 (Figure 13) validates that $\alpha = 1.0$ is optimal because communication imbalance is hidden by overlap, meaning the tradeoff collapses to a single objective in practice—but the formulation correctly models the possibility that this might not always be true (e.g., on hardware with less effective overlap).

  3. It is a one-time offline cost. Appendix D.1 emphasizes that the algorithm completes in milliseconds for real models on real hardware—the $\mathcal{O}(N \log N)$ complexity is dominated by the bucket sorting, not the discretization. This means the formulation is not merely theoretically elegant but practically deployable without runtime overhead.

Evidence. Figure 3c provides the direct evidence of the problem and solution: naive static partitioning produces a FLOPs imbalance ratio of 3.24× (Max/Avg) and a Memory imbalance ratio of 2.46×, while the α-Balanced strategy reduces these to 1.43× and 1.11× respectively. Figure 3a shows that this improved balance translates directly to reduced makespan (eliminated computation bubbles). Figure 13 (α sensitivity) validates that the tunable tradeoff collapses cleanly in practice: $\alpha = 1.0$ is optimal across the tested configurations, meaning the formulation's generality correctly anticipates that one objective (DP balance) dominates in realistic settings.

Significance. This is a fundamental advance in formulation rather than an incremental algorithmic improvement. By casting load balancing as a constrained optimization over slicing boundaries, the paper provides a precise mathematical language for what was previously an ad-hoc engineering concern. The formulation is general enough to accommodate non-linear cost functions (Appendix D.5), alternative communication primitives, and different parallelism strategies. The separation of the optimization (offline) from the execution (runtime) is also a design pattern that will likely influence future distributed training systems.

Innovation 4: The Unified Framework as a Category-Resolving Contribution

The paper's claim to provide a "Unified, Asynchronous, and Load-Balanced" framework is not merely marketing language—it represents a genuine category-resolving contribution that dissolves the previously accepted tradeoff between mathematical exactness, system throughput, and optimizer generality. This innovation is about the coherence of the design philosophy across parallelism dimensions rather than any single mechanism.

The trilemma that prior work accepted. Before Canzona, practitioners faced an implicit trilemma when deploying matrix-based optimizers at scale:

  • System-level exactness (preserve the optimizer's mathematical update rule) was achievable only through approaches that sacrificed throughput—either redundant compute (DDP-style, which violates memory scaling) or layer-wise partitioning (which violates communication geometry, incurring 2× All-Reduce volume and exposed optimizer-step communication).
  • System throughput (efficient communication and overlap) was achievable only through approaches that sacrificed mathematical exactness—algorithmic approximations like block-diagonal preconditioners, shard-local orthogonalization, or low-rank subspace projections, all of which alter the update rule and introduce convergence fidelity risk.
  • Optimizer generality (support for arbitrary matrix-based optimizers without per-algorithm engineering) was absent from both categories—every prior solution was optimizer-specific (MuonBP targets Muon, Distributed Shampoo targets Shampoo, layer-wise mapping requires per-optimizer adaptation) or required re-engineering for each new algorithm.

The accepted wisdom was that you could achieve at most two of these three simultaneously. Appendix E systematizes this taxonomy, showing how every prior approach occupies exactly two axes while compromising the third.

How Canzona resolves the trilemma. The paper's design achieves all three through a combination of architectural choices that work in concert:

  • Exactness is preserved by the decoupling principle: the optimizer's internal mathematics are never modified. The system handles atomicity through memory layout and scheduling, not through algorithmic approximation. The precision experiments (Section 5.3, Appendix C.4) confirm this: training loss trajectories for Muon, Shampoo, and SOAP are indistinguishable from the synchronous baseline across 400B tokens. This is a stronger claim than "approximately equivalent"—it is bit-exact convergence equivalence.

  • Throughput is achieved through the combination of zero-communication DP updates (enabled by static partitioning with geometric alignment) and asynchronous TP execution (enabled by Micro-Group batching and load-balanced Host Rank assignment). The 1.57× end-to-end speedup and 5.8× optimizer-step speedup (Figure 4) quantify this, but the key architectural property is that these gains come from eliminating redundant work and communication rather than from optimizing within a compromised paradigm.

  • Generality is achieved through the task abstraction: tensor updates are treated as generic computational tasks defined by cost metrics ($\mathcal{W}(p) = \text{numel}(p)$). The partitioner and scheduler are optimizer-agnostic by construction—they need no knowledge of whether the MatrixOp is Newton-Schulz, SVD, eigendecomposition, or Kronecker factorization. The experiments in Appendix C.4 validate this: the identical Canzona codebase, with zero modification, achieves comparable efficiency gains and convergence fidelity for Muon, Shampoo, and SOAP.

Why this is a conceptual innovation rather than just good engineering. The trilemma-resolution is the result of a specific architectural insight (decoupling) combined with a specific optimization formulation (static load balancing with geometric constraints) and a specific execution model (asynchronous compute across both parallelism dimensions). None of these elements alone resolves the trilemma; their combination does. The paper's Section 4.3 explicitly synthesizes these elements into a "cohesive design philosophy," arguing that the DP and TP strategies are not independent optimizations but share four unifying principles: unified abstraction, asynchronous execution, load-balanced static planning, and preservation of communication primitives. This synthesis is what elevates Canzona from a collection of techniques to a framework—a design template that future systems can instantiate for new optimizers, new model architectures, and new hardware topologies.

Evidence of generality beyond the primary experiments. The paper's scaling analysis (Appendix C.3) demonstrates that the load-balancing algorithms remain effective as DP size scales from 16 to 128 (Figure 8a), as TP size scales from 2 to 8 (Figure 8b), and as model size scales from 1.7B to 32B (Figure 9). The performance gap versus NV-layerwise widens with model size (Figure 6, bottom row), suggesting that the decoupling advantage becomes more pronounced as optimizer computation dominates. The full comparison against layerwise_optimizer (Appendix C.2, Figure 6) shows consistent superiority across all parallelism configurations tested (DP16-TP8, DP32-TP4, etc.), indicating robustness to topology. These results collectively support the claim that the framework is not fine-tuned to a specific scale or configuration.

Significance beyond the specific system. The trilemma-resolution establishes a new baseline for the field: future distributed optimizer frameworks will be evaluated against the standard Canzona sets—can they simultaneously preserve exactness, match or exceed this throughput, and support arbitrary optimizer algorithms without modification? The paper's explicit connection to FSDP (Appendix D.4) suggests the design principles generalize beyond Megatron, though the authors are careful to note that hybrid FSDP+TP configurations introduce additional communication complexity not directly addressed.

This innovation is fundamental because it redefines what is achievable. Before Canzona, the trilemma was a constraint on system design. After Canzona, it is a solved problem, and the research frontier shifts to questions like: how cheaply can difficulty estimation be amortized? Can the offline planning be made dynamic for elastic scaling? How does the framework perform on tasks without clean correctness signals? The paper did not just optimize an existing approach—it opened a new category of solutions.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper evaluates on the Qwen3 model family (Yang et al., 2025) spanning 1.7B to 32B parameters, trained with matrix-based optimizers using a sequence length of 4096 and a batch size per DP rank of 1. This is not a benchmark dataset in the traditional sense—the paper measures training throughput and optimizer latency rather than task accuracy. The primary model for detailed performance analysis is Qwen3-32B with the Muon optimizer unless otherwise specified. For precision verification (convergence fidelity), models are trained on 400B tokens.

  • Base model(s). The Qwen3 family (Yang et al., 2025) at scales of 1.7B, 4B, 8B, 14B, and 32B parameters. The choice is motivated by Qwen3's status as a contemporary, publicly documented LLM family representative of modern transformer architectures. The paper also verifies generality across three matrix-based optimizers—Muon (Jordan et al., 2024), Shampoo (Gupta et al., 2018), and SOAP (Vyas et al., 2024)—to demonstrate optimizer-agnostic behavior.

  • Metrics. Two categories of metrics are used. For performance: (1) Wall-clock time for different training phases—Fwd-Bwd time (forward and backward pass), Target Optimizer time (e.g., Muon step latency), and total iteration time—all averaged over 10 runs × 10 steps to ensure stability. (2) Load-Balance Ratio $R_{LB}$ for both memory footprint and computational FLOPs, defined as $R_{LB} = \max_r(\mathcal{v}_r) / \text{avg}_r(\mathcal{v}_r)$ where $\mathcal{v}_r$ is the metric value (Peak Memory or FLOPs) on rank $r$. An ideal balanced system yields $R_{LB} \approx 1.0$; higher indicates more severe imbalance. For precision verification: training loss trajectories are compared to confirm that Canzona's system-level optimizations do not alter the mathematical convergence behavior relative to the synchronous baseline.

  • Baselines. Four distinct strategies are compared:

    • SC (Synchronous Compute, Baseline): Represents the naive practice (DDP for DP, All-Gather for TP). Employs non-partitioning and synchronous collective communication, resulting in redundant computation and blocking synchronization. Corresponds to Paradigm 1 in Section 3.1.
    • NV-layerwise (Baseline): NVIDIA's layerwise_optimizer implementation (an open-source PR to Megatron-LM). Assigns optimizer states at the granularity of whole layers to respect tensor boundaries, but suffers from the ZeRO Geometric Incompatibility analyzed in Appendix D.2. Corresponds to Paradigm 2 in Section 3.1.
    • ASC (Asynchronous Compute): Adopts Canzona's decoupled architecture to enable asynchronous execution (DP-ASC + TP-ASC) but utilizes naive partitioning without load-aware scheduling. This ablation isolates the contribution of load balancing from the decoupling architecture.
    • LB-ASC (Load-Balanced Asynchronous Compute, Canzona's Core Strategy): The complete implementation integrating the α-Balanced Static Partitioning for DP and the Micro-Group Greedy Rollback scheduling for TP. Corresponds to Paradigm 3 in Section 3.1.

  • Hardware and scale. The primary experiments (Section 5.2) use a cluster of 256 GPUs with a distributed topology of DP size = 32 and TP size = 8. Scaling analyses (Appendix C.3) extend to up to 512 GPUs, with DP sizes ranging from 16 to 128 and TP sizes from 2 to 8. For Shampoo/SOAP experiments (Appendix C.4), the parallelism configuration adjusts to PP=2, DP=32, TP=4 on Qwen3-14B to accommodate the additional memory overhead of preconditioners.

  • Generation budget / compute accounting. Unlike the reference example (which measures test-time compute in "generations" for an inference workflow), this is a training systems paper, so compute is measured in wall-clock time for specific training phases and in FLOPs/Memory load per rank. The load function $\mathcal{W}(p) = \text{numel}(p)$ is used as a proxy for computational cost, with Appendix D.5 validating that FLOPs-based and numel-based scheduling produce negligible difference ($\approx 10^{-4}$ seconds for Qwen3-32B). All-to-All communication capacity is parameterized by $C_{\text{max}}$ (buffer size limit).

  • Cross-validation / statistical protocol. No cross-validation is reported (this is a systems measurement paper, not a statistical learning evaluation). Timings are averaged over 10 runs and 10 steps within each run to ensure measurement stability. Load-balance ratios are computed deterministically from the static partition map $\Pi$. Precision verification trains on 400B tokens with loss curves tracked throughout.

Main Quantitative Results

Effectiveness of Load Balancing Strategies (Figure 3)

The paper first establishes that load imbalance is severe under naive static partitioning and that Canzona's algorithms effectively neutralize it. Results are from Qwen3-32B with Muon on 256 GPUs (DP=32, TP=8).

Data Parallelism load balancing (Figure 3c). Without DP load balancing, the FLOPs imbalance ratio (Max/Avg) reaches 3.24× and the Memory imbalance ratio reaches 2.46×—meaning the most heavily loaded rank performs over three times the computation of the average rank, creating severe stragglers. The α-Balanced Partitioning reduces these ratios to 1.43× (FLOPs) and 1.11× (Memory), effectively flattening the workload distribution. The Memory ratio improves more than FLOPs because memory cost is strictly proportional to numel, while computational cost for matrix operations has additional shape-dependent factors that numel-based balancing cannot fully capture—though Appendix D.5 shows this residual imbalance is practically negligible.

Tensor Parallelism load balancing (Figure 3b). The naive TP baseline similarly suffers from high variance with a FLOPs imbalance ratio of 3.24×. The Micro-Group Scheduling improves this to 2.46× for FLOPs and 1.16× for Memory. The FLOPs improvement is more modest for TP than DP because TP tensors are already fragmented (each rank holds only a shard of each weight matrix), and the Host Rank assignment within a Micro Group must balance a smaller set of remaining tensors. The Memory improvement to 1.16× reflects the effectiveness of the greedy rollback in preventing any single Host Rank from accumulating excessive state.

Efficiency impact (Figure 3a). The load balancing directly translates to reduced makespan. The figure reports step time in TFLOPS-equivalent units (a normalized metric where higher values indicate better GPU utilization). LB-ASC achieves the lowest maximum step time (1.05 TFLOPS equivalent) compared to SC and ASC, with the visual representation showing effectively eliminated computation bubbles—the "Load Balance" state from Figure 1's bottom-right panel.

End-to-End Comparison with layerwise_optimizer (Figure 4)

The headline comparison on Qwen3-32B with Muon on 256 GPUs (DP=32, TP=8) yields:

  • Total iteration time: 0.877s (Canzona) vs. 1.381s (NV-layerwise) — a 1.57× speedup.
  • Optimizer step time: 0.066s (Canzona) vs. 0.383s (NV-layerwise) — a 5.8× speedup.
  • Fwd-Bwd time: 0.811s (Canzona) vs. 0.998s (NV-layerwise) — a 1.23× speedup.

The optimizer-step speedup is the dominant factor (5.8×), driven by Canzona's zero-communication DP updates (no Broadcast/All-Gather needed during the optimizer step) and asynchronous TP pipeline (hiding reconstruction communication). The Fwd-Bwd improvement (1.23×) stems from Canzona's use of efficient Reduce-Scatter (standard ZeRO-1 volume) versus NV-layerwise's forced All-Reduce (2× communication volume), as structurally analyzed in Appendix D.2.

The breakdown reveals an important subtlety: NV-layerwise's optimizer step alone (0.383s) is nearly half of Canzona's total iteration time (0.877s). This means the layer-wise approach's exposed optimizer communication is not a minor overhead—it is a dominant cost that would substantially erode any convergence benefits from using a matrix-based optimizer in the first place.

Full Performance Comparison Across Model Sizes and Configurations (Appendix C.2, Figure 6)

The paper extends the comparison across the entire Qwen3 family (1.7B to 32B) under multiple parallelism configurations. Key findings:

  • Consistent superiority: Canzona outperforms NV-layerwise in every tested configuration, with the total iteration time reduction ranging from modest to dramatic.
  • Scaling with model size: The performance gap widens as model size increases. For Qwen3-1.7B, the absolute time savings are bounded by the low optimizer cost. For Qwen3-32B (DP16-TP8), the optimizer latency reduction reaches approximately 8.3×. This is because larger models have more heavy matrix operations (wider attention projections, larger FFN layers) that expose the communication bottleneck of the layer-wise approach more severely.
  • Robustness to parallelism strategy: The advantage persists regardless of the DP/TP split. For Qwen3-32B, configurations of DP16-TP8 and DP32-TP4 both show substantial speedups, indicating that the load-balancing algorithms are not sensitive to the specific parallelism decomposition.

Verification of Fwd-Bwd Communication Efficiency (Appendix C.2, Figure 7)

To isolate the communication efficiency claim, the paper designs a controlled comparison using standard AdamW as a reference:

  • AdamW with All-Reduce (DDP) establishes the latency upper bound (inefficient, 2× communication volume).
  • AdamW with Reduce-Scatter (ZeRO-1) establishes the theoretical lower bound (optimal overlapping).

Under this comparison:

  • NV-layerwise closely tracks the AdamW All-Reduce baseline, confirming its Fwd-Bwd latency is structurally bottlenecked by 2× communication volume from forced All-Reduce.
  • Canzona closely tracks the AdamW Reduce-Scatter baseline, confirming that the static partitioning successfully preserves the efficient communication primitives.

The paper notes a minor qualification: Canzona's Fwd-Bwd time is "marginally higher" than the AdamW Reduce-Scatter baseline in some configurations because the load-balancing strategy introduces variable-sized parameter chunks (rather than perfectly equal chunks), creating slight communication imbalances. However, this impact is characterized as "negligible compared to the significant performance penalty of the inefficient All-Reduce primitive" (Appendix C.2).

Scalability Analysis (Appendix C.3, Figures 8 and 9)

Parallelism scaling (Figure 8). With Qwen3-32B and Muon fixed:

  • Data Parallelism (Figure 8a): As DP size scales from 16 to 128 (TP=4 fixed), the baseline ASC shows a "linear degradation" in load-balance ratio for both memory and FLOPs—the statistical variance in parameter sizes naturally increases with more ranks. In contrast, DP LB-ASC maintains a load-balance ratio close to the ideal 1.0, neutralizing the straggler effect even at 128 DP ranks. Optimizer time remains stable under LB-ASC while increasing under ASC.

  • Tensor Parallelism (Figure 8b): Scaling TP size from 2 to 8 (PP=4, DP=4 fixed) exacerbates fragmentation of weight matrices. The baseline shows a sharp increase in computational imbalance. TP LB-ASC (Micro-Group Scheduling) keeps the FLOPs load-balance ratio significantly lower, preserving low optimizer latency as TP scales.

Model size scaling (Figure 9). With parallelism fixed at DP=16, TP=4:

  • DP Load Balance (Figure 9a): Larger models contain wider variance in tensor shapes (large embedding layers vs. small projection heads). The naive ASC's load-balance ratio increases significantly with model size (1.7B to 32B). LB-ASC adapts to this heterogeneity, maintaining a flat profile—meaning the algorithm's effectiveness does not degrade as model architecture becomes more diverse.

  • TP Load Balance (Figure 9b): Interestingly, the TP baseline imbalance does not strictly increase with model size but fluctuates based on specific architecture (e.g., hidden dimension alignment). However, the greedy scheduling consistently finds near-optimal packing, sustaining the performance advantage across all model sizes.

Generalization to Shampoo and SOAP (Appendix C.4)

Efficiency (Figures 10a, 11a). On Qwen3-14B with 256 GPUs (PP=2, DP=32, TP=4—adjusted for the higher memory of Shampoo/SOAP preconditioners):

  • Shampoo: SC baseline optimizer step takes 3.313s; Canzona's LB-ASC reduces this to 0.110s — a ~30× speedup. The SC baseline is dramatically slower for Shampoo than Muon because Shampoo involves cubic-complexity operations (SVD) on larger preconditioner matrices that are redundantly computed on every rank in the synchronous approach.
  • SOAP: Similar drastic reduction (exact numbers in Figure 11a), confirming that the framework's optimizer-agnostic design handles diverse matrix operations.

Precision (Figures 10b, 11b). Qwen3-1.7B trained for 400B tokens (DP=8, TP=4):

  • Both Shampoo and SOAP training loss curves under LB-ASC overlap perfectly with the standard synchronous baseline (SC). There is zero mathematical precision loss—the convergence trajectory is identical because Canzona implements a system-level optimization that does not modify the optimizer's update rule.

Load Balance (Figure 12). The load distribution analysis for Shampoo/SOAP shows the same pattern as Muon: naive partitioning creates significant computational bubbles (FLOPs load-balance ratio $>$ 2.0), while LB-ASC flattens the workload variance to ~1.05, achieving near-perfect balance. This is notable because Shampoo and SOAP have different computational characteristics from Muon (preconditioner construction vs. Newton-Schulz orthogonalization), yet the numel-based cost model generalizes effectively.

Precision Verification (Section 5.3, Figure 5)

The primary precision experiment trains Qwen3-1.7B on 400B tokens using Muon (DP=8, TP=4). The training loss trajectory of LB-ASC is indistinguishable from the synchronous baseline (SC). This confirms the paper's central claim that Canzona's decoupled partitioning and scheduling strategies act as purely system-level optimizations, preserving the exact convergence behavior of the original matrix-based optimizer. The paper emphasizes this is not "approximately equivalent" but mathematically exact—no approximations are introduced, no update rules are modified, and the optimizer sees the identical gradient and state information it would see in a single-device execution.

Ablation Studies and Robustness Checks

DP Load-Balance Factor α (Appendix C.5, Figure 13): On Qwen3-32B with 128 GPUs (PP=8, DP=16), α is swept from 0.0 to 1.0. As α increases (prioritizing compute balance over communication uniformity), the Muon optimizer time decreases monotonically. The Fwd-Bwd time remains "relatively stable" across all α values, confirming that the communication imbalance introduced by non-uniform shard sizes is effectively hidden by the standard Megatron communication-computation overlap. Setting α = 1.0 yields the best end-to-end performance, validating the design choice to prioritize computational load balancing. This ablation is significant because it demonstrates that the dual-objective formulation, while theoretically general, effectively collapses to a single objective in practice for this hardware and model configuration.

TP Micro-Group Fusion Capacity (Appendix C.6, Figure 14): On 128 GPUs (DP=16, TP=8), the capacity constraint $C_{\text{max}}$ (buffer size for fused All-to-All) is varied. The "No-Fuse" baseline (per-tensor communication) suffers from high latency (~0.11s) due to many small kernel launches and poor bandwidth utilization. Enabling fusion immediately drops latency to ~0.073s. Performance improves slightly as $C_{\text{max}}$ increases and plateaus around 512 MB, indicating that All-to-All bandwidth is fully saturated beyond this point. This confirms both that fusion is essential for efficiency and that the system is robust to the specific capacity choice as long as it exceeds the saturation threshold.

Cost Metric Ablation (Appendix D.5, Figure 16): Comparing numel-based scheduling vs. exact FLOPs-based scheduling for Qwen3-32B (DP=16, TP=8) with Muon: the execution makespan differs by only ~10⁻⁴ seconds (0.0717s vs. 0.0718s). This validates that numel serves as an accurate proxy for load balancing transformer architectures—shape correlates strongly with computational cost, and the minute discrepancy does not materially affect the optimizer-step latency.

Fwd-Bwd Communication Isolation (Appendix C.2, Figure 7): As discussed above, this controlled experiment using AdamW baselines confirms that Canzona's Fwd-Bwd latency tracks the optimal Reduce-Scatter baseline while NV-layerwise tracks the suboptimal All-Reduce baseline. The experiment isolates the geometric incompatibility as the structural cause of NV-layerwise's Fwd-Bwd penalty, not implementation quality or configuration.

Generalization Across Optimizers (Appendix C.4, Figures 10-12): This serves as both a main result and an ablation confirming the optimizer-agnostic claim. The identical Canzona codebase, with identical hyperparameters, achieves comparable efficiency gains and zero-fidelity-loss precision for Muon, Shampoo, and SOAP—three optimizers with distinct computational patterns (Newton-Schulz iterations, SVD-based preconditioners, eigendecomposition-based preconditioners). No per-optimizer tuning was performed.

Model Size Robustness (Appendix C.3, Figure 9): The load-balance algorithms maintain effectiveness across model sizes from 1.7B to 32B parameters, with the DP algorithm (Figure 9a) showing a flat load-balance profile even as model heterogeneity increases. This is non-trivial: larger models typically have more diverse tensor shapes (wider ranges of hidden dimensions across layers, larger embedding matrices relative to other parameters), so a load-balancing strategy that works at 1.7B might degrade at 32B if it relies on assumptions about parameter homogeneity. The results show it does not.

Parallelism Configuration Robustness (Appendix C.2, Figure 6): The framework's advantage over NV-layerwise persists across multiple DP/TP splits (DP16-TP8, DP32-TP4, etc.), indicating that the load-balancing algorithms are not sensitive to the specific decomposition. This is important because practitioners choose parallelism configurations based on model size, hardware topology, and batch size constraints—a solution that only works for specific splits would be of limited practical value.

Critical Assessment

Does the 1.57× End-to-End Speedup Represent a Fair Comparison?

The headline 1.57× speedup (Figure 4) compares Canzona against NVIDIA's layerwise_optimizer. This is a legitimate and relevant baseline—it represents the current state-of-the-art for deploying matrix-based optimizers with mathematical exactness in Megatron—but several qualifications are necessary:

The comparison assumes identical optimizer mathematics. Both systems compute the exact same matrix update. This is validated by the precision experiments, but it means the speedup is purely from system efficiency, not algorithmic differences. The question is whether layerwise_optimizer is the strongest possible synchronous-exact baseline. The paper's own analysis (Appendix D.2) demonstrates that layerwise_optimizer is structurally suboptimal due to the Geometric Incompatibility, so it is in some sense a "weak" baseline by design. However, this is precisely the point: the paper argues that the geometric incompatibility is inherent to layer-wise approaches, not an implementation flaw, so any exact approach using layer-wise assignment will suffer this penalty. If there exists an alternative exact approach that avoids both the geometric incompatibility and the load imbalance without Canzona's decoupling, it is not presented or cited.

The 5.8× optimizer-step speedup dominates but may not fully generalize. The optimizer-step speedup (0.066s vs. 0.383s) is the primary driver of the total speedup. This gap arises from NV-layerwise's exposed All-Gather/Broadcast during the optimizer step (to redistribute updated parameters) versus Canzona's zero-communication DP updates. The magnitude of this gap depends on the ratio of optimizer computation time to communication time. For very small models where optimizer computation is negligible, the gap would narrow. For optimizers with even heavier computation (Shampoo shows a ~30× gap in Figure 10a), the gap would widen. The paper's scaling results (Figure 6) confirm the gap widens with model size, which is the regime of practical interest for matrix-based optimizers, so the headline number is representative of large-scale deployment.

The Fwd-Bwd improvement (1.23×) is structural but modest. This improvement comes from Canzona's use of Reduce-Scatter vs. NV-layerwise's forced All-Reduce. A 1.23× Fwd-Bwd speedup is meaningful but not transformative—the bulk of the total improvement comes from the optimizer step. The paper's Figure 7 convincingly demonstrates that this gap is due to the 2× communication volume difference, not overlapping quality. However, the 1.23× figure is measured at a specific parallelism configuration (DP=32, TP=8). At different DP sizes, the All-Reduce vs. Reduce-Scatter gap might be larger or smaller depending on how communication time compares to computation time in the Fwd-Bwd pass.

Does Canzona Actually Resolve the "Trilemma" (Exactness, Throughput, Generality)?

Exactness: Strongly supported. The precision experiments (Figures 5, 10b, 11b) show indistinguishable loss curves across 400B tokens for three different optimizers. This is the gold standard for claiming zero fidelity loss—if the convergence trajectory matches exactly, the system-level optimization has not altered the mathematics in any way.

Throughput: Supported with boundary conditions. The 1.57× improvement is demonstrated at scale (256 GPUs, 32B parameters) against a legitimate baseline. However, the comparison does not include algorithmic approximation approaches (MuonBP, block-diagonal Shampoo) that sacrifice exactness for throughput. Such a comparison would test whether the exactness-throughput tradeoff is genuinely resolved or whether approximations remain faster at the cost of convergence quality. The paper's framework provides a basis for such a future comparison but does not conduct it.

Generality: Strongly supported within the tested scope. The identical codebase achieves comparable efficiency gains for Muon, Shampoo, and SOAP—three optimizers with substantially different computational patterns. The scaling analysis shows robustness across model sizes and parallelism configurations. However, "generality" is tested only within the Megatron framework and the Qwen3 model family. The paper's Appendix D.4 argues the principles extend to FSDP, but this is a design argument, not an empirical demonstration. Similarly, only transformer-based language models are tested; the load-balancing behavior might differ for CNNs, mixture-of-experts architectures, or models with heterogeneous layer types not present in Qwen3.

Are the Load-Balancing Algorithms Validated to Be Necessary (Not Just Helpful)?

The ASC ablation (decoupled architecture without load balancing) is critical for this claim. The paper shows that ASC reduces optimizer time compared to SC but still suffers from stragglers, and that LB-ASC further reduces optimizer time (Figures 3a, 8). However, the quantitative contribution of load balancing versus decoupling alone is not always cleanly separated:

  • For DP scaling (Figure 8a), ASC optimizer time degrades as DP size increases while LB-ASC remains stable—clearly demonstrating that load balancing is necessary for strong scaling.
  • For the main Figure 4 comparison against NV-layerwise, the LB-ASC result is the only Canzona configuration shown. An ASC-only comparison would reveal how much of the 5.8× optimizer-step speedup comes from decoupling (zero-communication DP, async TP) versus load balancing (eliminating stragglers within those paradigms). The straggler elimination presumably contributes to the optimizer-step time but not the Fwd-Bwd time (since the Fwd-Bwd uses standard communication patterns regardless of load balance). The paper could strengthen this by reporting the ASC optimizer-step time in the Figure 4 configuration to decompose the contribution.

What Experiments Are Missing?

Comparison against algorithmic approximations. The paper critiques MuonBP, block-diagonal Shampoo, and low-rank approximations (Appendix E.3) for sacrificing mathematical fidelity, but never compares throughput against them. A head-to-head comparison would quantify the "fidelity tax"—how much slower is Canzona (exact) versus MuonBP (approximate), and whether the convergence benefit of exactness justifies the throughput difference in end-to-end training time to reach a target loss. This is a complex experiment (it requires training multiple models to convergence under each approach), but it would directly test the paper's central claim that the exactness-throughput tradeoff is false.

Memory footprint comparison. The paper focuses on throughput (iteration time) and load-balance ratios, but does not report the peak memory usage of Canzona versus baselines. The load-balance ratio for Memory (e.g., 1.11× after α-Balanced partitioning) indicates variation across ranks, but the absolute memory consumption per GPU is not reported. For practitioners, memory is often the binding constraint—if Canzona's static partitioning increases per-GPU memory usage compared to layerwise_optimizer (e.g., because larger contiguous parameter intervals reduce opportunities for memory optimizations), this would be a significant practical consideration. The α-Balanced algorithm optimizes for memory balance across ranks but does not explicitly minimize total memory consumption.

Impact of the α = 1.0 finding on communication imbalance. The α ablation (Figure 13) shows that α = 1.0 is optimal because communication imbalance is hidden by overlap. However, this hiding depends on the overlap window being sufficiently large. At extreme scales (e.g., 512 GPUs with very large models where the Fwd-Bwd computation per bucket is short relative to communication time), the overlap might be less effective, and the communication imbalance from α = 1.0 might become exposed. The paper's scaling analysis (up to 128 DP ranks) does not reveal such degradation, but testing at larger scales or with slower interconnects would probe the boundary of this finding.

Offline planning cost amortization. Appendix D.1 states the offline planning completes in milliseconds, which is negligible compared to training time. However, this is reported as a qualitative claim without measurement data. A simple timing of the α-Balanced algorithm and Micro-Group construction for Qwen3-32B would make this claim falsifiable and provide practitioners with concrete expectations for their own model sizes.

Dynamic scheduling or mid-training rebalancing. The paper's static planning is computed once before training. If the model architecture changes during training (e.g., staged training where certain layers are unfrozen progressively), the static partition map would need recomputation. The paper does not address how the framework handles such scenarios. Similarly, elastic scaling (adding or removing GPUs mid-training) would require recomputing the partition map—the paper's offline assumption makes this a non-trivial operational constraint that is acknowledged but not deeply explored.

Do the Results Support the Claim That Canzona "Opens the Door for Matrix-Based Optimizers to Become the Default"?

This is a forward-looking claim, not directly tested. The experiments demonstrate that Canzona makes matrix-based optimizers feasible at scale without throughput sacrifice, which is a necessary condition for them becoming the default—but not sufficient. The paper does not demonstrate that the convergence benefits of Muon/Shampoo/SOAP over AdamW (which are cited from prior work but not experimentally replicated here) persist when combined with Canzona at scale. A complete demonstration would show end-to-end training of a production-scale model with Canzona + Muon reaching a target loss in fewer total GPU-hours than AdamW with standard ZeRO-1. This is beyond the paper's scope (which focuses on system efficiency, not optimizer convergence comparisons), but it is the experiment that would validate the implied promise.

Summary of Evidentiary Strength

The paper's core system claims—that decoupling logical assignment from physical distribution enables exact matrix-based optimizers to run efficiently, that the resulting load imbalance can be addressed through static optimization, and that the geometric compatibility with ZeRO-1 primitives is preserved—are well-supported by the presented experiments. The 1.57× end-to-end speedup and 5.8× optimizer-step speedup are robust and consistent across model sizes and configurations.

The paper's broader claims—that this framework resolves the exactness-throughput-generality trilemma, that the principles generalize beyond Megatron, and that this opens the door for matrix-based optimizers to become the default—are plausible and consistent with the evidence but not fully tested. The missing comparisons against algorithmic approximations, the limited exploration of memory footprint, and the absence of end-to-end convergence-efficiency benchmarks against AdamW leave these claims as well-motivated hypotheses rather than demonstrated facts. The paper provides the system infrastructure that makes such comparisons possible, which is a genuine contribution even if the comparisons themselves remain future work.

6. Limitations and Trade-offs

6.1 The Load-Balancing Optimization Assumes Static, Known Model Architectures (No Mid-Training Structural Changes)

The assumption or constraint. Canzona's entire workflow depends on a one-time offline planning phase during model initialization (Section 3.3, Appendix D.1) that computes the Global Partition Map Π (for DP) and the Micro-Group sequence 𝕄 (for TP) based on the complete parameter list of the model. These schedules are "determined during the initialization phase" (Section 4.1) and remain fixed for the entire training run. The paper explicitly notes that the offline planning is "a one-time offline planning step during model initialization" (Appendix D.1), completing in milliseconds—but this assumes the model architecture is fully known and static at initialization time.

The consequence. Several increasingly common training paradigms violate this assumption:

  • Staged training (progressive layer unfreezing, gradual width/depth scaling): If layers are added, removed, or unfrozen mid-training, the parameter list changes. The partition map Π would need recomputation, and more critically, the optimizer states for the new parameters would need to be allocated and potentially redistributed across ranks. The paper provides no mechanism for incremental rebalancing—recomputing Π from scratch would require reallocating the param_and_grad_buffer, which in Megatron is a heavyweight operation typically done once at startup.

  • Mixture-of-Experts (MoE) architectures: MoE models route tokens to different experts per batch, meaning the "active" parameter set varies per iteration. The static partition map assumes every parameter is updated every step (or at least that the set of parameters needing updates is fixed). In MoE training, expert parameters may receive gradients only when tokens are routed to them, creating dynamic load patterns that static balancing cannot anticipate.

  • Elastic scaling (adding/removing GPUs mid-training): If the DP or TP size changes (e.g., due to node failures or dynamic resource allocation), the partition map Π is invalidated because the assignment logic depends on R (the number of ranks). Recomputing Π would require redistributing optimizer states across the new rank topology—a global synchronization operation that is not part of Canzona's runtime workflow. The paper's design assumes a static cluster topology for the training duration.

What evidence exists in the paper. The paper does not evaluate any of these scenarios. All experiments use a fixed model architecture (Qwen3 family) with a fixed number of GPUs throughout training. The precision verification runs 400B tokens without architectural changes. The scaling analysis (Appendix C.3) varies parallelism configuration between runs, not dynamically within a run. The paper does not cite or discuss MoE architectures, staged training, or elastic scaling. Appendix D.4 discusses adapting Canzona to FSDP but does not address dynamic topology changes.

Mitigation status. Not addressed. The paper's design philosophy explicitly embraces static planning as a feature (enabling zero runtime scheduling overhead and fixed optimizer state placement), but does not discuss the tradeoff this creates for dynamic training regimes. The authors acknowledge that the offline nature is a design choice (Appendix D.1 emphasizes the negligible initialization cost), but do not characterize which training paradigms it excludes or how the framework could be extended to support them.

6.2 The Load-Balancing Cost Model (numel) Is a Proxy that May Fail for Non-Transformer Architectures or Future Optimizers

The assumption or constraint. The α-Balanced Greedy LPT algorithm (Algorithm 1) and the Micro-Group Greedy Rollback algorithm (Algorithm 2) both rely on a load function \mathcal{W}(p) to estimate the computational cost of updating each parameter. The paper adopts \mathcal{W}(p) = \text{numel}(p) (number of elements) as the universal cost metric, justifying it on three grounds (Appendix D.5): (1) optimizer-agnostic universality—the system doesn't need to know whether an optimizer is \mathcal{O}(N) or \mathcal{O}(N^3), (2) shape-cost correlation: for standard transformer architectures, parameter count correlates strongly with actual computation, and (3) a controlled experiment (Figure 16) shows the difference between numel-based and exact FLOPs-based scheduling is negligible (≈ 10^{-4} seconds for Qwen3-32B with Muon).

The consequence. The numel proxy can fail in two scenarios:

  • Non-transformer architectures with irregular tensor shapes: The shape-cost correlation the paper relies on emerges from the regularity of transformer architectures (attention projections, FFN layers, embedding matrices all have predictable relationships between d_model, d_ff, and num_heads). For architectures with heterogeneous operators—CNNs with varying kernel sizes and spatial dimensions, graph neural networks with irregular adjacency structures, or models mixing dense and sparse layers—numel may not correlate with actual FLOPs. A 1024×1024 convolution and a 1024×1024 dense matrix have identical numel but vastly different computational costs due to differing FLOPs per element.

  • Future optimizers with non-standard complexity scaling: The paper's numel proxy implicitly assumes that all matrix-based optimizers have roughly similar complexity scaling with tensor dimensions—specifically, that the dominant cost is proportional to the number of elements times some optimizer-specific constant factor. This holds approximately for Muon (Newton-Schulz iterations scale with matrix dimensions), Shampoo (preconditioner construction scales with gradient dimensions), and SOAP (eigendecomposition scales similarly). However, an optimizer whose complexity scales with a different function of tensor shape—e.g., one that is \mathcal{O}(d_{\text{in}}^2 \cdot d_{\text{out}}) for rectangular matrices—would cause numel-based scheduling to systematically underestimate the cost of wide, shallow matrices relative to narrow, deep ones of the same element count, creating load imbalance that the algorithm cannot correct.

What evidence exists in the paper. The paper tests numel-based scheduling on three optimizers (Muon, Shampoo, SOAP) on one model family (Qwen3, a standard decoder-only transformer). The Figure 16 ablation comparing numel vs. FLOPs-based scheduling is only shown for Muon on Qwen3-32B—not for Shampoo or SOAP, and not for non-transformer architectures. The load-balance ratios (Figure 3b, 3c) show residual imbalance even after LB-ASC (FLOPs ratio 1.43× for DP, 2.46× for TP), which the paper attributes to shape-dependent factors that numel cannot capture. This residual is small enough to not materially affect throughput in the tested configurations, but its magnitude for untested architectures is unknown.

Mitigation status. The paper acknowledges the limitation explicitly (Appendix D.5): "our framework formulation fully supports the generalized non-linear complexity." The mathematical formulation of the optimization problems (Sections 3.2, 4.2) uses \mathcal{W}(p) as an abstract cost function without committing to numel. The practical implementation uses numel for simplicity, and the paper frames this as a pragmatic choice justified by evidence rather than a fundamental constraint. However, no experiments demonstrate how to plug in a custom cost function, and no guidance is provided for practitioners who might need to define optimizer-specific or architecture-specific cost models.

6.3 Memory Footprint Is Not Evaluated, Despite Being a Binding Constraint for Large-Scale Training

The assumption or constraint. The paper's evaluation focuses exclusively on throughput (iteration time, optimizer-step latency) and load-balance ratios (a dimensionless metric of distribution quality), but never reports absolute per-GPU memory consumption. The load-balance ratios for Memory (e.g., 1.11× for DP after α-Balanced partitioning, Figure 3c; 1.16× for TP after Micro-Group scheduling, Figure 3b) measure variation across ranks within a configuration, not the total memory required per GPU compared to baselines.

The consequence. For practitioners, memory is often the binding constraint that determines whether a model can be trained at all on available hardware—more fundamental than throughput. Several aspects of Canzona's design could increase peak memory usage in ways not captured by load-balance ratios:

  • Atomic parameter assignment (DP): By enforcing that entire parameters are owned by single ranks rather than sharded uniformly, Canzona may create situations where a rank must allocate memory for a large contiguous parameter block that, under standard ZeRO-1 equal-chunk partitioning, would be split across multiple ranks. The α-Balanced algorithm optimizes for balance of memory across ranks (minimizing the Max/Avg ratio), but does not minimize the total memory per rank. If the total optimizer state size is fixed, balancing across ranks cannot reduce the mean—it can only prevent some ranks from exceeding the mean dramatically. But if the atomicity constraint forces certain ranks to hold parameters that were previously split across ranks (increasing the per-rank maximum), the peak memory per GPU could be higher than standard ZeRO-1, even if the distribution is balanced.

  • Micro-Group All-to-All buffers (TP): The fused All-to-All communication in the TP pipeline requires buffer space for gathering gradient shards and scattering update shards. The capacity parameter C_{\text{max}} constrains this, but the absolute buffer size needed depends on the group size. The paper shows that performance plateaus at C_{\text{max}} ≈ 512 MB (Figure 14), meaning the TP pipeline requires at least this much additional buffer memory per rank that is not needed in synchronous TP approaches (where reconstruction is done in-place or with smaller temporary buffers).

  • Offline planning memory: The partition map Π and scheduling plan 𝕄 themselves consume memory, though this is presumably negligible compared to model parameters.

What evidence exists in the paper. No absolute memory measurements are reported. The load-balance ratios for memory (Figures 3b, 3c, 8, 9, 12) show relative improvement over naive partitioning but do not compare against standard ZeRO-1 (equal chunks) or NV-layerwise baselines. The paper does not report whether Canzona's peak GPU memory is higher, lower, or equivalent to these baselines for any model size. The experiments on Shampoo/SOAP (Appendix C.4) use a smaller model (Qwen3-14B) with added pipeline parallelism (PP=2) specifically "to accommodate the state requirements"—suggesting that memory constraints are a practical concern, but no quantitative comparison is provided.

Mitigation status. Not addressed. The paper's focus is on throughput optimization, and the load-balancing formulation optimizes memory distribution (Equation 2 for DP, the makespan minimization for TP) but does not formulate memory capacity as an objective. The α-Balanced algorithm's Memory load-balance ratio of 1.11× (Figure 3c) indicates successful equalization, but the absolute memory consumption relative to baselines remains unknown. Future work would need to measure and report per-GPU peak memory to enable practitioners to assess whether Canzona fits within their hardware constraints.

6.4 The Framework Is Validated on a Single Model Family (Qwen3) and a Single Training Framework (Megatron)

The assumption or constraint. All experiments use the Qwen3 model family (Yang et al., 2025) at scales from 1.7B to 32B parameters, trained within the Megatron distributed training framework (Shoeybi et al., 2019). The paper's design principles are argued to generalize—Appendix D.4 sketches how the approach would apply to FSDP, and Section 4.3 claims the framework is "optimizer-agnostic" and compatible with "any current or future matrix-based optimizers"—but the empirical validation is restricted to one architecture family, one parallelism paradigm (Megatron's ZeRO-1 + TP), and three optimizers from the same broad class.

The consequence. Several generalization questions remain empirically open:

  • Architecture diversity: Qwen3 is a standard dense decoder-only transformer with uniform layer structure (self-attention followed by FFN, repeated). Architectures with heterogeneous layer types—encoder-decoder models, mixture-of-experts, vision transformers with varying spatial resolutions, or multimodal models combining different backbone types—may exhibit different parameter-size distributions that challenge the numel-based load-balancing heuristics. The paper's finding that the DP load-balance ratio degrades with model size for naive ASC but remains flat for LB-ASC (Figure 9a) is encouraging, but it is tested only on Qwen3's particular scaling pattern (where larger models primarily increase d_model and number of layers, preserving architectural regularity).

  • Training framework semantics: Canzona's DP strategy depends intimately on Megatron's specific buffer management: the param_and_grad_buffer, bucket-based communication overlapping, and the Start_Index-anchored geometric constraints (Equation 1). PyTorch FSDP (Zhao et al., 2023) uses a different sharding mechanism (per-parameter FlatParameter wrappers with shard placement) that does not directly expose the contiguous-buffer-with-geometric-slicing abstraction. The paper's Appendix D.4 acknowledges these differences and argues that the TP-style asynchronous compute pipeline could be adapted, but crucially notes that "hybrid FSDP+TP configurations introduce additional communication complexity not directly addressed." DeepSpeed ZeRO has yet different semantics. A practitioner using a non-Megatron framework cannot assume Canzona's algorithms transfer without adaptation.

  • Hardware topology assumptions: The TP strategy's viability depends on "the high-bandwidth Intra-node domain (e.g., via NVLink)" (Section 4). For clusters where TP spans inter-node connections (e.g., due to small node sizes or large TP degrees), the All-to-All communication in the Micro-Group pipeline would traverse slower links, potentially making the reconstruction overhead non-hidable and eroding the asynchronous compute benefits. The paper's experiments use TP=8 (fitting within a single 8-GPU node) and TP=4, both within the intra-node regime. No experiments test cross-node TP.

What evidence exists in the paper. The model-size scaling (Figure 9) covers 1.7B to 32B within Qwen3, and the parallelism scaling (Figure 8) varies DP up to 128 and TP up to 8. This is thorough for the chosen model family and framework but does not constitute cross-architecture or cross-framework validation. The optimizer-agnosticism claim is validated across Muon, Shampoo, and SOAP—three optimizers with different computational patterns—lending credibility, but all three are from the same broad class (second-order/preconditioned methods for matrix parameters). The paper does not test whether the framework works for optimizers that target non-matrix parameters (e.g., embedding-specific optimizers) or that combine matrix-based and element-wise updates within the same optimizer.

Mitigation status. The paper is transparent about scope: the title and abstract specify Megatron-compatible distributed training, and the framing is explicitly about resolving the "system-algorithm conflict" in "modern massive-scale training stacks" exemplified by Megatron. The appendix discussion of FSDP (Appendix D.4) acknowledges limitations and sketches a path to generalization without claiming it is solved. However, the strong language about being a "Unified framework" and "optimizer-agnostic" (Section 4.3, Appendix E.4) implies broader generality than the experiments directly support. A reader should understand that "unified" means unified across the three tested optimizers within Megatron on Qwen3, not validated across diverse frameworks and architectures.

6.5 The ZeRO Geometric Compatibility Tradeoff: Load-Balanced Shards Create Variable-Size Communication That Could Become Exposed at Extreme Scale

The assumption or constraint. The α-Balanced Static Partitioning trades off uniform communication shard sizes for computational load balance. When α = 1.0 (the optimal setting per Figure 13), the algorithm aggressively prioritizes DP compute balance, producing non-uniform shard sizes S_{i,r} within each bucket. The paper argues this is acceptable because "the slight latency variance caused by the non-uniform data sizes is effectively hidden" by Megatron's standard overlap of Reduce-Scatter/All-Gather with backward/forward computation (Section 3.3). This claim is validated at up to DP=128 with Qwen3-32B (Figure 8a, Appendix C.5), where Fwd-Bwd time remains "relatively stable" as α varies from 0.0 to 1.0.

The consequence. The "hiding" of communication imbalance depends on the overlap window—the amount of backward/forward computation available to mask the variable-duration communication. If the overlap window shrinks relative to the communication variance, the imbalance becomes exposed and adds to the critical path. This can happen in several regimes:

  • Very large DP degrees: As DP size increases, each rank's shard size S_{i,r} decreases (since total bucket size is fixed but divided among more ranks). The absolute variance in shard sizes also decreases, but the number of ranks participating in the collective increases, which can amplify the tail latency of Reduce-Scatter/All-Gather operations. At extreme DP scales (e.g., 512 or 1024 ranks), the communication time becomes dominated by the slowest rank's contribution, making any shard-size imbalance more impactful.

  • Small per-bucket computation: If the model architecture or micro-batch size results in very little computation per bucket (e.g., narrow models with tiny FFN layers), the backward/forward computation completes quickly and the overlap window may not fully absorb the communication variance. Megatron's overlap is effective when communication time is less than or comparable to computation time; if communication dominates, variances are directly exposed.

  • Slower interconnects: The paper's experiments use high-bandwidth intra-node and inter-node interconnects typical of cloud GPU clusters. On clusters with lower-bandwidth interconnects (e.g., older InfiniBand generations, Ethernet-based clusters), the absolute communication time is larger, making the variance from non-uniform shards a larger absolute quantity that is harder to hide.

What evidence exists in the paper. Figure 13 (α sensitivity, DP=16 on 128 GPUs) shows that Fwd-Bwd time is stable across α values, confirming the overlap is effective at this scale. Figure 8a (DP scaling up to 128 ranks) shows LB-ASC optimizer time remains stable, but does not separately report Fwd-Bwd time at different DP scales with α = 1.0 versus α = 0.0 to quantify whether the overlap effectiveness degrades at larger DP. Figure 7 compares Canzona's Fwd-Bwd against AdamW baselines and notes it is "marginally higher" than the ideal Reduce-Scatter baseline—this gap is the exposed portion of the communication imbalance, but its scaling behavior with DP size is not characterized.

Mitigation status. The paper acknowledges the communication imbalance as a "secondary concern" (Section 3.1) and empirically validates that it is hidden in the tested configurations. The α parameter provides a tunable mechanism to trade off compute balance for communication uniformity if the imbalance becomes exposed—a practitioner experiencing Fwd-Bwd degradation at extreme scale could reduce α (e.g., to 0.5 or 0.0) to recover communication uniformity at the cost of increased optimizer-step stragglers. The paper does not provide guidance on how to detect when communication imbalance has become the bottleneck (e.g., profiling Fwd-Bwd time versus optimizer time at different α values), leaving this as an engineering judgment call. The lack of scaling data for Fwd-Bwd time versus DP size with α = 1.0 means the "safe" regime for aggressive load balancing is not characterized.

6.6 Comparison Is Limited to Exactness-Preserving Baselines—No Throughput Comparison Against Algorithmic Approximations

The assumption or constraint. All baseline comparisons—SC (Synchronous Compute), NV-layerwise (NVIDIA's layerwise_optimizer), and ASC (ablation of Canzona without load balancing)—are approaches that preserve the exact mathematical update rule of the optimizer. The paper's narrative (Section 1, Appendix E) explicitly contrasts these "system-level exact" approaches with "algorithmic approximations" (MuonBP, block-diagonal Shampoo, low-rank methods) that sacrifice mathematical fidelity for throughput, arguing that Canzona resolves this tradeoff by achieving both exactness and efficiency simultaneously. However, no head-to-head throughput comparison against any algorithmic approximation is conducted.

The consequence. A practitioner deciding whether to adopt Canzona faces an incomplete comparison. The paper demonstrates that Canzona is ~1.57× faster than the best exact baseline (NV-layerwise) and ~5.8× faster in optimizer-step time. But it does not answer: how much faster (or slower) is Canzona compared to MuonBP, block-diagonal Shampoo, or low-rank approximations? If an approximation approach achieves, say, a 3× speedup over NV-layerwise (compared to Canzona's 1.57×) but with a 2% degradation in convergence quality, a practitioner must weigh throughput against fidelity. Without this comparison, the paper's implicit claim—that exactness can be achieved without a meaningful throughput penalty—remains an untested hypothesis relative to the approximation approaches the paper critiques.

This is significant because the paper's primary contribution is demonstrating that the exactness-throughput tradeoff can be avoided, not merely optimized. If approximations remain substantially faster than Canzona despite their fidelity loss, the practical value proposition shifts: Canzona becomes the choice for practitioners who cannot tolerate any convergence risk (e.g., training a production model where reproducibility is paramount), while approximations remain preferable for throughput-maximizing scenarios (e.g., rapid experimentation). The paper does not help practitioners make this distinction quantitatively.

What evidence exists in the paper. The paper's Appendix E.3 catalogs the limitations of algorithmic approximations—MuonBP's "directional drift" from local vs. global Newton-Schulz, block-diagonal Shampoo's loss of off-diagonal correlations, low-rank methods' subspace constraints—but provides no runtime measurements for any of these approaches. The precision experiments (Figures 5, 10b, 11b) show Canzona matches the synchronous exact baseline, confirming fidelity preservation, but do not show the throughput of alternative approximate methods trained to the same loss. The reference to MuonBP (Khaled et al., 2025) notes that it "may degrade convergence speed or solution quality" but does not cite quantitative evidence for the magnitude of this degradation relative to the throughput gain.

Mitigation status. Not addressed. The paper explicitly scopes itself to system-level optimizations that preserve exactness (Section 4.3: "neither the DP partitioner nor the TP scheduler requires modification to the optimizer's internal mathematics"), making the comparison against exact baselines the appropriate primary evaluation. However, the strong claims in the introduction and Appendix E about "bridging the gap" and "resolving the fundamental conflict" create an expectation that the practical tradeoff has been empirically characterized, which it has not. This is a missing experiment rather than a methodological flaw—the paper provides the infrastructure to conduct such a comparison but does not itself perform it. A complete evaluation would involve training models with MuonBP, block-diagonal Shampoo, and Canzona+Muon to a fixed validation loss and comparing total GPU-hours, thereby quantifying the "fidelity tax" (if any) of approximations versus the "throughput tax" (if any) of exactness.

7. Implications and Future Directions

How This Work Changes the Landscape

Canzona fundamentally reframes the relationship between optimizer algorithm design and distributed system architecture. Before this work, the prevailing assumption was that matrix-based optimizers and efficient distributed training were in fundamental tension—practitioners had to choose between mathematical exactness (preserving the optimizer's full update rule) and system throughput (exploiting ZeRO-1's coalesced Reduce-Scatter and asynchronous overlap). This tension was not merely an engineering inconvenience; it was treated as a structural constraint that forced two distinct research communities to pursue incompatible goals: algorithm designers developed increasingly sophisticated update rules (Shampoo, Muon, SOAP) assuming single-device execution, while systems builders optimized communication primitives assuming element-wise updates (AdamW). The resulting gap meant that convergence-optimal optimizers were often throughput-suboptimal to deploy, and throughput-optimal systems were unable to support the most advanced algorithms.

The paper's core reframing is that this tension is diagnosable and resolvable through a specific architectural diagnosis: the conflict arises not from any fundamental incompatibility between matrix operations and distributed execution, but from a coupling of two independent concerns—logical optimizer-task ownership (which rank computes the update for each parameter) and physical communication geometry (how data is sliced across ranks for efficient collectives). By decoupling these and re-introducing a monotonic ordering constraint, Canzona demonstrates that both exactness and throughput can be preserved simultaneously.

The magnitude of this contribution is a new design paradigm for distributed optimizer frameworks, not an incremental optimization. It is analogous to the shift from hand-tuned operator fusion to compiler-based kernel generation in deep learning compilers: it replaces ad-hoc, per-optimizer engineering with a principled, optimizer-agnostic scheduling abstraction. The paper's explicit connection to ZeRO-1 Geometric Compatibility creates a falsifiable criterion for evaluating future distributed optimizer proposals—any system that scrambles the monotonic rank-to-position mapping within contiguous parameter buffers will incur the 2× All-Reduce penalty documented in Figure 7, regardless of implementation quality. This transforms system design from empirical benchmarking ("try various heuristics and measure") to constraint verification ("check that your assignment preserves geometric alignment").

This work also reconciles the apparently contradictory findings in prior distributed training literature. On one side, practitioners observed that layer-wise partitioning could deploy matrix-based optimizers at scale (NVIDIA's layerwise_optimizer). On the other side, benchmarks showed that these deployments incurred substantial communication overhead that eroded convergence benefits. Canzona's diagnosis of the ZeRO Geometric Incompatibility (Appendix D.2) explains both observations simultaneously: layer-wise partitioning does preserve exactness (explaining the positive deployment reports), but it structurally cannot use efficient Reduce-Scatter (explaining the throughput penalties). The contradiction was not in the observations but in the implicit assumption that the throughput penalty was inherent to atomicity preservation. By showing that a geometrically-aligned static partition achieves the Reduce-Scatter baseline (Figure 7), the paper establishes that the penalty was specific to the layer-wise coupling, not to the atomicity requirement—and that an alternative architecture can avoid it entirely.

The research landscape shifts in several concrete ways:

  • Verifier/optimizer co-design becomes less critical for throughput. Prior to Canzona, a major research thrust was modifying optimizers to fit system constraints (block-diagonal approximations, shard-local orthogonalization, low-rank projections). These approaches treated the system as fixed and the optimizer as malleable. Canzona demonstrates that the opposite approach—treating the optimizer as fixed and the system as malleable—can achieve exactness without throughput sacrifice. This does not make algorithmic approximations obsolete (they may still be valuable for memory-constrained or latency-critical regimes), but it removes the necessity of compromising fidelity for large-scale training. The research frontier shifts from "how do we approximate this optimizer to make it fast?" to "how do we schedule this (exact) optimizer to minimize its overhead?", which is a fundamentally different and more general problem.

  • Optimizer-agnostic system design becomes a realistic goal. Before Canzona, each new matrix-based optimizer required bespoke distributed engineering—Muon needed MuonBP, Shampoo needed Distributed Shampoo, and deploying a new algorithm like SOAP at scale meant starting the system integration from scratch. Canzona's task abstraction (Section 4.3), where optimizer updates are treated as generic computational tasks defined by cost metrics, means that supporting a new optimizer requires zero system code changes—only the optimizer's step() function needs to be provided. This lowers the barrier to entry for optimizer innovation: algorithmic researchers can design new update rules without worrying about distributed execution, and systems researchers can optimize the scheduling framework without understanding optimizer internals. The paper validates this for three optimizers (Muon, Shampoo, SOAP) with distinct computational patterns, but the principle extends to any future algorithm that can be expressed as a MatrixOp on a tensor.

  • The "trilemma" (exactness, throughput, generality) becomes a solved problem in principle. The paper's Section 4.3 and Appendix E.4 explicitly frame Canzona as resolving a previously inescapable tradeoff. While the empirical validation is limited to Megatron on Qwen3 (a limitation discussed in Section 6.4 of the prior analysis), the architectural pattern—decoupled assignment, monotonic geometric constraint, offline load balancing, asynchronous execution, and task abstraction—provides a template that can be instantiated in other frameworks (FSDP, DeepSpeed) and for other model architectures. Future systems that violate this template will need to justify why they accept one axis of the trilemma as a tradeoff, rather than treating it as inevitable.

  • Offline static planning as a first-class design choice gains credibility. The paper's approach of computing partition maps and scheduling plans once during initialization (in milliseconds, Appendix D.1) and then executing them without runtime overhead challenges the prevailing preference for dynamic, runtime-adaptive scheduling in distributed systems. The success of this approach—achieving near-perfect load balance (FLOPs ratio 1.43× for DP, 2.46× for TP) without runtime monitoring or rebalancing—suggests that for the specific domain of transformer training with known model architectures, static planning is not merely adequate but optimal: it eliminates runtime scheduling overhead, enables fixed optimizer state placement (zero state transmission), and provides deterministic, reproducible performance. This pattern may influence other domains where the workload is known in advance and the cost of dynamic adaptation exceeds the benefit.

However, the landscape is shifted rather than fully transformed. The paper's experimental scope—one model family (Qwen3, a dense decoder-only transformer), one framework (Megatron), and three matrix-based optimizers—means that the generality of the paradigm remains to be demonstrated. The limitations identified in Section 6 (no memory footprint measurements, no comparison against algorithmic approximations, no dynamic architecture support, no cross-framework validation) define the boundaries of what this work establishes versus what it enables. The paper provides the system infrastructure and design principles that make such broader validation possible, but the validation itself is future work.

Follow-Up Research This Work Enables

End-to-end convergence-throughput benchmarking against algorithmic approximations. The paper demonstrates that Canzona achieves 1.57× end-to-end speedup over the best exact baseline (NV-layerwise) and 5.8× optimizer-step speedup, but does not compare against algorithmic approximations like MuonBP, block-diagonal Shampoo, or low-rank subspace methods that sacrifice mathematical fidelity for throughput. A definitive follow-up would train Qwen3-32B (or an equivalently scaled model) to a fixed validation loss using Canzona+Muon, MuonBP, and the Synchronous Compute baseline, measuring total GPU-hours to convergence. This experiment would directly test the paper's central claim that exactness can be preserved without meaningful throughput penalty: if Canzona+Muon achieves comparable or better total-time-to-convergence than MuonBP (because Canzona's exact updates require fewer training steps, offsetting any per-step throughput disadvantage), the case for exactness-preserving systems is strengthened. If MuonBP remains substantially faster end-to-end despite requiring more steps, the practical value of exactness is context-dependent. The paper's precision experiments (Figures 5, 10b, 11b) establish that Canzona's loss curves match the exact baseline, providing the fidelity baseline for such a comparison; the missing piece is measuring the step efficiency of approximations at the same scale.

Memory footprint characterization and optimization. Section 6.3 of the prior analysis identified that the paper reports load-balance ratios for memory (1.11× after α-Balanced DP partitioning) but never absolute per-GPU memory consumption compared to standard ZeRO-1 or NV-layerwise. A necessary follow-up would instrument Canzona to report peak GPU memory during training for Qwen3 models at multiple scales (1.7B through 32B) and compare against: (a) standard ZeRO-1 with AdamW (the memory floor for element-wise optimizers), (b) NV-layerwise with Muon (the exactness-preserving baseline), and (c) theoretical minimum memory given the optimizer state sizes. This would answer whether the atomic parameter assignment in Canzona's static partitioning creates per-rank memory pressure that could prevent training models that fit under standard ZeRO-1. If memory overhead is significant, a natural extension would be to incorporate memory capacity constraints directly into the α-Balanced optimization (Equation 2), adding a term that penalizes assignments where any rank's memory exceeds a hardware limit. The current formulation minimizes memory variance across ranks but does not constrain the maximum.

Cross-framework implementation and validation (FSDP, DeepSpeed). Appendix D.4 sketches how Canzona's principles could extend to PyTorch FSDP but acknowledges that hybrid FSDP+TP configurations introduce additional complexity. A concrete follow-up would implement Canzona's TP-style asynchronous compute pipeline (Micro-Group Scheduling with fused All-to-All and Host Rank assignment) within FSDP for the ZeRO-2/3 scenarios where gradients are sharded. The key engineering challenge is adapting the geometric constraint to FSDP's per-parameter FlatParameter sharding (which does not expose a contiguous-buffer-with-buckets abstraction) while preserving the coalesced communication properties. A successful implementation would validate the paper's claim of framework-agnostic design principles and provide a migration path for the large FSDP user base. The evaluation should replicate the Figure 4 comparison (end-to-end iteration time versus layerwise_optimizer-equivalent for FSDP) and the Figure 7 communication efficiency analysis (Fwd-Bwd latency versus AdamW Reduce-Scatter baseline) in the FSDP context.

Dynamic difficulty estimation and elastic scheduling for heterogeneous architectures. The paper's static offline planning assumes a fixed, known model architecture throughout training. For Mixture-of-Experts models, staged training (progressive unfreezing), or elastic scaling (changing GPU count mid-training), this assumption breaks. A research direction is to develop incremental rebalancing algorithms that can recompute the partition map Π after architectural changes without requiring a full buffer reallocation. One approach: when new parameters are added (e.g., unfreezing layers), treat them as a new "virtual bucket" and run the α-Balanced algorithm only on the new parameters, assigning them to ranks based on current deficits d_r without disturbing existing assignments. When GPUs are added or removed, formulate the rebalancing as a minimum-movement migration problem: compute the new optimal partition map for the new topology, then find the minimum set of parameter reassignments to transition from the old map to the new map, minimizing optimizer state transfer. The paper's discrete optimization formulation (Section 3.2) provides the mathematical foundation; the extension is to make it incremental and movement-aware.

Verifier-guided scheduling for optimizer-specific cost models. The paper adopts numel as a universal cost proxy, justified by the negligible performance difference from exact FLOPs-based scheduling (Figure 16, ~10⁻⁴ seconds) and by the optimizer-agnostic design principle. However, for optimizers with highly non-uniform complexity scaling (e.g., an algorithm that is O(d_in² · d_out) for rectangular matrices), numel may systematically misestimate cost for certain tensor shapes. A follow-up would develop a cost model calibration procedure: before training, profile the actual execution time of MatrixOp for a representative sample of tensor shapes from the model, fit a parameterized cost function (e.g., C(d_in, d_out) = a · d_in^p · d_out^q), and plug this calibrated function as \mathcal{W}(p) in Algorithms 1 and 2. The experiment would compare load-balance ratios and optimizer-step makespan using calibrated cost models versus numel for Shampoo and SOAP at scale (since these have more complex preconditioner construction than Muon's Newton-Schulz), quantifying whether calibration provides meaningful improvement. The paper's abstract formulation (using \mathcal{W}(p) without committing to numel) makes this extension straightforward to implement.

Stress-testing the geometric compatibility at extreme DP scales and slower interconnects. The paper's claim that communication imbalance from non-uniform shard sizes is "effectively hidden" by overlap (Section 3.3) is validated up to DP=128 on high-bandwidth interconnects (Appendix C.5, Figure 13). A stress-test would push this to DP=512 or DP=1024 on clusters with varying interconnect quality (e.g., 100 Gbps Ethernet vs. 400 Gbps InfiniBand vs. NVLink-only), measuring Fwd-Bwd time at α=1.0 versus α=0.0 to determine the DP scale at which communication imbalance outlasts the overlap window and becomes exposed. The experiment would also test whether the α parameter can be dynamically tuned—starting at α=1.0 and reducing to 0.5 or 0.0 if profiling detects that Fwd-Bwd time is increasing due to communication stragglers. This would establish the operating envelope for the aggressive load-balancing strategy and provide practitioners with a diagnostic for when to back off.

Canzona as a substrate for optimizer co-design research. The paper positions Canzona as an infrastructure layer that makes optimizer innovation independent of system engineering. This enables a new class of research: optimizer algorithms designed with awareness of Canzona's scheduling but without being constrained by it. For example, an optimizer could be designed to produce update rules that are computationally heavier but converge faster (reducing total steps), knowing that Canzona's load-balanced asynchronous execution will minimize the per-step overhead. Conversely, an optimizer could dynamically choose between an exact expensive update and an approximate cheap update based on a signal about current load imbalance—if a particular rank is already the straggler, use the cheap update for its assigned parameters to avoid falling further behind. Canzona's task abstraction (treating updates as cost-labeled operations) provides the interface for such co-design: the optimizer exposes a cost estimate, Canzona schedules to minimize makespan, and the optimizer can query scheduling decisions to adapt its computation. This is a more principled version of the ad-hoc approximations critiqued in Appendix E.3.

Practical Applications and Downstream Use Cases

Large-scale LLM pretraining with convergence-optimal optimizers. The most direct application is replacing AdamW with Muon, Shampoo, or SOAP in production pretraining pipelines without sacrificing throughput. The paper's results (Figure 4) demonstrate that on Qwen3-32B with 256 GPUs, Canzona reduces the optimizer step from 0.383 seconds (NV-layerwise) to 0.066 seconds, bringing the overhead to a level where the optimizer is no longer the dominant training cost. For a training run requiring, say, 500,000 iterations, this saves approximately 44 hours of wall-clock time (500k × 0.317 seconds saved per iteration) on the same hardware, or equivalently reduces the GPU-hours required to reach a target loss—assuming the matrix-based optimizer achieves comparable or better per-step convergence progress than AdamW (a claim validated in prior work but not in this paper). The practical benefit scales with model size: Figure 6 (bottom row) shows the gap widens for larger models (14B, 32B), making Canzona increasingly valuable as models grow. For organizations training 100B+ parameter models on thousands of GPUs, the cumulative savings are substantial enough to affect budgeting and scheduling decisions.

Cost-efficient hyperparameter exploration and optimizer comparison. The optimizer-agnostic design of Canzona means that switching between Muon, Shampoo, SOAP, or any future matrix-based optimizer requires no system code changes—only providing a different step() function. This dramatically reduces the engineering cost of optimizer ablation studies. A research team can run controlled comparisons of multiple optimizers at scale (e.g., training a 14B model with each optimizer for 50B tokens and comparing loss curves) without needing to develop, debug, and optimize distributed implementations for each one. The paper's Appendix C.4 validates this: the identical Canzona codebase achieved ~30× Shampoo optimizer speedup and zero-fidelity-loss convergence without any optimizer-specific tuning. This shifts optimizer selection from an engineering-constrained decision (which optimizers can we afford to implement at scale?) to a purely algorithmic decision (which optimizer converges best?), accelerating the pace of optimizer innovation and adoption.

On-premise and cloud-agnostic deployment of advanced training algorithms. Many organizations training LLMs operate on heterogeneous or shared GPU clusters where hardware topology (NVLink availability, interconnect bandwidth, node sizes) varies across allocations. Canzona's static planning can be recomputed for each specific topology at initialization time (in milliseconds, Appendix D.1), meaning the same training script automatically adapts its load balancing to the available hardware without manual tuning. The partition map Π depends on the number of ranks R and the parameter list; recomputing it for a new allocation is a one-time cost. For a cloud user who provisions a different number of GPUs for each training run, or for an on-premise cluster where node availability fluctuates, this adaptivity eliminates the manual effort of re-tuning sharding strategies and communication parameters. The α parameter (ablation in Figure 13) can be set once based on the cluster's typical overlap characteristics and left alone, or profiled at job start.

Self-improvement and distillation pipelines using exact optimizer updates. In self-improvement loops (e.g., ReST, STaR, or rejection sampling fine-tuning), a model generates training data, which is then used to fine-tune the model itself. The quality of the optimizer update in each iteration affects how quickly the model improves. Using an exact matrix-based optimizer (rather than an approximation like block-diagonal Shampoo or shard-local Muon) ensures that the model's parameters move in the true steepest-descent direction according to the optimizer's curvature model, potentially improving the stability and efficiency of the self-improvement process. Canzona makes this feasible by removing the throughput penalty that would otherwise make exact updates impractical in iterative pipelines where many fine-tuning rounds are needed. The paper's precision experiments (Figures 5, 10b, 11b) confirm that Canzona's updates are mathematically identical to single-device execution, so the self-improvement loop sees exactly the optimization trajectory the algorithm designer intended, without system-induced drift.

When to Prefer This Method

The paper explicitly positions Canzona against two named alternatives: system-level exact approaches (NVIDIA's layerwise_optimizer, Paradigm 2 in Section 3.1) and algorithmic approximation approaches (MuonBP, block-diagonal Shampoo, low-rank methods catalogued in Appendix E.3). Based on the paper's evidence and analysis, the decision framework is:

Prefer Canzona over layerwise_optimizer (NV-layerwise) when:

  • You are training within Megatron's ZeRO-1 + Tensor Parallelism paradigm and the model uses a standard transformer architecture (dense decoder-only or encoder-decoder) where parameter shapes are regular and the param_and_grad_buffer abstraction is available.
  • The optimizer step constitutes a non-trivial fraction of total iteration time. Figure 4 shows NV-layerwise spends 0.383 seconds in the optimizer step for Qwen3-32B (28% of the 1.381s total iteration); Canzona reduces this to 0.066 seconds (7.5% of 0.877s). The advantage scales with model size (Figure 6): at 1.7B the absolute gap is small, at 32B it is 8.3×.
  • You cannot tolerate any mathematical deviation from the optimizer's intended update rule. Canzona's precision experiments (Figures 5, 10b, 11b) demonstrate bit-exact convergence equivalence across 400B tokens for Muon, Shampoo, and SOAP. NV-layerwise also preserves exactness, but at higher throughput cost.
  • Your training regime is static: fixed model architecture, fixed GPU count throughout training. Canzona's offline planning assumes this; layerwise_optimizer also assumes static topology.

Prefer Canzona over algorithmic approximations (MuonBP, block-diagonal Shampoo) when:

  • Convergence fidelity is paramount—for example, training a production model where unexpected convergence degradation from approximations could require costly re-runs, or where the model will be fine-tuned and deployed in safety-critical applications where training trajectory reproducibility matters.
  • You are using an optimizer for which no mature approximation exists. Canzona's optimizer-agnostic design means any future matrix-based optimizer is immediately deployable at scale without waiting for the approximation literature to catch up.
  • You want to run controlled optimizer ablation studies where the only variable is the update rule, not the system implementation. Canzona provides identical system infrastructure for all optimizers, eliminating confounding system-level differences.

The comparison against algorithmic approximations remains an open empirical question. The paper does not measure the throughput of MuonBP or block-diagonal Shampoo against Canzona, so the tradeoff between convergence quality (where exact methods are provably faithful and approximations may degrade) and per-step throughput (where approximations may have an advantage, magnitude unknown) cannot be resolved from this paper's data. Until a head-to-head convergence-throughput benchmark is conducted (as proposed in Follow-Up Research above), the choice between Canzona and algorithmic approximations must be guided by the practitioner's tolerance for convergence risk versus the value of additional throughput.

Prefer standard ZeRO-1 with AdamW over Canzona with matrix-based optimizers when:

  • Your model architecture violates the assumptions of the numel-based cost model—for example, CNNs with heterogeneous spatial dimensions, graph neural networks with irregular tensor shapes, or architectures where parameter count and computational cost are poorly correlated. The paper validates numel on Qwen3's regular transformer structure (Figure 16) but provides no evidence for irregular architectures.
  • Your training regime requires dynamic architecture changes (staged unfreezing, elastic scaling, MoE routing) that invalidate the static partition map. Canzona would require recomputing Π and redistributing optimizer states, which is not supported in the current implementation.
  • You are memory-constrained to the point where any increase in per-GPU peak memory (even if balanced across ranks) would prevent training. The paper does not report absolute memory consumption, and the atomic parameter assignment could increase per-rank memory relative to standard ZeRO-1's equal-chunk sharding.
  • You operate at extreme DP scales (>512 ranks) on slow interconnects where the communication imbalance from α=1.0 might outlast the overlap window, eroding the Fwd-Bwd advantage. The paper validates up to DP=128; beyond that, the tradeoff is uncharted.