See also

• Megatron-DeepSpeed

General resources

• https://lilianweng.github.io/posts/2021-09-25-train-large/
• https://github.com/rwitten/HighPerfLLMs2024 - Rafi’s excellent series based in Jax

DDP: simple, but requires full model on each node

Model parallelism: ambiguous term, refers to either pipeline or tensor, usu. tensor

Thinking about efficiency

• Training: 6PD (for the 3 matmuls in fwd + bwd) is our expected TFLOP. Compute 6PD/time for a single step and see what percent of device’s max TFLOP/s this is (e.g. TPU v4 has 275TFLOP/s)
• Inference throughput is based on how many sequences you can parallel process per chip. E.g. device memory (TPU v4-8 has 32GB*4=128GB) = weights (7B params = 14GB) + 0.94GB per sequence’s KV (for Gemma) * num sequences. Solve for num sequences ~= 121. For margin of error, round down to 110 which totals 117GB not the full 128GB. So load 117GB to generate 110 tokens!
• Balancing sharding (source)
  • Training: usu FSDP with a bit of tensor parallel (FSDP can scale out a lot, but at some point you’re consuming too many tokens for a single step)
  • Inference: usu tensor parallel—FSDP loads all params over ICI, and in inference we’re already bottlenecked on loading all params from HBM so loading from ICI would be even slower

Pipeline parallelism: inter-layer

• Split batches into microbatches
• Various algorithms that minimize time bubbles

• DeepSpeed does this:

Tensor parallelism: intra-layer, chatty

• Clever: the MLP has 2 GEMMs. There are 2 ways to split A for the $XA$ mult: columns or rows. Splitting rows requires sync-summing before nonlinearity activation. Cols produces complete elements per half, but requires full input token. So, use cols for the first GEMM and rows for the second, saving the sync (g in the figure below) for the end!
• Attention: simply parallelize by head.

• Piece by piece (source)

  tensor parallel.mov

3D parallelism: DP, MP, PP combined

• Maybe helpful:

  

ZeRO/FSDP: partitions optimizer state, weights, gradients

• This is a form of data parallelism

• Partitioning weights may interfere with MP/PP?

• Animations (source)

  fsdp.mov

  fsdp overlap.mov

• Stage 1: just optimizer state

• Stage 2: gradients

• Stage 3: weights, very chatty?

https://youtu.be/kW-uQ-1pm38

Sequence parallelism aka context parallelism (e.g. in DS Ulysses): partition the input sequence. (paper)

• Doesn’t partition any model memory, so more similar to data parallelism (pure compute parallelism / chopping up of inputs and activations), but with the communication of tensor parallelism for attention (since all-to-all communication needed)

• Ring self-attention algorithm to compute attention by passing around keys then values in a circle

  • Video explainer, in-depth (haven’t watched)
  • Kinda blockwise, but not as blockwise as Ring Attention?
  • Algo
    • Each partition starts off with just its own tokens’ Q/K/V embeddings.
    • Each partition computes its own A=QK’
    • Each partition then sends its K to its neighbor, so that neighbor can compute more tokens’ attentions.
    • Each partition then repeats to get final embedding=AV

  

  

  

• Actually overcomputes (explainer)—see red not needed below

  

Ring attention

• Like ring self-attention introduced in sequence parallelism work, but uses block-wise

The use of a ring topology for computing self-attention has also been studied in prior work [21] but it incurs non-overlapped communication overheads similar to sequence parallelism, making it infeasible for large context sizes …

Prior work has also proposed leveraging a ring topology to compute self-attention [21], aiming to reduce communication costs. Our work differs by utilizing blockwise parallel transformers to substantially reduce memory costs. As we show in the next section, this enables zero-overhead scaling of context size during both training and inference and allows arbitrarily large context size. …

