对于llama3.1 8B模型，FP32和BF16混合精度训练，用的是AdamW优化器，模型训练时占用显存分析

samoyan

已于 2024-08-22 11:30:51 修改

阅读量770

点赞数 6

分类专栏： LLM 面试文章标签：人工智能机器学习深度学习

于 2024-08-22 11:28:36 首次发布

本文链接：https://blog.csdn.net/baoyan2015/article/details/141423192

版权

LLM 面试专栏收录该内容

35 篇文章 35 订阅

订阅专栏

显存主要被用在四个模块上：

模型权重本身
梯度
优化器
激活值

其中，

静态：优化器状态、模型参数
动态：激活值、梯度值

也就是说，我们其实没法特别准确的计算出我们实际运行时候的显存大小，如果在面试的时候，就可以忽略掉激活值的计算，梯度当做静态计算就好。

为什么先不考虑激活值的显存占用

在计算显存占用时，我们通常会区分模型参数、梯度参数和优化器状态的显存占用，以及激活值的显存占用。以下是具体原因：

模型参数、梯度参数和优化器状态：
- 这些部分的显存占用是相对固定的，取决于模型的大小和优化器的选择。
- 在混合精度训练中，我们可以明确计算这些部分的显存占用。
激活值：
- 激活值的显存占用与批量大小（batch size）和序列长度（sequence length）密切相关，且在不同的训练任务和配置下变化较大。
- 激活值的显存占用往往是动态的，取决于具体的训练过程和数据流动。

因此，在讨论显存占用时，我们通常会先计算固定部分（模型参数、梯度参数和优化器状态）的显存占用，而不考虑激活值的显存占用。这是因为激活值的显存占用是高度可变的，需要根据具体的训练配置进行动态调整。具体变化在最后简单介绍一下。

在模型训练中，显存占用主要包括模型参数、梯度参数和优化器状态。对于LLaMA 3.1 8B模型，使用混合精度训练（FP32和BF16）和AdamW优化器时，显存占用的计算如下：

1. 模型参数

含义

模型参数是神经网络的权重和偏置等参数。对于8B参数的模型：

BF16（Brain Floating Point 16-bit）：每个参数占用16位（2字节）
FP32（Floating Point 32-bit）：每个参数占用32位（4字节）

计算

假设模型的所有参数都存储为BF16和FP32两种格式：

BF16：8B参数 * 2字节 = 16GB
FP32：8B参数 * 4字节 = 32GB

总的模型参数显存占用为： 16𝐺𝐵+32𝐺𝐵=48𝐺𝐵

2. 梯度参数

含义

梯度参数是用于反向传播更新模型参数的梯度值。在混合精度训练中，梯度通常以BF16格式存储：

BF16：8B参数 * 2字节 = 16GB

总的梯度参数显存占用为： 16𝐺𝐵

3. 优化器参数

含义

AdamW优化器需要存储额外的状态参数，包括一阶动量（momentum）和二阶动量（variance）。这些参数通常以FP32格式存储：

一阶动量（FP32）：8B参数 * 4字节 = 32GB
二阶动量（FP32）：8B参数 * 4字节 = 32GB

总的优化器参数显存占用为： 32𝐺𝐵+32𝐺𝐵=64𝐺𝐵

4. 较固定总显存占用

计算

不考虑激活值的情况下，总显存占用为： 48𝐺𝐵（模型参数）+16𝐺𝐵（梯度参数）+64𝐺𝐵（优化器参数）=128𝐺𝐵

详细解释

模型参数（48GB）：
- BF16：模型的所有参数以16位格式存储，占用16GB显存。
- FP32：模型的所有参数以32位格式存储，占用32GB显存。
梯度参数（16GB）：
- BF16：用于反向传播的梯度参数以16位格式存储，占用16GB显存。
优化器参数（64GB）：
- 一阶动量（32GB）：AdamW优化器的一阶动量参数以32位格式存储，占用32GB显存。
- 二阶动量（32GB）：AdamW优化器的二阶动量参数以32位格式存储，占用32GB显存。