Llama中模块参数大小

LLama2中,流程中数据大小的变换如下

Transformer模块

第一次输入,进行prefill,输入x维度为[1, 8, 4096]

  1. 构建wq,wk,wv,wo,尺寸均为[4096,4096], 与x点乘,得到xq, xk, xv

  2. 构建KV cache, 尺寸为 [batch size, max_seq_len, local_kv_heads, head_dim],对应 [1, 8, 32, 128]

3.基于kv cache构造 keys, alues,对应的尺寸还是[1,8,32,128]

  1. 在最后两个维度对于xq和key进行点乘,得到scores,维度变成【1, 32, 8, 8】

  2. 将mask与scores相加

  3. 对于scores进行softmax

  4. 将scores [1, 32, 8, 8]与values [1, 32, 8, 128]进行乘法

  5. 得到output [1, 8, 4096]

  6. 将output再与wo进行乘法[1, 8, 4096]

  7. 接下来对于输出进行 ffn_norm的操作

Feedforward模块

11.然后进行feed_forward.得到当前transformer模块的输出 [1, 8, 4096]

feed_forward的操作如下,虽然代码很小,但是计算量却很大。

复制代码
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

其中,w1的维度为[11008, 4096], w2的维度为[4096, 11008], w3的维度为[11008, 4096]

kv cache的表达如下

python 复制代码
        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )

关于kv cache的细节讨论

llama2设定 local_kv_heads为32,head_dim为128。所以,kv cache的尺寸为 [1, 512,32, 128] * 2

对于一个batch的数据来说哦,因为llama2 7B 包含32个transformer,所以,当使用FP32表达时, 对应一个batch的kv cache的大小为128 * 32 * 128 *2 * 32 * 4byte= 0.5GB.

这里,也可以看到几个变量:

* 当batch变大时,kv cache线性增长

* 当batch 的最大长度增大时, Kv cache线性增长。

参考链接:

https://arxiv.org/pdf/1911.02150

相关推荐
AndrewHZ2 分钟前
【图像处理基石】如何入门色彩评估?
图像处理·人工智能·深度学习·色彩科学·hvs·色彩评估·颜色工程
TomatoSCI2 分钟前
聚类的可视化选择:PCA / t-SNE丨TomatoSCI分析日记
人工智能·机器学习
大咖分享课4 分钟前
深度剖析:最新发布的ChatGPT Agent 技术架构与应用场景
人工智能·openai·智能助手·ai代理·chatgpt agent·自主任务执行
lucky_lyovo14 分钟前
卷积神经网络--网络性能提升
人工智能·神经网络·cnn
liliangcsdn19 分钟前
smolagents - 如何在mac用agents做简单算术题
人工智能·macos·prompt
nju_spy23 分钟前
周志华《机器学习导论》第8章 集成学习 Ensemble Learning
人工智能·随机森林·机器学习·集成学习·boosting·bagging·南京大学
静心问道1 小时前
TrOCR: 基于Transformer的光学字符识别方法,使用预训练模型
人工智能·深度学习·transformer·多模态
说私域1 小时前
基于开源AI大模型、AI智能名片与S2B2C商城小程序源码的用户价值引导与核心用户沉淀策略研究
人工智能·开源
亲持红叶1 小时前
GLU 变种:ReGLU 、 GEGLU 、 SwiGLU
人工智能·深度学习·神经网络·激活函数
说私域1 小时前
线上协同办公时代:以开源AI大模型等工具培养网感,拥抱职业变革
人工智能·开源