Training memory and parallelism planner
How much memory each GPU needs to train a model under data, tensor, pipeline, expert and context parallelism with ZeRO 0–3 / FSDP and activation recomputation, what each axis sends over the network, the pipeline bubble with a 1F1B timeline, MFU vs HFU, an animated activation-checkpointing walk-through, and the memory and FLOPs of PPO, GRPO and DPO.
Added
How to use it
Five tabs:
-
Planner — pick a model (or type its size) and a GPU, then set the degrees of data (DP), tensor (TP), pipeline (PP), context (CP) and expert (EP) parallelism. The planner adds up what one GPU holds: 16 bytes per parameter of model state (BF16 weights and gradients, FP32 master weights and Adam’s two moments), sharded by ZeRO-1, -2 or -3 / FSDP over the data-parallel ranks, plus activations for the micro-batches the pipeline schedule keeps in flight, under no checkpointing, √n checkpoints, selective or full recomputation. A table compares the four ZeRO stages for the same layout, and another lists what each parallel axis sends per step, whether it fits in the NVLink domain, and how long it would take if nothing overlapped it.
-
Pipeline bubble — a timeline of GPipe or 1F1B on p stages with m micro-batches, the bubble (p − 1)/(m + p − 1), and how many micro-batches of activations each stage holds.
-
MFU & HFU — model FLOPs utilisation from one measured step (6·N·D ÷ (peak × chips × seconds)), from tokens per second, or from a reported FLOP rate, optionally with attention FLOPs; HFU adds the recomputed FLOPs.
-
Checkpointing — step through (or play) the forward and backward passes of an n-layer network checkpointed every k layers: which activations are kept, which are discarded and recomputed, and how many are held at each moment, against the no-checkpointing line.
-
RLHF — the resident memory of PPO (actor, critic, reward, reference), GRPO (no critic, and no reward model with a rule-based reward) and DPO (policy and reference), and the FLOPs and time share of one RL step: rollout generation, scoring, and training.
The accelerator figures (memory and dense BF16 peak, each with its source) are in one table at the bottom of the page.
Limits
Every figure is an estimate. Model state assumes mixed-precision Adam with no extra buffers (real runs add gradient-accumulation and communication buffers: 18–20 B/param). Activations follow Korthikanti et al. (2022) with sequence parallelism and leave out the embedding and output layers. TP and PP divide the parameters evenly; MoE experts are also split by EP, with the expert parallel group carved out of the DP × CP ranks. Communication volumes are ring-algorithm bytes sent per GPU per step; ZeRO-3’s per-micro-batch re-gathers and overlap with compute are not modelled. The pipeline timeline treats every stage as equally fast. RLHF FLOPs use 2·N per token for inference and 6·N for training, and leave out the rollout engine’s own weight copy and KV cache.
Related: AI hardware calculator (training compute and time, LLM memory, roofline), failures, checkpoints and goodput, floating-point inspector (BF16, FP8 and the block-scaled formats).
Companion to AI Acceleration: From Transistors to Trillion-Parameter Models: parallelism topologies, peak FLOPs, MFU and MLPerf, parallelism as an optimization problem, ZeRO, FSDP and the optimizer-memory arithmetic, activation checkpointing and the infrastructure of post-training.
Sources
-
Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2019) — 16 bytes per parameter, the three stages and their communication.
-
Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021) — 3D parallelism, the interleaved schedule, 502 PFLOP/s on 3,072 A100s.
-
Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models (2022) — activation memory, selective recomputation.
-
Chen et al., Training Deep Nets with Sublinear Memory Cost (2016) — √n checkpointing.
-
Huang et al., GPipe (2019) and Harlap et al., PipeDream (2018) — micro-batching and 1F1B.
-
Chowdhery et al., PaLM (2022) — MFU and HFU (46.2% / 57.8%; 238.3K tokens/s on 6,144 TPU v4).
-
Meta, The Llama 3 Herd of Models (2024) — the 4D layout (TP 8, CP 1 or 16, PP 16, DP 128 or 8).
-
Schulman et al., PPO; Rafailov et al., DPO; Shao et al., DeepSeekMath (GRPO); Hu et al., OpenRLHF.
Linking to an example
Links can be written by hand. Add the tab after #: plan, pipeline, mfu, ckpt or rlhf.
Parameters you leave out take their defaults. Numbers may be written 4e6, 4000000 or
4,000,000. In a URL + means a space, so write a plus sign as %2B (D=4e%2B6); a 4e+6
typed straight into the address bar also works, because the tool reads 4e 6 back as 4e6.
| Parameter | Meaning | Example |
|---|---|---|
|
Planner model: |
|
|
Custom model: parameters (billions), layers, hidden size, attention heads, KV heads |
|
|
|
|
|
Sequence length, micro-batch and global batch (sequences) |
|
|
Parallel degrees |
|
|
|
|
|
Activations: |
|
|
FlashAttention ( |
|
|
Planner schedule |
|
|
GPUs per NVLink domain; NVLink and scale-out GB/s per GPU |
|
|
Pipeline tab: stages, micro-batches, |
|
|
MFU tab: |
|
|
Parameters (billions), tokens in the step, step seconds |
|
|
Tokens per second, or achieved petaFLOP/s |
|
|
Accelerator, dense peak TFLOPS per chip, number of chips |
|
|
Count attention FLOPs ( |
|
|
Recomputed FLOPs %, for HFU |
|
|
Checkpointing tab: layers, checkpoint every k, labels |
|
|
Start the walk-through at this step |
|
|
RLHF: |
|
|
Policy and reward-model sizes (billions); |
|
|
Prompts per step, samples per prompt, prompt and completion tokens, epochs |
|
|
Decode and training MFU % |
|
|
One of the presets below; parameters given with it override the preset’s |
|
Presets:
example= |
Shows |
|---|---|
|
70B × 16 bytes = 1.12 TB of training state on one GPU (open) |
|
70B over 64 data-parallel GPUs with ZeRO-3 / FSDP; the table shows ZeRO-1 ≈ 4×, ZeRO-2 ≈ 8×, ZeRO-3 = 64× at 1.5× the traffic (open) |
|
The ZeRO paper’s 7.5B on 64 GPUs: 120 → 31.4 / 16.6 / 1.9 GB (open) |
|
Llama 3 405B: TP 8 × PP 16 × DP 128 = 16,384 GPUs (open) |
|
The same at 128K context with CP 16 (open) |
|
Megatron-LM’s 1T model: TP 8 × PP 64 × DP 6 on 3,072 A100s (open) |
|
DeepSeek-V3 with EP 64: expert all-to-all (open) |
|
4 stages, 8 micro-batches, 1F1B: bubble 3/11 = 27% (open) |
|
The same with GPipe: same bubble, 8 micro-batches of activations per stage (open) |
|
8 stages, 64 micro-batches: bubble ≈ 10% (open) |
|
6 · 70e9 · 4e6 ÷ (989e12 · 512 · 7 s) = 47.4% (open) |
|
502 PFLOP/s ÷ (3,072 × 312 TFLOPS) = 52% (open) |
|
PaLM 540B: 46.2% MFU (45.7% without attention FLOPs), 57.8% HFU (open) |
|
The book’s 9-layer network checkpointed every 3 layers (open) |
|
The same, labelled by layer inputs (open) |
|
100 layers every √100 = 10: 20 tensors held instead of 101, +33% compute (open) |
|
PPO: four resident 70B networks, 36 bytes per policy parameter (open) |
|
GRPO: 3 networks instead of 4, ≈44% fewer bytes (open) |
|
GRPO with a rule-based reward, as DeepSeek-R1 (open) |
|
DPO: policy + reference, no rollout (open) |
|
Long reasoning rollouts: a fifth of the FLOPs, over 70% of the time (open) |
Feedback is a gift. Something to improve, a feature to add, or a bug? Tell me.
Runs entirely in your browser; nothing you type leaves this page.