Sharded matrix multiplication and result ownership

The mechanism

A compiler must preserve both a tensor’s values and who owns them. Splitting matrix output rows gives different ranks independent output rows; splitting the contracted inner dimension gives each rank only part of every output sum. Those are different meanings of a local matrix product, even if the local arrays have similar shapes. JAX’s explicit sharding example forms local products and then reduces their partial sums. A sum all-reduce makes the completed result available at every rank. Reduce-scatter instead finishes the sum while distributing different output pieces, and all-gather can later assemble those pieces when a consumer needs a replica. Concatenating unfinished partial sums does not perform their missing arithmetic. This distinction connects tensor-axis information in a compiler to the collective scheduled over the network. It also prevents allocating replicas merely because the preceding stage happened to use all-reduce. A valid next operator may consume distributed rows directly, while another needs a different layout and an explicit conversion. The example below uses exactly representable small integers; its equality does not imply arbitrary floating-point reduction trees are bitwise identical.

1. Declare C = AB, the complete tensor shapes, stored element size and required consumer ownership before partitioning.
2. Split the inner dimension into disjoint sets Kᵣ; each rank receives matching columns of A and rows of B.
3. Compute Cᵣ locally. Each Cᵣ covers all output coordinates but contains only the products indexed by Kᵣ.
4. Add contributions with matching output coordinates. Use all-reduce for replicated C or reduce-scatter for completed disjoint output shards.
5. Retain row shards when the next operation can use them; all-gather only when the consumer requires the whole tensor.
6. Check every output against a separate dense reference, and count a sent payload once at its chosen link boundary. Input distribution, temporary buffers and protocol overhead are additional costs.

Science and assumptions

Engineering example

Use A = [[1,2],[3,4]] and B = [[5,6],[7,8]]. Rank 0 owns k=0 and forms C₀ = [[5,6],[15,18]]; rank 1 owns k=1 and forms C₁ = [[14,16],[28,32]]. Reduce-scatter with row ownership sends C₀ row 1 [15,18] from rank 0 to rank 1, and C₁ row 0 [14,16] in the opposite direction. Rank 0 adds [5,6]+[14,16]=[19,22]; rank 1 adds [15,18]+[28,32]=[43,50]. Each sends 8 payload bytes if elements are float32. If both consumers need all of C, all-gather exchanges these completed rows for another 8 bytes sent per rank. Both then hold [[19,22],[43,50]]. Total sent payload is 16 bytes per rank, 32 across the two senders; counting the matching receives again would double-count this network boundary. Concatenating the two original 2×2 partial matrices would create four rows and would leave their sums unfinished.

Distinctions and tradeoffs

A shard can be an unfinished contribution to every output, not necessarily a disjoint piece of the finished output.

A shard-preserving collective may save result replication, but a later consumer can reintroduce an all-gather or resharding cost. Logical bytes alone do not predict network latency.

Distributed linear layers become correct only when the compiler and runtime finish the same tensor operation and deliver its expected ownership. Avoiding an unnecessary replica can also save memory and communication.

Specifications and interfaces

Contracted-axis split

Partial sums overlap in output coordinates

Requires addition, not concatenation.

Output-row split

Disjoint completed rows

Concatenation can reconstruct the global output without summing corresponding rows.

Consumer contract

Replicated or sharded result

The useful collective is determined by the required layout.

Interfaces

Compiler sharding plan Input

Tensor shapes, partitioned axes, dtype and input ownership

Consumer-ready output Output

Completed values with explicit per-rank coordinates

Check your understanding

Evidence and scope

The sources establish sharding and collective semantics. All matrices, the two-row exchange schedule and byte counts here are synthetic. This is not an observed NCCL implementation, a topology model or a performance comparison.

The ideal ring all-reduce lab varies communication payload, ranks, startup and bandwidth. It does not execute these matrix shards or alter the optimizer; the exact computation is worked through in this lesson.

Sources

• JAX — Manual parallelism with shard_map
• NVIDIA NCCL — Collective Operations

Published lesson data · Atlas releases and licensing

Continue exploring

Foundations

• Compiler layout, lowering and buffer planning
• Collective communication library
• Data, tensor, pipeline, and state parallelism

Related lessons

• Systolic arrays and tensor dataflow
• Collective traffic and congestion
• Runtime, driver and device dispatch
• Floating-point reductions and compiler permissions