• 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

• 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
    • Sequence Parallelism (2021): Focuses purely on distributing the sequence across devices, keeping full key-value blocks in memory on each device
    • Ring Attention (2023): Uses blockwise computation from Flash Attention AND ring communication pattern, meaning each device only needs to store small blocks at a time

  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

Mixture of depths

The key innovation is a method to dynamically allocate computation across tokens in a transformer model by letting some tokens skip certain transformer blocks entirely. Here's how it works:

1. Budget Setting: