提出4维并行内存估算公式,提前避免训练溢出。
Accelerating Large Language Model Training with 4D Parallelism and Memory Consumption Estimator
- 基于Llama架构推导4维并行内存消耗公式
- 实验验证80%内存利用率下无溢出风险
- 助力高效配置搜索,适合大模型训练优化者
在大规模语言模型训练中,张量并行(TP)、流水线并行(PP)、数据并行(DP)以及上下文并行(CP)等多维并行策略被用于分布模型参数、激活值和优化器状态。为在不同环境下确定最优并行配置并避免显存溢出,本文针对Llama架构,给出了4维并行(DP, TP, PP, CP)下参数、梯度、优化器状态和激活值的精确内存消耗估算公式。我们在A100和H100 GPU上进行了454次实验,纳入了临时缓冲区和内存碎片等常被忽略的因素。结果显示,当估算内存使用低于可用显存的80%时,训练从未发生显存溢出。该简单有效的公式可提前识别可能导致溢出的配置,显著缩小配置搜索空间。通过对454组实验结果的全面分析,本研究还提供了4维并行配置的实证优化建议。
原文摘要 · Abstract (English)
In large language model (LLM) training, several parallelization strategies, including Tensor Parallelism (TP), Pipeline Parallelism (PP), Data Parallelism (DP), as well as Sequence Parallelism (SP) and Context Parallelism (CP), are employed to distribute model parameters, activations, and optimizer states across devices. Identifying the optimal parallelization configuration for each environment while avoiding GPU memory overflow remains a challenging task. In this study, we provide precise formulas to estimate the memory consumed by parameters, gradients, optimizer states, and activations for 4D parallel training (DP, TP, PP, CP) in the Llama architecture. We conducted 454 experiments on A100 and H100 GPUs, incorporating often neglected factors such as temporary buffers and memory fragmentation into our analysis. Results indicate that when the estimated memory usage is below 80\% of the available GPU memory, the training never encounters out-of-memory errors. This simple yet effective formula allows us to identify parallelization configurations that could lead to memory overflow in advance, significantly reducing the configuration search space. Additionally, through a comprehensive exploration of optimal configurations in 4D parallelism, our analysis of the 454 experimental results provides empirical insights into optimal 4D parallelism configurations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。