大模型参数估算

大模型参数估算

引用：

大模型训练需要多少GPU？大模型显存怎么算？ #大模型 #AI系统（本文主要基于该视频的内容，ZOMI老师牛逼）

LLaMA: Open and Efficient Foundation Language Models

ML 2021 Spring (ntu.edu.tw)

[1810.04805] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (arxiv.org)

符号约定

参数	含义
l	解码器的层数
h	隐藏层维度/特征向量维度
a	注意力头数
V	词表大小/词向量维度
b	批大小
s	序列长度

大模型架构

以以下架构为准进行估算

嵌入层

这里假设位置嵌入采用不可学习的方式，因此只计算词嵌入矩阵大小为Vh

Transformer层

自注意力层

矩阵形式表达如下：

1. 将三个权重矩阵QKV拆分为a份，分别进行注意力计算。
2. 最后的输出结果被concat起来并通过$W^0$矩阵映射到输出

QKV每个头的每个矩阵大小为$\frac{h^2}{a}$，一共3a个矩阵，参数一共为$3h^2$，注意这个结果与不分头是一样的。concat后的形状为$V\times (\frac{h}{a}\times a)$，$W_0$的形状为$h\times h$，因此所有的参数加起来为$4h^2$

MLP层

包含两个线性层，分别从$h\to 4h$和从$4h\to h$，总参数为$h\times 4h+4h\times h=8h^2$

归一化层

可训练参数有两个，为缩放参数$\gamma$和平移参数$\beta$，对于$V\times h$的输入矩阵，对列进行归一化。归一化层共两个，因此参数量为$2\times h\times 2=4h$

总参数量与模型大小估算

transformer层的总参数量为
$$l\times (4h^2+8h^2+4h)\approx l(12h^2)$$
对于采用FP16格式保存的模型，其大小约为$24l{h}^2$B。

模型名称	h	l	$12l{h}^2$	实际参数量	模型大小(FP16)
LLAMA-6B	4096	32	6442450944	6.7B	12GB
LLAMA-13B	5120	40	12582912000	13.0B	23.4GB
LLAMA-33B	6656	60	31897681920	32.8B	59.4GB
LLAMA-65B	8192	80	64424509440	65.2B	120GB

这里的计算忽略了Transformer层之外的层，但可以看到，也已经非常接近了。

输出层

总体参数

总体参数为：
$$Vh+l\times (4h^2+8h^2+4h)+输出层参数$$
实际上估算可用如下公式：
$$Vh+l(12h^2)$$

计算尝试

from transformers import LlamaConfig
from transformers import LlamaModel

config=LlamaConfig()
Llama=LlamaModel(config) # 注意这里默认参数很大，可以把config中参数调小

print(Llama) # 查看模型架构