深入了解大型语言模型参数和内存需求

计算参数数量

用于高效 Transformer 训练的预训练模型

计算基于 Transformer 的 LLM 中的参数

让我们分解一下基于 Transformer 的 LLM 的每个组件的参数计算。我们将使用原始论文中的符号，其中 d_model 表示模型隐藏状态的维度。

1. 嵌入层:
   • 参数 = vocab_size * d_model
2. 多头注意力:
   • 针对 h 头部， d_k = d_v = d_model / h:
   • 参数 = 4 * d_model^2（用于 Q、K、V 和输出预测）
3. 前馈网络:
   • 参数 = 2 * d_model * d_ff + d_model + d_ff
   • 其中 d_ff 通常是 4 * d_model
4. 层标准化:
   • 参数 = 2 * d_model （针对规模和偏差）

一个Transformer层的总参数：

• Parameters_layer = Parameters_attention + Parameters_ffn + 2 * Parameters_layernorm

对于具有 N 层数：

• 总参数 = N * Parameters_layer + Parameters_embedding + Parameters_output

示例计算

让我们考虑具有以下规格的模型：

• d_model = 768
• h （注意力头数量）= 12
• N （层数）= 12
• vocab_size = 50,000

1. 嵌入层:
   • 50,000 * 768 = 38,400,000
2. 多头注意力:
   • 4 * 768^2 = 2,359,296
3. 前馈网络:
   • 2 * 768 * (4 * 768) + 768 + (4 * 768) = 4,719,616
4. 层标准化:
   • 2 * 768 = 1,536

每层总参数：

• 2,359,296 + 4,719,616 + (2 * 1,536) = 7,081,984

12层的总参数：

• 12 * 7,081,984 = 84,983,808

模型参数总计：

• 84,983,808 + 38,400,000 123,383,808

该模型大约有 123 亿个参数。

内存使用类型

使用 LLM 时，我们需要考虑两种主要类型的内存使用情况：

1. 模型记忆：存储模型参数所需的内存。
2. 工作记忆：推理或训练期间所需的内存，用于存储中间激活、梯度和优化器状态。

计算模型内存

模型内存与参数数量直接相关。每个参数通常存储为 32 位浮点数，尽管有些模型使用 16 位浮点数进行混合精度训练。

模型内存（字节）= 参数数量 * 每个参数的字节数

对于具有 123 亿个参数的示例模型：

• 型号内存（32 位）= 123,383,808 * 4 字节 = 493,535,232 字节 ≈ 494 MB
• 型号内存（16 位）= 123,383,808 * 2 字节 = 246,767,616 字节 ≈ 247 MB

评估工作记忆

根据具体任务、批次大小和序列长度，工作记忆要求可能会有很大差异。推理过程中的工作记忆粗略估计为：

工作记忆 ≈ 2 * 模型记忆

这用于存储模型参数和中间激活。在训练期间，由于需要存储梯度和优化器状态，​​内存要求可能更高：

训练内存 ≈ 4 * 模型内存

对于我们的示例模型：

• 推理工作内存 ≈ 2 * 494 MB = 988 MB ≈ 1 GB
• 训练内存 ≈ 4 * 494 MB = 1,976 MB ≈ 2 GB

稳定状态内存使用情况和峰值内存使用情况

在基于 Transformer 架构训练大型语言模型时，了解内存使用情况对于高效分配资源至关重要。让我们将内存需求分为两大类：稳定状态内存使用情况和峰值内存使用情况。

稳定状态内存使用情况

稳定状态内存使用情况包括以下部分：

1. 型号重量：FP32 模型参数的副本，需要 4N 字节，其中 N 是参数的数量。
2. 优化器状态：对于 Adam 优化器，这需要 8N 字节（每个参数 2 个状态）。
3. 渐变：渐变的FP32副本，需要4N字节。
4. 输入数据：假设输入为 int64，这需要 8BD 字节，其中 B 是批量大小，D 是输入维度。

总的稳定状态内存使用量可以近似为：

• M_steady = 16N + 8BD 字节

峰值内存使用率

在反向传递期间，当激活被存储用于梯度计算时，内存使用量达到峰值。造成内存峰值的主要因素有：

1. 层标准化：每层标准需要 4E 字节，其中 E = BSH（B：批量大小，S：序列长度，H：隐藏大小）。
2. 注意力障碍:
   • QKV计算：2E字节
   • 注意力矩阵：4BSS 字节（S：序列长度）
   • 注意输出：2E字节
3. 前馈模块:
   • 第一线性层：2E 字节
   • GELU 激活：8E 字节
   • 第二线性层：2E 字节
4. 交叉熵损失:
   • Logits：6BSV 字节（V：词汇量）

总激活记忆可以估算为：

• M_act = L * (14E + 4BSS) + 6BSV 字节

其中 L 是 Transformer 层的数量。

总峰值内存使用量

训练期间的峰值内存使用量可以通过结合稳态内存和激活内存来近似：

• M_peak = M_steady + M_act + 4BSV 字节

额外的 4BSV 项解释了在后向传递开始时的额外分配。

通过了解这些组件，我们可以优化训练和推理期间的内存使用情况，确保高效的资源分配并提高大型语言模型的性能。