How to split a model across chips
Laying out a dense model for prefill and decode, through the Orbit cluster
Slide contents
1. How to split a model across chips
Laying out a dense model for prefill and decode, through the Orbit cluster
2. The layout sets latency and cost
How to cut a dense model across chips, and when to change the cut
3. From one GPU to a GPU group
Continuing the vLLM deck (episode 32) and the inference-systems longread
4. A dense 70B exceeds one H100's memory
Orbit: 16 nodes × 8 H100 80 GB, NVLink inside a node, 400 Gb/s InfiniBand between nodes
5. Small batches are weight-bound, large ones KV-bound
Memory time, compute time and chip-to-chip communication
6. Prefill and decode are different workloads
Latency, throughput, MFU and cost in chip-seconds per token
7. Every layout is a set of collectives
BLE_xyz notation, all-gather, reduce-scatter and all-to-all on a TPU v4 torus
8. 1D layout hits a constant communication floor
FFN weights split along one axis, as in Megatron; activations are gathered whole
9. In 2D, communication falls as 1/√n
Weights split along both E and F; PaLM 540B, decode, batch 512
10. With large batches, moving weights pays off
Weight-gathered: PaLM 540B, prefill, 64 chips, length 2,048
11. Multiquery saves memory when sharded by batch
One K/V head for all query heads; decode by batch, prefill by heads
12. Same memory, 32× longer context
PaLM 540B, 64 chips, 30% of memory for the KV cache
13. Parallel block drops one of two all-reduces
Fused matrices, overlapped communication and int8 weights
14. Goals change batch, layout and weight format
PaLM 540B, 64 TPU v4, context 2,048; prefill 2,048 tokens, decode 64
15. Low latency costs more per token
Pareto frontier of cost versus decode latency, PaLM 8B–540B
16. 64-way TP kept 44% MFU
Against FasterTransformer on A100, by MFU
17. Findings are tied to TPU v4's torus
Uniform network, static batches and PaLM's architecture
18. DistServe: separate pool and layout per phase
The paper's split-server idea and DistServe (2024)
19. TP over InfiniBand eats the prefill budget
1D-layout formula for a 70B: d_model 8,192, 80 layers, bf16, 50 GB/s
20. Orbit: 70B copies stay inside a node
The layout depends on the phase and tokens per batch
21. Only measurement shows collective latency
A TP × batch × prompt-length grid on one node and on two
22. Formulas transfer; numbers do not
What the paper showed, what we transfer and what needs testing
23. Layout follows the phase, not the model
How many tokens per batch, where is the fast network, what stays put?