必知必会：大模型训练显存计算与优化详解

显存消耗的组成与模型权重计算1.1 核心问题大模型训练时显存被什么占满了不同量化精度下模型权重需要多少显存1.2 原文核心要点深度神经网络训练的显存消耗主要包括两大部分模型状态模型权重、梯度、优化器状态和激活值各个非线性模块的中间激活值。不同量化精度下的显存占用差异巨大。1.3 显存消耗的两大组成部分换句话说显存就像你的工作台空间一部分放置工具箱和材料模型状态一部分作为临时加工区激活值。前者大小固定后者随工作量波动。组成部分具体内容说明模型状态模型权重参数、梯度、优化器状态与模型参数量Φ成正比是固定开销激活值各个非线性模块的中间激活值与batch_size和序列长度相关是动态开销1.4 模型权重与量化精度的关系假设模型参数量为Φ单位参数个数不同量化精度下的显存占用如下量化程度每参数字节数显存占用1B参数模型7B参数模型FP324字节4Φ4GB28GBFP16/BF162字节2Φ2GB14GBINT81字节1Φ1GB7GBINT40.5字节≤1Φ0.5GB3.5GB1.5 模型参数量的计算公式以Llama-3模型为例其参数量由以下符号定义符号含义n_vocab词表中词的个数d_hidden隐藏层维度嵌入向量的维度n_head注意力头的数量n_kv-head分组查询注意力中的键值头数量n_layerTransformer的层数d_FFN前馈神经网络的隐藏层维度b输入数据的批次大小batch sizes输入序列长度模型总参数量公式$$\Phi n_{\text{vocab}} \times d_{\text{hidden}} n_{\text{layer}} \times \left[ d_{\text{hidden}} \left(2 2 \cdot \frac{n_{\text{kv}}}{n_{\text{head}}}\right) d_{\text{hidden}}^2 d_{\text{hidden}} 3 \cdot d_{\text{hidden}} \cdot d_{\text{FFN}} \right] d_{\text{hidden}} d_{\text{hidden}} \times n_{\text{vocab}}$$组成部分公式项说明词嵌入层$n_{\text{vocab}} \times d_{\text{hidden}}$词表大小 × 隐藏维度Transformer层×$n_{\text{layer}}$含 QKV 投影 FFNGQA 时 KV 头数 注意力头数输出层$d_{\text{hidden}} d_{\text{hidden}} \times n_{\text{vocab}}$LayerNorm 输出投影注意当 n_kv-head 1 时为多查询注意力MQA当 n_kv-head n_head 时为多头注意力MHA当 1 n_kv-head n_head 时为分组查询注意力GQA1.6 通俗理解直观类比想象你在搬家需要把所有家当装上卡车GPU显存。模型权重 你的家具沙发、床、桌子——这些是固定的搬多少次都一样重。梯度 每件家具的搬运说明书——和家具数量一一对应同样多。优化器状态 每件家具的维修记录和使用日志——Adam优化器需要记录每个参数的动量和方差所以额外占用2倍的家具重量。激活值 搬运过程中的临时存放点——搬的批次batch_size越多需要的临时空间越大。量化精度就像选择不同精度的包装方式FP32 用厚实的防震泡沫包裹每件家具4字节/参数最安全但最占空间FP16 用薄一些的包装2字节/参数空间减半INT8 只用塑料薄膜简单裹一下1字节/参数核心要点显存 模型状态固定 激活值动态两者都需要关注量化精度每降一档模型权重显存减半7B模型仅权重FP32就需要28GB整体训练显存远超单卡容量1.7 小结维度说明两大组成模型状态权重梯度优化器 激活值量化关系FP324Φ, FP162Φ, INT81Φ参数计算含词嵌入层 n_layer个Transformer层 输出层关键认知7B模型FP32权重28GB训练总显存约112GB2. FP32训练与混合精度训练2.1 核心问题FP32训练需要多少显存混合精度训练能节省显存吗2.2 原文核心要点使用AdamW优化器进行FP32训练模型状态总显存为16Φ。混合精度训练并没有节省模型状态的显存其真正优势是加速计算和降低激活值显存。2.3 FP32训练的显存占用通俗来讲训练模型不仅要存模型本身还要存每个参数的更新历史优化器状态这才是显存的大头。使用AdamW优化器进行FP32训练时