[bold{x}=[x_1,x_2,x_3] \ bold{W}_q= begin{bmatrix} w_{11} & w_{12} & w_{13}\ w_{21} & w_{22} & w_{23}\ w_{31} & w_{32} & w_{33} end{bmatrix} \ bold{Q}=[q_1,q_2,q_3] ]

那么具体的：

[q_1=w_{11}x_1+w_{21}x_2+w_{31}x_3 \ q_2=w_{12}x_1+w_{22}x_2+w_{32}x_3 \ q_3=w_{13}x_1+w_{23}x_2+w_{33}x_3 ]

设损失函数为(L)，这是一个实值函数，可以将(L)理解为一个标量。我们知道，梯度的定义是损失函数对某个权重的偏导，而梯度可以理解为：某个权重改变了一个单位长度后，损失函数变化的程度。也就是说，我们要求出损失函数对所有可更新参数的偏导，这样才能进行参数更新（梯度下降）。

而在这个过程中，(W_q)是要更新的权重矩阵，(x)是下层输入（随样本数据的变化而变化）。对(W_q)中一个参数的具体的求导过程如下：

[frac{partial{L}}{partial{w_{12}}}= frac{partial{L}}{partial{q_2}} · frac{partial q_1}{partial{w_{12}}} = frac{partial{L}}{partial{q_2}} ·x_1 ]

可以更抽象的解释一下上面的结果：(w_{12})表示第1个位置的输入(x_1)对第2个位置的输出(q_2)的贡献权重。因此先计算(q_2)对(L)的影响，再计算(w_{12})对(q_2)的影响（根据公式的形式是后计算(w_{12})对(q_2)的影响，实际上在前向过程中先计算），根据链式求导法则，二者相乘得到(w_{12})对(L)的影响。

相似的，对(W_q)中各权重的求导结果如下：

[begin{matrix} frac{partial{L}}{partial{w_{11}}}= frac{partial{L}}{partial{q_1}} ·x_1 & frac{partial{L}}{partial{w_{12}}}= frac{partial{L}}{partial{q_2}} ·x_1 & frac{partial{L}}{partial{w_{13}}}= frac{partial{L}}{partial{q_3}} ·x_1 \ frac{partial{L}}{partial{w_{21}}}= frac{partial{L}}{partial{q_1}} ·x_2 & frac{partial{L}}{partial{w_{22}}}= frac{partial{L}}{partial{q_2}} ·x_2 & frac{partial{L}}{partial{w_{23}}}= frac{partial{L}}{partial{q_3}} ·x_2 \ frac{partial{L}}{partial{w_{31}}}= frac{partial{L}}{partial{q_1}} ·x_3 & frac{partial{L}}{partial{w_{32}}}= frac{partial{L}}{partial{q_2}} ·x_3 & frac{partial{L}}{partial{w_{33}}}= frac{partial{L}}{partial{q_3}} ·x_3 \ end{matrix} ]

为了便于书写，现在引入一种新的形式——对矩阵求导：

[frac{partial{L}}{partial bold{W_q}}= begin{bmatrix} frac{partial{L}}{partial w_{11}} & frac{partial{L}}{partial w_{12}} & frac{partial{L}}{partial w_{13}} \ frac{partial{L}}{partial w_{21}} & frac{partial{L}}{partial w_{22}} & frac{partial{L}}{partial w_{23}} \ frac{partial{L}}{partial w_{31}} & frac{partial{L}}{partial w_{32}} & frac{partial{L}}{partial w_{33}} \ end{bmatrix} \ frac{partial{L}}{partial bold{Q}}= begin{bmatrix} frac{partial{L}}{partial q_1} & frac{partial{L}}{partial q_2} & frac{partial{L}}{partial q_3} \ end{bmatrix} ]

就是按元素位置对应求导，向量也是一样（数学形式上，向量就是行为1的二维矩阵）。

那么对(W_q)中各权重的求导结果就可简单的表示为：

[frac{partial{L}}{partial bold{W_q}}= begin{bmatrix} frac{partial{L}}{partial w_{11}} & frac{partial{L}}{partial w_{12}} & frac{partial{L}}{partial w_{13}} \ frac{partial{L}}{partial w_{21}} & frac{partial{L}}{partial w_{22}} & frac{partial{L}}{partial w_{23}} \ frac{partial{L}}{partial w_{31}} & frac{partial{L}}{partial w_{32}} & frac{partial{L}}{partial w_{33}} \ end{bmatrix} \ frac{partial{L}}{partial bold{Q}}= begin{bmatrix} frac{partial{L}}{partial q_1} & frac{partial{L}}{partial q_2} & frac{partial{L}}{partial q_3} \ end{bmatrix}frac{partial{L}}{partial bold{W_q}}= begin{bmatrix} x_1 \ x_2 \ x_{3} end{bmatrix} · begin{bmatrix} frac{partial{L}}{partial q_1} & frac{partial{L}}{partial q_2} & frac{partial{L}}{partial q_3} \ end{bmatrix} = bold{x}^T·frac{partial{L}}{partial bold{Q}} ]

