AI-Infra

整理自 CMU 15-779 Lect 12，从数据并行到 ZeRO，解答了是如何把大模型塞进有限的显存里的问题。 Data Parallelism 训练本质上是计算梯度并更新权重。数据并行直观理解就是如果有 $ N $ 个 GPU，那么把总的 Batch Size 分成 $ N $ 份。 在第 $ k $ 个 GPU上的计算过程： Local Gradient：基于本地数据 $ x^{(k)} $ 计算梯度 $$ g^{(k)} = \nabla L(w_t, x^{(k)}) $$ All-Reduce：计算全局平均梯度 $$ \bar{g} = \frac{1}{N} \sum_{k=1}^{N} g^{(k)} $$ Weight Update：所有 GPU 使用相同的 $ \bar{g} $ 更新权重 $$ w_{t+1} = w_t - \eta \cdot \text{Optimizer}(\bar{g}, \text{State}_t) $$ 这时每个 GPU 上都会有一样的完整模型副本，各算各的梯度，最后在汇总计算。 实现方法 ...

推理效率对于 llm 是一个至关重要的问题。当模型生成文本时，尤其是以自回归方式逐词生成时，效率瓶颈会变得非常明显。KV Cache 就是为了解决这一问题而诞生的技术。 1. What is KV Cache? KV Cache，全称 Key-Value Cache，是一种优化技术，用于加速 Transformer 架构在自回归生成过程中的推理速度。它的核心思想是缓存并重用在注意力机制中计算得到的 Key (K) 和 Value (V) 向量。 2. Transformer Attention Mechanism Review 要理解 KV Cache，首先需要对 Transformer 架构中的自注意力机制有一个基本认识。自注意力机制允许模型在处理序列中的某个词时，考虑序列中所有其他词的重要性。 每个输入 token（词或子词）在进入注意力层时，都会被转换成三个不同的向量： Q 向量：代表当前 token 的“查询”信息 K 向量：代表所有 token 的“键”信息，用于与 Query 进行匹配 V 向量：代表所有 token 的“值”信息，用于加权求和，得到最终的输出 自注意力机制的计算过程为以下步骤： 计算 Query 与所有 Key 的点积，得到注意力分数 将注意力分数进行缩放，除以 $ \sqrt{d_k} $（$ d_k $ 是 Key 向量的维度) 对缩放后的分数进行 Softmax，将其转换为注意力权重，表示每个 token 对当前 token 的重要性 将注意力权重与 Value 向量进行加权求和，得到当前 token 的注意力输出 公式为： $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ 其中矩阵 $ Q,K,V \in \mathbb{R}^{L \times d} $ ，$ L $ 为当前上下文长度 ...