Distributed Training Techniques and Quantitative Analysis for Large Language Models (GPT‑175B)

The presentation is organized around four main topics: (1) the latest SOTA training techniques for transformer‑based large language models, (2) a quantitative analysis of these techniques using GPT‑175B as an example, (3) a detailed breakdown of memory, communication, and compute costs during model scaling, and (4) concluding thoughts on how to combine the optimizations effectively.

Large language models have grown dramatically in both data volume and parameter count, bringing challenges in computation, GPU memory, and communication. For GPT‑3/175B trained on 1.5 T tokens, a naïve estimate requires about 128 DGX‑A100 nodes (1024 A100 GPUs) for roughly 120 days of continuous training, ignoring practical overheads such as checkpointing, node failures, and debugging.

Memory consumption is split into three parts: model parameters (including optimizer states), activation tensors, and temporary buffers. Optimizer states dominate memory, occupying 16 × parameter‑size bytes. Distributed optimizer techniques like ZeRO‑1 partition optimizer states across data‑parallel ranks, reducing memory from ~22 GB to ~2.7 GB per GPU when DP size = 8.

Activation memory grows quadratically with sequence length. Full checkpointing stores only layer outputs and recomputes intermediate results during back‑propagation, cutting activation memory to 2 × batch × seq × hidden. Selective checkpointing further reduces recomputation overhead by targeting high‑cost ops such as self‑attention.

Parallelism strategies are illustrated with Megatron‑LM and NVIDIA NeMo frameworks. Tensor parallelism (TP) splits linear layers across GPUs, pipeline parallelism (PP) distributes layers across stages, and sequence parallelism (SP) partitions sequence‑dimension operations (e.g., LayerNorm) to lower activation memory. Combining TP, PP, and DP yields a communication hierarchy where DP handles gradient All‑reduce, TP performs frequent All‑reduce within a stage, and PP incurs the least communication.

Pipeline parallelism uses a 1F1B (one forward, one backward) schedule; warm‑up and cool‑down phases introduce bubble time, which can be reduced by interleaved pipeline scheduling or by increasing mini‑batch granularity, provided the batch is not so small that GPU kernel overhead dominates.

Distributed optimizer (ZeRO‑1) partitions optimizer states across DP ranks, while ZeRO‑2/3 further split gradients and parameters but increase communication and computation overhead; thus they are not recommended together with pipeline parallelism.

Practical recommendations: start with mixed‑precision (prefer BF16 for models > 20 B), enable FlashAttention, ZeRO‑1, and basic activation checkpointing. If memory pressure persists, add selective checkpointing, then gradually increase TP (ensuring hidden‑size/TP ≥ 1024) and PP. Avoid combining PP with ZeRO‑2/3 unless memory is extremely constrained.

The final takeaway emphasizes a step‑wise approach: enable default optimizations, then address memory bottlenecks with selective checkpointing and ZeRO‑1, followed by TP and PP scaling, and finally expand data parallelism when the batch size permits.