注意，这里的(x)是一个一维向量，形状([3])，在attention中，每个序列的输入(x)的形状是([s,d])，这里假设为([2,3])，提升了一个维度上式同样成立。简单说一下就是(w_{12})表示：(x_{11})与(q_{12})、(x_{21})与(q_{22})之间的权重，于是:

[frac{partial{L}}{partial{w_{12}}}=frac{partial{L}}{partial{q_{12}}} ·frac{partial q_{12}}{partial{w_{12}}} + frac{partial{L}}{partial{q_{22}}} ·frac{partial q_{22}}{partial{w_{12}}} \ =frac{partial{L}}{partial{q_{12}}} ·x_{11} + frac{partial{L}}{partial{q_{22}}} ·x_{21} \ = begin{bmatrix} x_{11} & x_{21} end{bmatrix} · begin{bmatrix} frac{partial{L}}{partial q_{12}}\ frac{partial{L}}{partial q_{22}} end{bmatrix} ]

总之，根据计算结果，当我们反向传播更新权重(W_q)时，需要两个参数(x^T)和(frac{partial{L}}{partial bold{Q}})，其中(frac{partial{L}}{partial bold{Q}})只能反向传播过程才能得到。而(x^T)在前向过程中，也(Q=xW_q)就是过程中，就可以计算得到了，于是(x^T)（程序中直接保存(x)和）就是(xW_q)和的中间激活值。

中间激活值显存计算

中间激活值也采用16位浮点数，占2bytes

首先应该是Embedding层的中间激活值，但是文章中说不需要，考虑到Embedding层和输出层参数贡献，我猜测是两种可能之一：

• 仅在Embedding层更新参数，输出层参数固定。假设(bold{x}=bold{Seqs}bold{W}_E)，中间激活就是(Seqs)，而(Seqs)可能已经保存在显存中了，不作为中间激活额外保存。
• 仅在输出层更新参数，Embedding层参数不更新。假设(logits=xW_E)，那么中间激活就是(x)。

这里假设是第二种。

然后考虑Multi-mask Self-attention：

• 对于(xbold{W}_q,xbold{W}_k,xbold{W}_v)，第一层block中输入attention层的(x_0)可能没有参与过可训练参数的计算，所以不用计算(frac{partial{L}}{partial bold{x_0}})，但是后续block中既要算(frac{partial{L}}{partial bold{x_i}})也要算(frac{partial{L}}{partial bold{W}_q^i})，需要保存(W_q)和(x)，但是(W_q)本身就是模型参数，不需要额外保存，因此不是中间激活。所以中间激活只有(x)，形状为([b,s,d])，占用显存大小(2bsd)bytes。

• 对于(c)，需要计算(frac{partial{L}}{partial bold{Q}})和(frac{partial{L}}{partial bold{K^T}})，各自需要保存Q和(K^T)，(Q,K)的形状都是([b,h,s/h,d])，共计占用显存大小(4bsd)bytes。

• 对于(Softmax(frac{QK^T}{sqrt{d}}))，设(S=Softmax(bold{t}))，其中(bold{t}=[t_1,...,t_n],S=[s_1,...,s_n])。则：

  [frac{partial s_i}{partial t_j}=frac{partial}{partial t_j}(frac{e^{t_i}}{sum_k{e^{t_k}}})= begin{cases} -s_i s_j& text{i != j} \ s_i(1-s_i)& text{i == j} end{cases} \ frac{partial S}{partial t}=[frac{partial s_i}{partial t_j}]_{i=0,j=0}^{nn}=diag(S)-S^TS ]

  按道理，需要保存的是(S=Softmax(bold{t}))的结果，但是我看文章中写的是保存(QK^T)，不管是哪个，形状都是([b,h,s/h,d,d])，占用显存大小(2bsd)bytes。

• 对于(S(score)·V)，保存(S(score))和(V)，形状分别是([b,h,s/h,d,d])和([b,h,s/h,d])，共占用显存(2bsd^2+2bsd)bytes。

• 对于(V_{out}·W_o,V_{out}=S(score)·V)，保存(V_{out})和(W_o)，但是(W_o)是模型参数不用额外保存，(V_{out})形状为([b,h,s/h,d])，共占用显存(2bsd)bytes。

• dropout，不太清楚，元素用1byte存储，占用显存(bsd)bytes。

