大模型训练资源评估：算力与显存需求详解
大模型训练资源评估：算力与显存需求详解
大模型训练资源评估是AI领域的重要课题。本文详细探讨了大模型训练所需的GPU资源，包括算力需求、显存需求以及A100与A800的换算关系。通过具体的技术细节和计算公式，为AI从业者和对大模型训练感兴趣的读者提供了有价值的参考。
1. 算力需求
根据OpenAI发布的缩放定律，算力需求C的计算公式为：C = 6 * P * D，其中P为模型参数量，D为数据量。
1.1 语料+34B模型需要的浮点运算数
假设在增量预训练中，语料总量预计5B token：
D = 5B tokens = 5 * 10^9 tokens
34B模型表示其参数量：
P = 3.4 * 10^10
则总计算量（FLOPs = Floating point operations）：
C ≈ 6 * P * D = 1.02 * 10^22 FLOPs（来源：OpenAI缩放定律论文）
1.2 浮点运算数折算GPU-hours（单位：每块GPU每小时）
1.2.1 显卡利用率计算
根据LLAMA2论文数据，可以推算出Meta训练策略中34B预训练模型的显卡利用率。如下图所示：
LLAMA2-7B的计算量为 6 * 7 * 10^9 * 2 * 10^12 = 8.4 * 10^22 FLOPs
其理论GPU-hours为（8.4 * 10^22 / 312 * 10^12) * （1/3600）= 74800 GPU hours
论文中的实际时间为184320，可以计算其A100显卡利用率为 74800/184320 ≈ 40.5%
同理得到不同规模模型的显卡利用率LLAMA2-34B（35.0%）、LLAMA2-70B（43.5%）。
1.2.2 A100峰值算力折算
上述A100稳态算力约为峰值算力的0.35，A100显卡算力：
312 TFLOPS * 0.35 = 109.2 TFLOPS（Tera Floating point operations per second）
则GPU-hours的计算如下
(1.09 * 10^22 FLOPs) / (109.2 * 10^12 FLOPS) * （1/3600）= 27700 GPU小时
即1块A100需要27700小时（80块A100 80G 需要14.5天）计算完成。
三、显存需求
除了算力需求，GPU显存需求往往是制约训练的瓶颈。仅放下34b的模型启动训练需要4-6张A100（80G），主要包括：
- 将模型加载到GPU
- 保存优化器状态、加载优化器参数
- 存储Transformer模型激活值
- 中间梯度保存
Memory = 模型内存 + 优化器状态使用量 + 激活值使用量 + 梯度使用量
实测34B模型的权重占用显存为70GB。采用AdamW优化器以及混合精度训练，优化器状态参数与中间梯度共占用至少四倍以上显存，因此需要4-6张A100（80G）。
训练过程的激活值占用约40-60张A100（80G）（来源chatGLM团队、Eleuther AI比例估算值）。
激活值所消耗的额外显存随batchsize的大小而增加，在batchsize较大时，此部分额外显存消耗远远大于模型参数的消耗，采用激活重计算技术可将中间激活显存从O(n)减少到O(√n)。
显存上限决定了batchsize的上限。业界计算激活值占用显存通常采用以下公式：
Llama 7b模型层数32层，隐藏层大小4096，中间隐藏层大小11008，注意力头数32，词表大小32000，上下文长度2048，那么当batchsize为32时，激活值占用显存为：
32 * （14 * 32 * 2048 * 4096 + 4 * 32 * 2048 * 11008 + 8 * 32 * 32 * 2048 * 4096 + 5 * 32 * 32 * 2048 * 2048）= 3.1e12 = 3100GB
可估算得40-60张A100最高支持CodeLlama34b模型的batchsize为64。
考虑到多卡部署时梯度更新过程的通信损耗，共需约（4466）* 1.2 = **（5280）张A100（80G）**。
补充说明——A100与A800的换算关系及影响估算
从官方给出的技术规格上来看，在模型训练中最关注的显存与算力参数完全一致，80GB版本GPU显存带宽一致。唯一的区别在于NVLink互联桥的带宽下降至400GB/s，不满足美国的“每秒600GB及以上”的出口限制（NVLink——GPU间高速大带宽直连通信的互连技术）。
在分布式训练中多GPU间数据传输频繁，各训练加速框架例如accelerate、deepspeed等支持的ZeRO算法涉及梯度、优化器状态、模型权重参数甚至激活值的通信，数据传输速率的下降提高了通信时长。
卡间通信延长对训练时长的影响未有明确测算公式，根据现有实例推测：H800相对于H100也仅有卡间通信速率的下降（约50%），其在执行训练任务时，H800消耗比H100多10%-30%的时间。
A800相对A100在数据传输速率有33%的下降，预计约有5%的训练耗时增加。
因此执行增量预训练任务时，鉴于数据规模较小且训练周期较短，认为可以与A100进行1：1的等量换算。