Transformer

Simplicity buried in Abstractions I’ve always loved the “Transformer Accounting” diagram from the JAX Scaling Book. It did a brilliant job of making the tensor shapes of a Transformer intuitive on a single device. But as we scale up, the complexity shifts. We stop worrying about just matrix dimensions and start worrying about the ‘alphabet soup’ of N-D parallelism (DP, TP, SP, CP, EP). Here is the irony: The core ideas behind these parallelisms are actually fundamentally easy. Conceptually, we are just decomposing a global tensor operation into local tensor compute chunks connected by communication collectives. It’s like an assembly line: instead of one worker building the whole car, we have a line of workers (GPUs) passing parts (tensors) back and forth. ...