How to read this

One training step of a 45-layer MoE on 1,024 GPUs, drawn as a city: 4 pipeline stages, each 64 nodes of 4 GPUs.

• A square is a GPU: 4 per node, 64 nodes per stage
• Blue: a micro-batch's forward pass
• Orange: its backward pass
• Blue and orange: one of each at once
• Grey: summing gradients, then the optimizer
• A dot: a token to its experts and back
• A line: activations or gradients, stage to stage
• Hatched: idle time, the pipeline bubble
• The lens: one node up close, the micro-batches it runs

Drag the schedule to scrub. Simulated timings.