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

相关推荐
拆房老料13 小时前
Transformer推理优化全景:从模型架构到硬件底层的深度解析
深度学习·ai·自然语言处理·transformer
CiLerLinux13 小时前
第四十九章 ESP32S3 WiFi 路由实验
网络·人工智能·单片机·嵌入式硬件
七芒星202315 小时前
多目标识别YOLO :YOLOV3 原理
图像处理·人工智能·yolo·计算机视觉·目标跟踪·分类·聚类
Learn Beyond Limits16 小时前
Mean Normalization|均值归一化
人工智能·神经网络·算法·机器学习·均值算法·ai·吴恩达
byzy16 小时前
【论文笔记】VisionPAD: A Vision-Centric Pre-training Paradigm for Autonomous Driving
论文阅读·深度学习·计算机视觉·自动驾驶
ACERT33316 小时前
5.吴恩达机器学习—神经网络的基本使用
人工智能·python·神经网络·机器学习
C嘎嘎嵌入式开发16 小时前
(一) 机器学习之深度神经网络
人工智能·神经网络·dnn
Aaplloo16 小时前
【无标题】
人工智能·算法·机器学习
大模型任我行16 小时前
复旦:LLM隐式推理SIM-CoT
人工智能·语言模型·自然语言处理·论文笔记
tomlone16 小时前
AI大模型核心概念
人工智能