站长收录平台网站后台图片编辑器
显存组成要素
在进行大模型全参数微调时，显存消耗主要来自四个关键部分：
模型权重存储：这是基础开销，取决于参数量和数据精度。例如FP16/BF16精度下，每个参数占用2字节。
优化器状态：以Adam优化器为例，需要存储：
这三部分合计通常是模型权重的2-5倍大小(主要取决于优化器、数值精度类型的选取)
梯度（与参数同形状）
一阶动量
二阶动量
中间激活值：即使batch size很小，前向传播和反向传播过程中产生的中间结果也需要显存空间。
框架开销：包括PyTorch等框架运行时的缓存、工作区等额外消耗。
13B参数模型显存估算示例
让我们以llama3 130亿参数（13B）模型为例，详细计算全参数微调时的显存需求：
模型权重部分：
假设采用FP16/BF16精度：
参数数量：13,000,000,000（FP16/BF16精度：每个参数2字节）
总大小 = 13B × 2B = 26GB
优化器状态部分：
假设使用Adam，采用常规混合精度训练(梯度FP16，动量FP32)
如果使用动量混合精度，则一阶动量改用FP16，优化器总显存约16+32=48GB
一阶动量：13*4=52GB（FP32，每个参数4字节）
二阶动量：13*4=52GB（FP32）【二阶动量一般FP32，因为值很小，如果量化，精度会受损】
总计 = 52 + 52 =104GB
现在在进行深度学习训练时，可以用 checkpoint 技术来节省内存。它的做法是：在前向计算时只保存一部分关键的中间结果（激活值），等到反向传播时再重新计算这些值，以此来减少内存使用。
主要有两种策略：
全量 checkpoint（full checkpointing）：对模型中的所有操作都做 checkpoint，这意味着在反向传播时需要重新执行一次完整的前向计算。虽然这样能把显存使用从比如 60GB 降到 8GB，但代价是计算量大了，大概有 36% 的额外开销。
选择性 checkpoint（selective checkpointing）：只对一些计算量小但内存占用大的操作（比如 attention 部分）做 checkpoint。这样可以把重新计算的开销从 36% 降低到 4% 左右，更高效。
量化代码示例
使用HuggingFace Transformers进行4-bit量化的完整配置：
量化代码示例
使用HuggingFace Transformers进行4-bit量化的完整配置：
from transformers import BitsAndBytesConfig
import torch
bnb_config = BitsAndBytesConfig(load_in_4bit=True, # 启用4-bit量化加载bnb_4bit_use_double_quant=True, # 使用双重量化减少误差bnb_4bit_quant_type="nf4", # 使用NormalFloat4量化类型bnb_4bit_compute_dtype=torch.bfloat16 # 计算时使用bfloat16
)
model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-7b",quantization_config=bnb_config,device_map="auto"
)
为什么量化后所需要的训练显存没有明显减少？
主要因为：量化虽然能减少模型权重的存储大小，但计算时仍需要反量化回高精度进行计算。
反量化计算开销几乎不变
训练时，量化参数需实时反量化回FP16/FP32才能参与前向和反向计算，导致计算显存仍依赖高精度格式。
显存峰值：反量化后的权重会短暂占用与原始精度相同的显存（如8B参数在FP16下仍需16GB）。
梯度存储需求
梯度通常以FP16/FP32存储（而非量化格式），以便数值稳定性和优化器更新，因此梯度显存占用与未量化时相同。
优化器状态限制
即使权重被量化，优化器（如Adam）的动量/方差状态仍需FP32（除非使用8-bit优化器），这部分显存仍占大头。
例如：8B参数的Adam优化器状态在FP32下仍需 32GB（动量） + 32GB（方差） = 64GB。
哪些层最应该被量化？
大矩阵乘法参数（如FFN层）
原因：FFN层（前馈网络）的参数规模通常占模型50%以上，且对量化噪声容忍度较高。
示例：8B模型的FFN层若用8-bit量化，可节省~8GB显存（假设50%参数量化）。
低敏感度的Embedding层
原因：词嵌入矩阵（Embedding）通常参数量大但对精度要求较低（尤其是低频词）。
注意：高频词或位置嵌入建议保留FP16，避免信息损失。
注意力层的Value/Projection矩阵
原因：Q/K矩阵对注意力分数计算敏感，建议保留FP16；而V（Value）和输出投影矩阵可量化。
实验支持：LLM训练中，V矩阵量化对最终性能影响通常<1%（相比Q/K矩阵）。
混合精度训练如何影响显存？
激活值显存减少（FP16）
前向传播时，激活值（Activations）以FP16存储，相比FP32 显存减少50%。
示例：若原始激活值占10GB（FP32），混合精度下可降至5GB（FP16）。
梯度存储（FP16/FP32混合）
梯度通常以FP16存储（节省显存），但在权重更新前可能转为FP32以保持数值稳定性。
优化器状态的限制（仍需FP32）
混合精度下，优化器状态（如Adam的动量/方差）仍需FP32，无法节省显存。
这是混合精度显存节省的瓶颈（例如8B模型的优化器状态仍占64GB）。
