Transformer 大模型架构深度解析（6）Attention 和 FFN 模型结构的核心优化方向概览

目录文章目录目录Attention 优化方向Softmax AttentionSparse AttentionFlashAttentionPageAttentionRadixAttentionFFN 优化方向 —— MoEMoE 的诞生背景和思想Dense 与 MoE 的对比MoE 训练特点MoE 推理特点MoE FFN 的结构专家网络Experts Network门控网络Gating Network负载均衡损失函数Load Balancing Loss专家并行Expert ParallelismMoE 整体计算流程Attention 优化方向模型算法层面针对 Attention 的优化主要有 2 个方向**Softmax Attention **从 KV 矩阵的方向出发共享 KV 矩阵或者压缩 KV 矩阵的维度。Sparse Attention利用注意力得分矩阵的稀疏性选择部分 token 进行注意力运算从而减少运算次数工程层面根据 GPU 的存储架构做软件适配比如 FlashAttention分块计算、在线 Softmax、Paged Attention虚拟显存管理、RadixAttention公共前缀。根据模型结构做的框架层的优化比如 KV Cache、PD 分离。Softmax AttentionSoftmax Attention 即使用了 softmax 求权重的 Attention目前的主流相对的还有 Linear Attention。Softmax Attention 目前主要有以下优化方式MHAMulti-Head Attention每个 Attention Head 都有独立的 Q、K、V。KV Cache 随 head 数线性增长推理时显存和带宽压力较大在当代大模型中常被 GQA/MLA 等替代但并非不再使用。注KV Cache 大小取决于层数 × head 数 × head_dim × seq_len × batch × dtype。MQAMulti-Query Attention所有注意力头共享同一组 K/V 投影Q 仍然各头独立。Q、O 投影参数量不变减少的主要是 K/V 投影参数。计算量和 KV Cache 也最小的但是模型效果差。GQAGrouped-Query Attention分组查询注意力对 Attention Head 进行分组一组 Attention Head 中的 Q 共享相同的 K、V。是 MHA 和 MQA 的折中有效减少了参数量例如 70B 有 64 个 Q 头但只有 8 个 K/V 头可以在保持性能的同时减少 30%-50% 注意力层参数。GQA 通常是从 MHA checkpoint 上 uptrain主要收益是推理效率和 KV Cache 降低。目前的主流之一。MLAMulti-head Latent Attention2024 年 DeepSeek V2 提出并在 DeepSeek V3 中沿用。核心思想是每个 Attention Head 都有独立的 K 和 V但它们可以投影和反投影到同样且共享的 Latent KV。KV Cache 和 MQA 相当效果和 MHA 相当但也会额外的增加一些计算量。主流之一。Sparse AttentionSparse Attention 是 Softmax Attention 的一种变体。区别在于Softmax Attention 会使用全部的 Q·K 元素稠密而 Sparse Attention 只会使用一部分 Q·K 元素。假设 seq_len 是 n那么 Self-Attention 的 QK 计算就会产生一个 形状为 [n, n] 的注意力得分矩阵相似度矩阵。所以从理论上来讲Self Attention 的计算时间和显存占用量都是 O(n^2)。也就是说围绕该矩阵的计算量和显存占用量Scaled、Masked、Softmax、乘 V 加权融合计算等等会随 n 呈平方级增长例如如果 seq_len 变成原来的 2 倍显存占用量就是原来的 4 倍计算时间也是原来的 4 倍。但实际上Child 等人2019的研究发现在训练好的 Transformer 模型中注意力得分矩阵往往是稀疏的这意味着并不是每个 token 都需要关注其他所有 token每个 token 只关注非常有限个其他 token。有些 token 之间的相互作用可能对最终的输出贡献不大可以被忽略。稀疏注意力机制的核心思想是在自注意力计算中引入稀疏性即不是让序列中的每个位置都与其他所有位置进行注意力计算而是仅选择部分位置进行计算。所以稀疏注意力具有以下优势减少计算量通过减少参与注意力计算的位置数稀疏注意力显著降低了计算复杂度使得模型能够处理更长的序列。减少显存占用量稀疏操作减少了需要存储的注意力权重的数量从而降低了模型的内存需求。提高长距离依赖学习能力某些稀疏模式如分层或跳跃连接可以帮助模型更有效地学习序列中的长距离依赖关系。下图是 Self-Attention 的一个注意力矩阵。左边显示了注意力矩阵右边显示了关联性这表明每个元素都跟序列内所有元素有关联。Atrous Self Attention空洞注意力启发于 “膨胀卷积Atrous Convolution”它对相关性进行了约束强行要求每个元素只跟它相对距离为 k, 2k, 3k 的元素关联其中 k1 是超参数。如此的运行效率和显存占用都变成了 O(n^2/k) 也就是说能直接降低到原来的 1/k。Local Self Attention局部自注意力约束每个元素只与前后 k 个元素以及自身有关联。保留了一个 2k1 大小的窗口每个元素只跟 2k1 个元素算相关性这样一来理想情况下运行效率和显存占用都变成了 O(kn)也就是说随着 n 而线性增长非指数增长。这是一个很理想的性质当然也直接牺牲了长程关联性。Sparse Self Attention稀疏自注意力将 Atrous Self Attention 和 Local Self Attention 合并为一个除了相对距离不超过 k 的、相对距离为 k,2k,3k,… 的注意力都设为 0这样一来 Attention 就具有了 “局部紧密相关和远程稀疏相关” 的特性。FlashAttention由 Stanford DAWN Lab 实验室提出https://arxiv.org/pdf/2205.14135。其核心思想是通过优化 GPU 内存访问来大幅提升速度和降低显存占用同时保持与标准注意力完全相同的计算结果不损失精度。具体而言将 Q/K/V 分块加载到 SRAM 中计算例如 GPU SM Shared memory避免将完整 N×N 注意力矩阵写入 HBM以减少访问 GPU 片外的全局内存 HBM 的频率。GPU 内存是分层的HBM高带宽显存 容量大但速度较慢SRAM片上缓存 速度极快但容量很小。标准注意力计算需要将庞大的中间矩阵如 N×N 的注意力分数矩阵N 为 seq_len写入 HBM 再读出这种频繁的数据搬运消耗了大量时间使得计算单元如 Tensor Core常常处于“等待数据”的空闲状态。FlashAttention 正是为了解决这一“内存墙”问题而设计的。核心是IO 感知IO-Awareness即算法设计需要充分考虑内存读写代价它通过两项关键技术实现这一点分块Tiling将大的注意力计算分解为小块。算法不再一次性计算并存储完整的 N×N 矩阵而是将 Query (Q)、Key (K)、Value (V) 矩阵分割成小块逐块加载到高速的 SRAM 中进行计算。这样可以避免在慢速 HBM 上 “物化” 那个巨大的中间矩阵极大减少了 HBM 的读写次数。重计算Recomputation在反向传播时FlashAttention 不存储前向传播中产生的大量中间激活值如注意力矩阵而是在需要时重新计算它们。这牺牲了少量计算量但换来了显存占用的大幅降低。为了在分块计算的同时保证 Softmax 的数值稳定性算法还使用了在线 SoftmaxOnline Softmax 技术在逐块处理时动态更新统计量。FlashAttention 已成为当前大模型训练与推理的事实标准被 PyTorch、Hugging Face 等主流框架广泛集成是支撑现代大语言模型处理长上下文的关键底层技术之一。PageAttentionPaged Attention页面注意力由 UC Berkeley 团队提出并作为 vLLM 推理引擎的核心技术https://arxiv.org/pdf/2309.06180。在 LLM 推理时为了不重复计算历史信息系统会缓存每个 token 的 Key 和 Value 向量这就是 KV Cache。随着对话进行KV Cache 会动态增长。传统系统要求为每个请求预留一块连续的显存空间但这会带来严重的浪费。内部碎片为应对可能的最大长度系统会过度预留空间但实际生成长度往往短得多导致预留空间大量闲置。实验表明在传统系统中KV Cache 的实际有效利用率可能低至 20.4%。外部碎片不同请求释放后留下的内存空洞因大小不一而难以被后续请求利用。这些浪费严重限制了同时处理的请求数量Batch Size从而拉低了 GPU 的整体吞吐量。PagedAttention 的核心思想是借鉴操作系统的虚拟内存分页技术来管理 KV Cache。它不再要求连续空间而是将每个请求的 KV Cache 分割成固定大小的块Block每个块包含固定数量如 16 个token 的 KV 向量。这些块在物理显存中可以非连续存储。系统通过一个 “Block Table” 来记录逻辑顺序与物理块的映射关系注意力计算时会根据块表动态地查找所需的 KV 块。RadixAttentionRadixAttention基数注意力由 SGLang 团队提出。其核心思想是高效地复用多个请求间的公共前缀Prefix的 KV Cachehttps://arxiv.org/pdf/2312.07104。对比来看PagedAttention 解决了 KV Cache 的碎片化问题而 RadixAttention 则更进一步解决了跨请求间重复计算的问题。所以RadixAttention 通常和 PagedAttention 互补。PagedAttention 提供了底层的、无碎片的 “块管理” 机制而 RadixAttention 则在其上增加了 “前缀组织” 的智能决定哪些块可以被复用。在许多真实场景中大量请求共享着相同的开头部分。例如多轮对话每一轮对话都需携带之前所有轮次的历史历史部分的 KV Cache 在后续每一轮中都被重复计算。长系统提示System Prompt许多应用会使用相同的系统指令或角色设定。RAG检索增强生成多个查询可能基于同一份检索到的长文档。传统系统在处理完一个请求后其 KV Cache 就会被丢弃。这意味着下一个共享相同前缀的请求必须从头开始计算整个序列造成了巨大的计算浪费和首 Token 延迟。RadixAttention 的核心思想是不再按请求为单位管理 KV Cache而是将所有请求的 token 序列组织在一棵全局的基数树Radix Tree一种压缩前缀树中。它将 token 前缀本身作为树上的路径公共前缀在树中只存储一次其对应的 KV Cache 也随之被共享。当一个新请求到达时系统会执行四个步骤遍历匹配在基数树中查找与新请求 token 序列匹配的最长公共前缀。复用缓存如果找到匹配直接加载该前缀对应的 KV Cache跳过这些 token 的 Prefill 计算。计算后缀只对未匹配的新后缀部分执行模型的前向计算。插入更新计算完成后将新的 token 序列及其 KV Cache 插入树中供后续请求复用。FFN 优化方向 —— MoEMoE 的诞生背景和思想标准 Transformer 采用的是 Dense稠密 FFN。所谓 “稠密” 指的是所有 token 的每层 FFN 都要完整计算所有参数包括权重和偏置项都会被激活所有参数都被用于计算输出不跳过任何部分。Dense FFN 的参数量通常占整个模型参数量的 2/3因此训练一个 Dense 模型往往需要庞大的 GPU 算力FLOPs这也在一定程度上限制了模型参数量的增长。想扩大模型就必须同比例增加算力。后来研究人员发现FFN 在计算过程中存在明显的神经元激活稀疏性。推理阶段对单个 tokenFFN 中大部分神经元的激活值接近于零或贡献极小。例如以 T5-LargeReLU 激活为例90% 的输入只激活了不到 5% 的神经元。又例如在采用 SwiGLU 的现代 LLM 中每个 token 约有 30%–40% 的神经元提供了不可忽略的贡献其余 60%–70% 的神经元贡献极小。训练阶段神经元激活稀疏性是动态演化的。研究表明FFN 的激活稀疏度在训练初期约为 0.5在大约 20,000 步后迅速上升并稳定在 0.9 左右并在后续训练中保持稳定。其中推理时的神经元激活稀疏性的特点是研究员设计 MoE 架构的核心动机 —— 既然大部分神经元对当前 token 没有贡献就可以将它们 “打包” 成不同的专家并设计路由机制每个专家在训练过程中学习不同的信息。而在推理时仅使用与当前任务最相关的特定专家即让每个 token 只激活少数专家从而大幅降低计算量。具体而言MoE 将原本单个巨大的 FFN 拆分成多个 experts FFN 子网络并通过一个可学习的 Router 路由网络为每个 token 选择最合适的 Top-K 个专家进行计算。这样虽然模型的总参数量很大但每个 token 实际参与计算的参数量却很小而且还能保持近似的效果。如下图DeepSeek MoE 16B 推理时候只用到了 2.8B 的参数整体的 FLOPs 是 LlaMA2 Dense 7B 的 39.6%。推理速度更快的同时效果也不差。可以说MoE 将 “模型总参数量” 与 “激活参数量单次计算量” 进行划分前者决定了显存资源后者决定了计算资源。推理时在模型总参数量相同的情况下MoE 的单次计算量显著低于 Dense。但需要注意的是MoE 的显存需求与 Dense 依旧相当因为 MoE 不省权重显存主要省激活计算整体显存仍与总参数量同量级但具体取决于并行和实现。Dense 与 MoE 的对比Dense 和 MoE 并不是简单的演进或替代两者各有特点和场景。Dense 模型的核心优势是 “简单与可预测”。它每个 token 都激活全部参数计算路径固定。因此它非常适合单 GPU 或小型多 GPU 部署、对延迟稳定性要求高、以及资源受限或需要快速迭代的场景。NVIDIA 的官方分析也指出Dense 模型 “通常有利于更简单、更可预测的部署”。MoE 模型的核心优势是 “以更低的计算成本换取更大的模型容量”。它通过稀疏激活在推理时只调用一小部分参数。因此它适用于大规模、高吞吐量的场景例如大规模多语言服务、知识密集型应用等在这些场景下它能在可控的内存成本下提供更高的吞吐量。当 Dense 模型和 MoE 模型的总参数量相同时显存组成DenseMoE权重参数量相同相同梯度、优化器状态数据量训练基本相同基本相同KV Cache 数据量推理数据量相同相同激活值较大全部 FFN 参与较小仅激活专家参与总体显存基本相当Dense 可能略大基本相当通信开销小大需要 all-to-all 等专家通信MoE 训练特点前向稀疏每个 token 只经过少数专家。反向稀疏只有被激活的专家收到梯度更新未被激活的专家在该次迭代中不参与学习。负载均衡通过 LB 辅助损失确保所有专家在训练过程中被均衡地激活防止专家坍缩。显存需求仍需存储全部专家参数、优化器状态和激活值。以 AdamW 为例优化器状态显存占用通常是模型参数本身的数倍。通用估算公式总显存 ≈ 参数量 × 16 字节2 字节权重 2 字节梯度 12 字节 AdamW 优化器状态。MoE 推理特点前向稀疏每个 token 只激活少数专家一个 token 的单次计算量大幅降低。显存需求由总参数量决定。系统无法预知下一个 token 会用到哪个专家因此所有专家权重通常都需要常驻显存所以权重内存不会因为稀疏激活而自动减少。推理优化实际部署时可通过专家并行、量化、CPU offload 等技术在较少 GPU 上运行 MoE但通常仍需要较大的显存或多卡支持。MoE FFN 的结构MoE 还可以细分为 2 种类型稀疏专家混合模型Sparse Mixture of Experts密集专家混合模型Dense Mixture of Experts两者都具有类似的机构都使用路由器来选择专家但前者只选择少数几个专家而后者则选择全部专家。显然目前的 LLM 中 MoE 通常指的是稀疏 MoE。如上图MoE FFN 由 2 个关键部分组成。专家网络Experts NetworkMoE 中的每个专家是一个独立的 FFN 神经网络所有专家参数量之和构成 MoE FFN 层的总参数量。DeepSeek V3 的 experts 总数为 2561 个每个专家的 Hidden 维度为 2048。又细分为 2 类专家路由专家Routed Experts每个 MoE 层包含 256 个路由专家这些专家主要负责处理输入中某些特定、专业化的特征。共享专家Shared Expert每个 MoE 层中还有 1 个共享专家用于捕捉通用的、全局性的知识为所有输入提供基本的特征提取支持。推理时每个 Token 激活 8 个路由专家并且确保每个 Token 最多被发送到 4 个节点。需要注意的是这些 “专家” 并不像人类生活中的 “心理学” 或 “生物学” 专家那样在特定学科上表现出高度专业化。也就是说专家并非人为预先分类而是按 token 类型或上下文中的语义功能划分的。例如Expert 1 专注于处理标点符号Expert 2 专注于处理动词Expert 3 处理连词Expert 4 处理视觉描述相关的词汇。在训练过程中不同专家逐渐在不同数据模式上分化形成了专业化分工。具体如下图Mixtral 8x7B 论文中每个 token 都被其选择的第一个专家进行了着色其中代码中的缩进标记总是分配给相同的红色或黄色专家。可见某些专家可能专门处理缩进相关的 token而另一些则专门处理关键字或变量名。门控网络Gating Network门控网络Gating Network是 MoE 的大脑负责为每个 token 决定 “派发给哪些专家” 进行计算。需要注意的是Router 本身也是训练得到的本质就是一个线性层 [hidden_size, number_of_experts]。训练时它与专家网络同步接收梯度更新所以 Router 在训练过程中逐渐学会将不同类型的 token 分发给更擅长处理它们的专家。具体而言在训练的早期门控网络的路由决策是近似随机的但随着专家逐步积累专长门控网络也会调整其路由策略正反馈循环如果某个专家因早期获得较多特定类型数据而表现出色门控网络便倾向于将更多此类数据路由给它。专家和路由网络的协同演化专家因接收到较多特定数据而 “专长”而门控网络根据反馈不断更新参数使得路由更加精准​。推理时一个 token 可以被发送到多个专家所以 Router 的输出是一个对所有专家的打分向量经过 Softmax 归一化。在 Top-K 稀疏路由中每个 token 只被分配给得分最高的 K 个专家。训练或推理时Router 将输入 x 与权重矩阵 W 相乘然后对输出应用 Softmax 操作为每个专家创建一个概率分布 G(x)。Router 利用这个概率分布来为给定的输入选择最匹配的专家最后将每个 Router 的输出与各自选择的专家输出相乘并将结果相加。负载均衡损失函数Load Balancing Loss实验中发现不同 experts 在竞争的过程中会出现 “赢者通吃” 的现象前期变现好的 expert 会更容易被 gating network 选择导致最终只有少数的几个 experts 真正起作用。即导致 “专家坍缩” —— 其余专家得不到充分训练。因此需要额外增加了一个 LB loss 来缓解这种不平衡现象。专家级别的负载均衡目的是防止训练时总是选择少数专家负载均衡确保所有专家都能被充分训练。设备级别的负载均衡确保计算负载均匀分布在多个 GPU 设备上防止个别设备计算负载过重影响性能。网络级别的负载均衡因为专家比较多分布在不同的 GPU 上通信成本会很高。所以最好确保每个 token 的目标专家分布在最多 M 个设备上。如下右图充分利用 GPU 机内互联带宽。专家并行Expert Parallelism由于总参数量大单卡显存往往装不下所有专家。专家并行将不同专家分散到不同 GPU 上token 通过 all-to-all 通信被发送到对应专家所在的 GPU。这引入了通信开销是 MoE 训练和推理的重要瓶颈之一。如图所示一个包含 6 个专家的 MoE 模型在 EP2 时的专家分布情况。DeepSeek-V3 的训练使用了 2048 张 H800 GPU下图中一共有 1024 张卡两个 DP 共计 2048 张卡。使用了 PP16即 64 层Embd、MTP、LMHead 也各算一层划分到 16 组机器上每组机器 8 卡每张卡 4 层。使用了 EP64将 64 个专家分 8 组每组 8 个。即每层的 256 个专家分布到 8 台机共 64 张卡上每层在一张卡上有 4 个专家。MoE 整体计算流程Routing选择最适合处理输入的专家模型。当输入 token 通过 MoE 层时Token 通过和 Router 的权重矩阵相乘得到一个 Expert Indices决策矩阵和一个概率张量即索引和概率Expert indices 是 expert-to-token 映射是形状 [num_tokens, top_k] 的张量用于指示每个 token 被分配给了哪个 expert即张量中第 i 个值代表本 token 应该分配到第 i 个专家。Probabilities 张量是分配置信度的概率其中第 i 个值代表这个专家对于该 token 最终结果的权重Permutation排列/置换根据路由决策expert-to-token 映射将 Token 分配给对应的专家中间可能会有 drop 操作。Computation每个专家网络并行处理其分配到的 token计算输出。把输入矩阵 与专家网络的权重矩阵相乘 × 。Un-Permutation收集专家的计算结果。这是 Permutation 的逆运算将从各个 experts 收集到的处理后的 tokens 组合成一个完整的序列这个序列保持了原始 tokens 的顺序。即将每个专家网络的输出根据原始的 token 顺序重新排列。接着使用 Routing 步骤生成的分配置信度概率对结果进行加权求和以得到最终的模型输出然后将这个结果继续向下游处理。