LLM 训练显存估算器

一个简化版估算器，用来观察 Transformer 规模、batch、序列长度、TP/DP/ZeRO 对单卡显存的影响。

estimating

Preset Layers

Hidden size

Sequence length

Micro batch size

Tensor parallel

Data parallel / ZeRO shard count

ZeRO stage Activation checkpointing GPU memory per card (GB)

Approx parameters

Per-GPU training memory

Global batch tokens

Rough train FLOPs / step

Parameters

Gradients

Optimizer

Activations

Approximation: parameters ≈ layers × 16 × hidden² + 2 × vocab × hidden (vocab = 128k). Parameter and gradient tensors use bf16/fp16 bytes; Adam states are approximated as 12 bytes per parameter before ZeRO sharding. Activation memory is intentionally rough: layers × sequence × micro-batch × hidden × 2 bytes × 6 × checkpointing factor.