Transformer & LLM

课程的第二部分主要讲特定领域的优化技术，第一块是对LLM系列的优化技术

前面一大部分是LLM的基础知识，在很多地方已经记录过了，不再做赘述

1. KV-Cache

KV cache 是这类自回归模型中最重要的一类优化技术之一，通过对于context的KV进行存储，使得在进行Attention计算的时候，只需要计算最新进来的token和之前的context的attention即可

而在LLM追求Long Context的时候，KV cache的所占的空间也会变得非常的大

一旦用户变多（Batch size变大）KV cache 占据的存储空间甚至很容易超过模型本身的大小

这里我们就引入了MHA MQA GQA等方案，GQA是这几个里面最优的解决方案，这个和早期对于Conv block的优化非常相似

这些方案是在Attention Head上做处理，通过减少Attention Head的数量，来减少KV cache的存储，标准的MHA的我们有h个Attention Head，也就有h个QKV，MQA做的比较极端，我们只有一个KV权重，但是有h个Q权重，因为Q的计算是逃不掉的，但是这种做法过于极端导致性能较差，GQA就相对做得更好，我们让几个Q权重共享一个KV权重，这样子可以做到一个比较好的权衡

class GroupedQueryAttention(nn.Module):
    """GQA：heads 分组，每组共享 1 组 KV"""
    def __init__(self, d_model, num_heads, num_kv_heads=None):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads or num_heads // 8  # 默认 8 组
        self.d_k = d_model // num_heads
        self.num_queries_per_kv = num_heads // self.num_kv_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, self.num_kv_heads * self.d_k)  # G 组 key
        self.W_v = nn.Linear(d_model, self.num_kv_heads * self.d_k)  # G 组 value
        self.W_o = nn.Linear(d_model, d_model)