Llama中模块参数大小

海滩游侠2024-07-26 8:32

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

Transformer模块

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

构建wq,wk,wv,wo,尺寸均为 $4096,4096$ ，与x点乘，得到xq, xk, xv
构建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$

在最后两个维度对于xq和key进行点乘，得到scores，维度变成【1， 32， 8， 8】
将mask与scores相加
对于scores进行softmax
将scores $1, 32, 8, 8$ 与values $1, 32, 8, 128$ 进行乘法
得到output $1, 8, 4096$
将output再与wo进行乘法 $1, 8, 4096$
接下来对于输出进行 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