Sharded matmuls and collectives on a torus

Tick each objective you can deliver out loud without notes. Then write the one-paragraph explanation you would give a room of GPU engineers. Saved attempts stay below, so you can compare them over the week.

• I can read and write sharding notation, such as A[I_X, J_Y] and the unreduced marker {U_X}, and the equivalent JAX PartitionSpec. I can say what each device holds for a given mesh and spec.
• I can classify any sharded matmul into the four cases and name the collective each needs.
• I can derive the bidirectional-ring AllGather time V / W_ici and explain why it does not depend on the number of chips. I can extend it to several axes, to AllReduce (2×) and to AllToAll (¼).
• I can say when a collective is latency-bound, and compute the threshold on v5e and Ironwood.
• I can explain collective matmul (decomposing an AllGather into collective-permutes overlapped with partial matmuls) and compare it with async tensor parallelism on GPUs.
• I can relate the book's formulas to NCCL's algbw and busbw, and compare how quickly TPU and GPU collectives reach peak bandwidth.

Prompts to answer out loud

1. Why doesn't AllGather time on a ring depend on how many chips are in the ring?
2. Explain to a Megatron user what XLA's collective matmul does, and what they would have to do on a GPU to get the same overlap.
3. Given A[B, D] sharded P('X', 'Y') and W[D, F] sharded P('Y', None), what happens and what does it cost on an Ironwood 4×4×4 cube?
4. Why is an AllToAll a quarter of an AllGather on a TPU ring, and what does that mean for MoE?