Prior work extends sequence parallelism for computing self-attention using a ring topology [21], which reduces the communication cost compared to standard sequence parallelism. However, overlapping communication with computation remains challenging due to the constraints of arithmetic intensity. The communication overheads render this approach infeasible for training and inference in largecontext scenarios. Our work leverages on blockwise parallel transformers to distribute blockwise attention and feedforward across devices and concurrently overlaps the communication of key-value blocks in a circular of hosts with the computation of query-key-value blocks and feedforward, reducing memory cost substantially and allowing device count times larger context size with zero overheads.

• Similar to blockwise attention (BlockBERT)—also computes blockwise, but distributed/parallel

MoE: mixture of experts (foundations)

• Training the router: the learning process is smeared over time if you are switching to just one or another.
  • Router receives gradients based on how well the chosen experts performed for each token.
  • If an expert performed well (reduced the loss), the router will receive a signal to increase the probability of choosing that expert for similar inputs in the future.
  • If an expert performed poorly, the router will receive a signal to decrease the probability of choosing that expert for similar inputs.
• Sparse Gradients and Straight-Through Estimator TODO
  • The top-k selection process is not differentiable. To allow gradients to flow, techniques like the Straight-Through Estimator are often used (see Machine learning foundations ). This allows the router to receive approximate gradients even though the selection process is discrete.
• Upcycling
  • Problem: random initialization means you are learning a weight and simultaneously learning a router over crappy experts
  • Instead, start training 1 MoE, then after say 10k steps, duplicate it
  • Not too many, since otherwise it’s become too engrained in its ways and no longer serves as a good general springboard for training diverse experts
  • Very briefly in video

More on MoE routing schemes

• top-k routing
  • Uses a network to learn routing
  • Usu. k is small to maintain sparsity
• Switch Transformers: like top-1 routing, but with just a softmax over a scaled dot-product of router parameters and token embeddings, rather than a full network
  • Aux loss: $\displaystyle n{\sum_{{{i}={1}}}^{{n}}}{{f}{{i}}{P}{{i}}}$ where $f_i$ are $1/n$ times frequency of terms going to expert i, and $P_i$ are $1/n$ times weight of terms for expert $i$, and ideally both are $1/n$
• Expert choice routing
  • Routers can be imbalanced
  • Instead, let each expert look at all tokens and choose its own top k somehow
    • Lets multiple experts process same token - feature not bug
    • If token chosen by nobody? Schemes: default/fallback expert, pad under-capacity expert, use all experts and average
  • But at inf time, you have to replace this with something else, maybe one of the other per-token routings
  • Very briefly in video
  • Heard xAI didn’t like this
• Sinkhorn routing https://arxiv.org/pdf/2202.01169.pdf
  • Like top-k, start by computing a matrix R of size (tokens, experts) → softmax, but don’t choose top-k now, and iteratively normalize rows and columns (sum to 1), ensuring balanced assignment
  • MOE review https://arxiv.org/abs/2209.01667
  • What if we didn’t use load bal https://arxiv.org/pdf/2103.16716.pdf
• RL routing
  • How it works
    • State Representation:
      • The state typically includes information about the current token, context, and possibly the load status of experts.
    • Action Space:
      • Actions correspond to selecting one or more experts for the current token.
    • Policy Network:
      • A neural network that takes the state as input and outputs a probability distribution over actions (experts).
    • Reward Function:
      • Defines the objective. Could include factors like:
        • Task performance (e.g., language modeling perplexity)
        • Load balancing
        • Computational efficiency
    • Training Process:
      • The policy is updated using RL algorithms (e.g., REINFORCE, PPO) to maximize expected rewards.
  • Advantages
    • Adaptive Behavior:
      • Can potentially learn complex routing strategies that adapt to different input patterns.
    • Multi-objective Optimization:
      • Can balance multiple objectives (performance, efficiency, load balancing) through reward design.
    • Long-term Planning:
    • Online Learning:
  • vs. Top-k: More flexible and potentially more powerful, but also more complex.
  • vs. Expert Choice: Can learn a centralized policy that still adapts to different inputs.
  • vs. Sinkhorn: More adaptive, but lacks the mathematical guarantees of balanced assignment.

MoE expert parallelism

How parallel groups work, worked out

Offload, ZeRO Offload