• Self-attention层总计显存占用(11bsd+5bsd^2)。

Layer Normalization：

不会算，根据资料，需要保存输入(x)，以及方差(sigma)和均值(mu)，共计(2bsd+2bs)bytes。一共有两层LN，并且省略方差和均值的显存占用，共计(4bsd)bytes。

MLP层：

• 线性层(drightarrow 4d)，保存输入，占用显存(2bsd)bytes。
• 激活层，不会算，保存输入，占用显存(8bsd)bytes。
• 线性层，保存输入，占用显存(8bsd)bytes。
• dropout，保留mask矩阵，占用显存(bsd)bytes。
• 总计(19bsd)bytes。

中间激活值占用显存总计((34bsd+5bsd^2))bytes。

最终(l)层block中间激活层共计(l*(34bsd+5bsd^2))bytes

于是总的显存占用为(16l*(12d^2+13d)+Vd+l*(34bsd+5bsd^2) + bsd)bytes.

计算量

一次矩阵运算，例如(QK^T)，一共有(b*s^2)个元素，每个元素的计算都进行了(d)次的加法和(d)次的乘法，浮点数的一次加法或者乘法运算就被称为一次浮点数运算，总共做了(2bs^2d)次浮点数运算。

阶段	运算	浮点数运算
Embedding	(x=SeqsW_E)	因为one-hot非常稀疏，浮点运算次数未知
Self-attention	(xbold{W}_q,xbold{W}_k,xbold{W}_v)	(3*bsd*2d=6bsd^2)
Self-attention	(QK^T)	(bs^2*2d=2bs^2d)
Self-attention	(Softmax(frac{QK^T}{sqrt{d}}))	(bs*4s=4bs^2)
Self-attention	(S(score)·V)	(bsd*2s=2bs^2d)
Self-attention	(V_{out}·W_o)	(bsd*2d=2bsd^2)
Layer Normalization	(a=frac{x_{in}-mu}{sqrt{(sigma)^2+epsilon}})	(bs*3d=3bsd?)
Layer Normalization	(bold{gamma}odot bold{a} + bold{beta})	(bs*2d=2bsd)
MLP	(xW_1)	(4bsd*2d=8bsd^2)
MLP	(GeLu(xW_1))	未知
MLP	(xW_2)	(bsd*8d=8bsd^2)
输出层	(logits=xW_E^T)	(bsV*2d=2bsdV)
总计	忽略复杂度较低的	(l*(24bsd^2+4bs^2d)+2bsdV)

训练时间

根据浮点计算次数以及显卡计算速度和利用率计算训练时间。

显卡利用率一般在0.35到0.5之间。

KV Cache

kv cache是推理时采用的技术，是一种空间换时间的方案。

没有kv cache的推理过程中有大量的重复计算，例如重复计算(xbold{W}_q,xbold{W}_k,xbold{W}_v)。

因为推理是自回归的，很自然的会把代码写成下面的形式：

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer


model = GPT2LMHeadModel.from_pretrained("/WORK/Test/gpt", torchscript=True).eval()

# tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("/WORK/Test/gpt")
in_text = "Lionel Messi is a" # 很多文章也叫做prompt
in_tokens = torch.tensor(tokenizer.encode(in_text))

# inference
token_eos = torch.tensor([198]) # 句段结束标志。
out_token = None
i = 0
with torch.no_grad():
    while out_token != token_eos:
        logits, _ = model(in_tokens)
        out_token = torch.argmax(logits[-1, :], dim=0, keepdim=True) # 取序列末尾的token对应的输出用来预测下一个词
        in_tokens = torch.cat((in_tokens, out_token), 0)
        text = tokenizer.decode(in_tokens) # 将tokens变成句子
        print(f'step {i} input: {text}', flush=True) # 输出句子
        i += 1

out_text = tokenizer.decode(in_tokens)
print(f' Input: {in_text}')
print(f'Output: {out_text}')

对于代码中的in_text，也就是prompt来说，每一次循环，都要计算(xbold{W}_q,xbold{W}_k,xbold{W}_v)，利用矩阵乘法的分块乘性质，将这些结果保存，只需要计算新的token的(x_ibold{W}_q,x_ibold{W}_k,x_ibold{W}_v)，就可以大大减少计算量。

参考资料：

分析transformer模型的参数量、计算量、中间激活、KV cache - 知乎 (zhihu.com)

[LLM]KV cache详解 图示，显存，计算量分析，代码 - 知乎 (zhihu.com)

反向传播算法推导过程（非常详细） - 知乎 (zhihu.com)

大模型推理性能优化之KV Cache解读 - 知乎 (zhihu.com)https://zhuanlan.zhihu.com/p/630832593)