大模型训练参数量计算

在现有的大模型训练方案中，通常会采用混合精度训练，模型参数和模型梯度通常以 16 位浮点数存储，而 Adam 或 AdamW优化器则需要额外存储 32 位浮点数的模型参数、动量参数以及动量二阶矩参数。大语言模型的参数量取决于多个关键因素，包括词表大小（𝑉）、解码器层数（𝐿）、中间状态维度（𝐻）、以及前馈网络层的中间维度（𝐻′）。假设模型的参数量为 𝑃，训练中配备有 𝐺 张 GPU，训练

钦臣

1239人浏览 · 2024-12-23 10:08:24

钦臣 · 2024-12-23 10:08:24 发布

大模型训练参数量计算

1.参数量计算

大语言模型的参数量取决于多个关键因素，包括词表大小（𝑉）、解码器层数（𝐿）、中间状态维度（𝐻）、以及前馈网络层的中间维度（𝐻′）。以LLaMA模型为例，可以按以下步骤计算参数量：

输入嵌入层：𝑉𝐻参数。
多头注意力层：每层4𝐻²参数。
前馈网络层：每层3𝐻𝐻′参数。
归一化层：每层2𝐻参数，外加最后一层的𝐻参数。
输出层：𝑉𝐻参数。
总参数量公式为：参数量 = 2𝑉𝐻 + 𝐻 + 𝐿 · (4𝐻^2 + 3𝐻𝐻′ + 2𝐻)

以LLaMA (7B)为例，设𝑉=32000, 𝐿=32, 𝐻=4096, 𝐻′=11008，代入公式得参数量为6,738,415,616。

2.训练运算量估计

模型训练运算量指的是模型在训练过程中，需要进行的浮点运算次数（FloatingPoint Operations, FLOP）。这里的浮点运算包括浮点数的加减乘除运算，以及浮点数的指数函数，对数函数，三角函数等运算操作。
进一步设定以下参数：模型总参数量为 𝑃，批处理大小为 𝐵，输入序列长度为 𝑇，因此训练词元总数为 𝐶 = 𝐵𝑇；多头注意力机制包含 𝑁 个头，每个头的维度为 𝐷，因此和中间状态维度 𝐻 满足关系 𝐻 = 𝑁𝐷。
模型训练的运算量（FLOP）主要由多头注意力和线性变换决定。对于多头注意力，运算量与批处理大小（𝐵）、序列长度（𝑇）、注意力头数（𝑁）、及头的维度（𝐷）相关。线性变换的运算量则直接关联于矩阵参数量。

总运算量公式近似为:运算量≈6CP，其中𝐶是词元总数，𝑃是模型参数量。

3.训练时间估计

训练时间受运算量和GPU性能影响。

训练时间= 运算量/（GPU数量*GPU每秒浮点运算数）估算。
例如，LLaMA (65B)使用2048张A100 GPU，在1.4T词元上的训练时间约20.6天。

4.大模型训练显存估计

在训练大模型时，显存管理是一个关键的考量因素。显存占用主要来源于三个方面：模型参数与优化器、训练中保存的激活值、以及其他显存占用。

在现有的大模型训练方案中，通常会采用混合精度训练，模型参数和模型梯度通常以 16 位浮点数存储，而 Adam 或 AdamW优化器则需要额外存储 32 位浮点数的模型参数、动量参数以及动量二阶矩参数。

假设模型的参数量为 𝑃，训练中配备有 𝐺 张 GPU，训练的数据并行数为 𝑁𝐷，流水线并行数为 𝑁𝑃，张量并行数为 𝑁𝑇。

模型参数与优化器显存占用情况分析如下：

不使用ZeRO优化: 每张GPU需要16P字节显存。一个 16 位浮点数需要 2 字节，一个32 位浮点数需要 4
字节，因此模型参数和模型梯度各需要 2𝑃 字节的显存，Adam 优化器的模型参数、动量参数以及动量二阶矩参数则各需要 4𝑃字节的显存。通过对于这些显存占用进行累和，每张 GPU 上会需要使用 (2+2+4+4+4) ·𝑃 =16𝑃字节的显存用于存储模型参数与优化器。
使用ZeRO-1: 会将优化器的参数进行平摊到每张 GPU 上，而模型参数和模型梯度需要每张显卡各自保留。每张GPU需要4P +12P/ND字节显存。
使用ZeRO-2: 每张GPU需要2P + 14P/ND字节显存。
使用ZeRO-3: 每张GPU需要16P/ND字节显存。
张量并行和流水线并行: 显存开销需额外除以NPxNT，NP为流水线并行数，NT为张量并行数。

训练激活值的显存占用：