Parallelism and Memory‑Optimization Techniques for Distributed Large‑Scale Transformer Training

1. Introduction Since the advent of the Transformer architecture, model sizes have exploded to trillions of parameters, making single‑GPU training infeasible. Distributed training across multiple GPUs and nodes is therefore essential, but it introduces challenges such as high memory consumption, long training times, and throughput constraints. The article outlines how a combination of parallelism strategies and memory‑optimization techniques can address these issues.

2. Related Technical Foundations

2.1 Parallelism Techniques The main forms of parallelism used in large‑model training are:

Data Parallelism (DP): each worker holds a full model replica and processes a distinct mini‑batch, synchronizing gradients across workers.

Pipeline Parallelism (PP): the model is split layer‑wise into stages that run on different devices, allowing overlapping of forward and backward passes.

Tensor Parallelism (TP): tensors inside a layer are partitioned (row or column) and computed on separate devices, requiring inter‑device communication.

Sequence Parallelism (SP): long sequences are divided into sub‑sequences processed in parallel.

Context Parallelism (CP): an enhanced SP that splits all input and output activations along the sequence dimension.

Figures illustrating data parallelism, pipeline stages, tensor row/column splits, and the combined TP‑SP‑CP configurations are included.

2.2 Memory‑Optimization

2.2.1 Memory Consumption Analysis Memory usage in training stems from model weights, gradients, optimizer states, and activations. The article provides formulas for estimating memory footprints under different precisions (FP32, FP16, BF16, INT8) and parallelism settings.

2.2.2 Optimization Techniques Two primary methods are discussed:

Activation recomputation: only store inputs and recompute intermediate activations on demand, reducing activation memory at the cost of extra compute.

ZeRO (Zero Redundancy Optimizer): partitions weights, gradients, and optimizer states across GPUs, offering three levels of redundancy reduction.

Figures show the memory savings achieved by TP, SP, and ZeRO at various levels.

3. QLM Acceleration Practice QLM (Qihoo Language Model) is a customized framework built on Megatron‑LM that supports model conversion, pre‑training, evaluation, and fine‑tuning. Using the described parallelism and memory‑saving techniques, QLM accelerates 128k‑length text fine‑tuning from 120 s/sample to 35.5 s/sample (≈3.4× speed‑up). The article details hardware configurations (H800 vs. A100), optimization steps (enabling DP, TP, PP, SP, CP), and performance results across several version iterations, highlighting trade‑offs such as OOM risks and recomputation overhead.

4. Conclusion The paper summarizes how parallelism and memory‑optimization jointly enable efficient training of trillion‑parameter models, discusses the three major bottlenecks (memory wall, communication wall, compute wall), and emphasizes ongoing research to further reduce resource consumption and improve scalability.

References

[1] Megatron‑LM Sequence Parallel training. [2] Narayanan et al., Efficient large‑scale language model training on GPU clusters using Megatron‑LM, 2021. [3] Korthikanti et al., Reducing activation recomputation in large Transformer models, MLSys 2023. [4] NVIDIA documentation on Context Parallelism. [5] Rajbhandari et al., ZeRO: Memory optimizations toward training trillion‑parameter models, SC20.