大模型参数估算
引用:
大模型训练需要多少GPU?大模型显存怎么算? #大模型 #AI系统(本文主要基于该视频的内容,ZOMI老师牛逼)
LLaMA: Open and Efficient Foundation Language Models
符号约定
参数 | 含义 |
---|---|
l | 解码器的层数 |
h | 隐藏层维度/特征向量维度 |
a | 注意力头数 |
V | 词表大小/词向量维度 |
b | 批大小 |
s | 序列长度 |
大模型架构
以以下架构为准进行估算
嵌入层
这里假设位置嵌入采用不可学习的方式,因此只计算词嵌入矩阵大小为Vh
Transformer层
自注意力层
矩阵形式表达如下:
- 将三个权重矩阵QKV拆分为a份,分别进行注意力计算。
- 最后的输出结果被concat起来并通过 W 0 W^0 W0矩阵映射到输出
QKV每个头的每个矩阵大小为 h 2 a \frac{h^2}{a} ah2,一共3a个矩阵,参数一共为 3 h 2 3h^2 3h2,注意这个结果与不分头是一样的。concat后的形状为 V × ( h a × a ) V\times (\frac{h}{a} \times a) V×(ah×a), W 0 W_0 W0的形状为 h × h h\times h h×h,因此所有的参数加起来为 4 h 2 4h^2 4h2
MLP层
包含两个线性层,分别从 h → 4 h h \rarr 4h h→4h和从 4 h → h 4h \rarr h 4h→h,总参数为 h × 4 h + 4 h × h = 8 h 2 h\times 4h+4h\times h=8h^2 h×4h+4h×h=8h2
归一化层
可训练参数有两个,为缩放参数 γ \gamma γ和平移参数 β \beta β,对于 V × h V\times h V×h的输入矩阵,对列进行归一化。归一化层共两个,因此参数量为 2 × h × 2 = 4 h 2\times h\times 2=4h 2×h×2=4h
总参数量与模型大小估算
transformer层的总参数量为
l
×
(
4
h
2
+
8
h
2
+
4
h
)
≈
l
(
12
h
2
)
l\times(4h^2+8h^2+4h)\approx l(12h^2)
l×(4h2+8h2+4h)≈l(12h2)
对于采用FP16格式保存的模型,其大小约为
24
l
h
2
24lh^2
24lh2B。
模型名称 | h | l | 12 l h 2 12lh^2 12lh2 | 实际参数量 | 模型大小(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层之外的层,但可以看到,也已经非常接近了。
输出层
总体参数
总体参数为:
V
h
+
l
×
(
4
h
2
+
8
h
2
+
4
h
)
+
输出层参数
Vh+l\times(4h^2+8h^2+4h)+输出层参数
Vh+l×(4h2+8h2+4h)+输出层参数
实际上估算可用如下公式:
V
h
+
l
(
12
h
2
)
Vh+l(12h^2)
Vh+l(12h2)
计算尝试
from transformers import LlamaConfig
from transformers import LlamaModel
config=LlamaConfig()
Llama=LlamaModel(config) # 注意这里默认参数很大,可以把config中参数调小
print(Llama) # 查看模型架构