AI and Machine Learning

Distributed Training

Splitting a model and its optimizer state across devices once neither fits on one.

Colour is the family; a dashed line is the second member of it.

Drag to pan · scroll to zoom · click a node to open it
G all-reduce All-Reduce data-parallelism Data Parallelism data-parallelism->all-reduce gradients must be summed across replicas every step fully-sharded-data-parallel Fully Sharded Data Parallel zero-sharding ZeRO Sharding fully-sharded-data-parallel->zero-sharding shards parameters as well as optimizer state, per layer pipeline-parallelism Pipeline Parallelism tensor-parallelism Tensor Parallelism pipeline-parallelism->tensor-parallelism splits between layers rather than inside one tensor-parallelism->data-parallelism one replica of the model must still fit on one device zero-sharding->data-parallelism every replica redundantly stores the same optimizer state gradient-checkpointing Gradient Checkpointing zero-sharding->gradient-checkpointing shards state across devices instead of recomputing it

6 nodes

Pipeline Parallelism

Assigns consecutive layers to different devices and streams micro- batches through them. Communication is only at stage boundaries, at the cost of bu… · 2018

Tensor Parallelism

Splits individual weight matrices across devices so one layer is computed jointly. Communication lands inside every layer, so it wants the fastest in… · 2019

ZeRO Sharding

Partitions optimizer state, gradients and parameters across data- parallel ranks instead of replicating them, gathering each shard only while it is n… · 2019

Fully Sharded Data Parallel

The sharded-everything form of ZeRO as a framework primitive, gathering each layer's parameters just before use and releasing them straight after. · 2023

All-Reduce

A collective that sums a tensor across every participant and returns the result to all of them. Ring and tree schedules make its cost depend on inter…

Data Parallelism

Every device holds a full replica and processes a slice of the batch, then the gradients are summed. The simplest way to scale, and it stops working…