现代 AI 的计算困境
大型语言模型 (LLMs) 作为当代人工智能的支柱，展现出非凡的能力，正在重塑各行各业和科学发现。它们在生成类人文本、驱动复杂的对话代理，甚至辅助复杂研究任务方面的熟练程度，使其成为不可或缺的工具。这些强大模型的核心是 transformer 架构，这是一种以其交替层为特征的设计。输入数据被分解为 tokens，流经一系列 attention 机制（用于衡量不同 tokens 的重要性），然后是 feed-forward networks (FFNs)（用于处理收集到的信息）。这种分层的、顺序的处理方式是 transformers 学习和生成输出的基础。
然而，正是这种架构，虽然有效，但随着模型规模和复杂性的膨胀，带来了日益严峻的挑战。其顺序性意味着每一层通常必须等待前一层完成计算后才能开始。这种逐步处理方式造成了固有的瓶颈，尤其是在推理阶段——即使用训练好的模型实际生成预测或文本的阶段。随着像驱动高级 AI 助手那样的模型包含数千亿甚至数万亿的 parameters，推理所需的计算资源和时间急剧增加。这种不断升级的需求转化为显著的 latency（响应延迟）、降低的 throughput（单位时间内处理的请求数量）和不断增加的运营成本，阻碍了最强大 LLMs 的广泛部署和实时应用。因此，提高推理效率已成为 AI 研究界的首要关注点，推动了对创新策略的探索，这些策略旨在简化计算，同时不损害这些模型提供的卓越性能。核心挑战在于减轻顺序执行带来的限制，尤其是在计算跨越多个 GPUs 的分布式环境中，这会给处理时间增加通信开销。
探索优化版图：现有工具及其局限性
在使 LLMs 更精简、更快速的持续努力中，研究人员开发了一套优化技术工具包。每种技术都提供了一条通往效率的途径，但往往伴随着自身的妥协，使得没有任何单一方法能成为通用解决方案。理解这些权衡对于认识到像 FFN Fusion 这样的新方法的必要性至关重要。
一种突出的技术是 quantization（量化）。这涉及降低用于表示模型权重和激活值的数值精度。模型可能使用 16 位、8 位甚至更低位的表示，而不是标准的 32 位浮点数。这直接缩小了模型的内存占用，并能显著加快计算速度，因为对较低精度数字的操作通常更快且能耗更低。然而，quantization 并非没有风险。降低精度可能导致信息丢失，从而可能降低模型的准确性。在非常低的位宽下，这种风险变得更加明显，需要仔细实施，有时还需要重新训练以减轻准确性下降。挑战在于找到能够最大化效率增益同时将性能下降控制在可接受范围内的最佳平衡点。
另一种常用策略是 pruning（剪枝）。该技术基于这样一个原则：大型神经网络中的许多 parameters 可能冗余或对最终输出贡献甚微。Pruning 算法识别并移除这些不太重要的连接或神经元，从而产生一个更小、更稀疏的模型。与 quantization 一样，pruning 减少了内存需求和计算负载。然而，精确识别哪些 parameters 是‘安全’可移除的非常复杂。激进的 pruning 可能无意中移除关键组件，导致显著的准确性损失。Pruning 后通常需要对模型进行微调以恢复性能，这增加了工作流程的复杂性。必须进行仔细校准，以确保剪枝后的模型仍然有效。
一种在架构上更独特的方法是 Mixture-of-Experts (MoE) 模型。MoE 模型不是让每个输入都通过整个网络进行处理，而是由多个‘专家’子网络（通常是 FFNs）组成。对于每个输入 token，一个门控机制会动态选择这些专家中的一小部分来执行计算。这种条件计算意味着对于任何给定的输入，只有模型总 parameters 的一小部分被激活，从而显著节省计算量，尤其是在训练和推理非常大的模型时。MoE 模型可以扩展到数万亿 parameters，同时保持合理的计算成本。然而，它们的效率高度依赖于工作负载。它们擅长处理非常大的批量大小，此时选择性激活模式能带来良好的硬件利用率。在较小或中等批量大小时，MoE 模型可能会因计算资源利用不足而受到影响，因为并行硬件可能无法被稀疏激活的专家持续占满。此外，实现和负载均衡 MoE 模型可能比部署标准的‘密集’架构更复杂。
虽然 quantization、pruning 和 MoE 模型代表了 LLM 优化方面的宝贵进展，但它们固有的局限性凸显了对替代或补充策略的需求。人们仍在继续寻求能够在各种场景下提供广泛效率改进的方法，理想情况下，这些方法对准确性或实现复杂性的妥协更少，特别是对于那些因训练和部署相对简单而仍然流行的密集模型架构。
FFN Fusion：重新思考 Transformers 中的并行性
在这一优化技术的版图中，NVIDIA 的研究人员引入了一种引人注目的新方法，称为 FFN Fusion。这项技术直接挑战了 transformer 架构固有的顺序瓶颈，不是通过改变 parameters 或选择性地激活部分，而是通过从根本上重新思考如何并行化计算序列。这项创新源于对深度 transformer 模型中 FFN 层行为的一个关键观察。
研究人员使用一个名为 Puzzle 的诊断工具分析了大型模型的内部工作机制。当他们实验性地移除 attention 层时，他们注意到模型通常会保留令人惊讶的长序列的连续 FFN 层。更重要的是，分析显示这些相邻 FFNs 执行的计算经常表现出极小的相互依赖性。本质上，序列中一个 FFN 的输出通常不会显著改变紧随其后的 FFN 所需的方向路径或核心信息。这表明，这些传统上一个接一个执行的 FFNs，可能具有同时、并行执行的潜力，而不会显著干扰模型的整体功能。
这一洞察构成了 FFN Fusion 的基石。核心思想优雅简单却功能强大：识别具有低计算依赖性的连续 FFN 层序列，并将它们合并成一个单一的、更宽的 FFN 层，该层并行执行等效的计算。结构不再是 Input -> FFN1 -> FFN2 -> FFN3 -> Output 这样的链条，而是变成了 Input -> Fused_FFN (等效于 FFN1+FFN2+FFN3 并行) -> Output。这种架构转换有效地缩短了网络的顺序深度，用一个单一、更宽的计算步骤取代了多个步骤。通过针对这些低依赖性的 FFN 序列，FFN Fusion 旨在减少 latency 和计算成本，同时保留模型的表示能力和准确性。从 Llama-3.1-405B-Instruct 开发出 Ultra-253B-Base，正是这项技术潜力的一个主要展示。
架构炼金术：FFN Fusion 如何工作
FFN Fusion 背后的魔力在于其对 feed-forward networks 底层数学结构的巧妙操纵。它不仅仅是让现有层并排运行；它涉及到创建一个新的、统一的层，该层复制原始序列的集体行为，但却是并发执行的。
考虑一个包含 k 个连续 FFN 层的序列。在标准 transformer 中，输入 x 通过 FFN1，其输出成为 FFN2 的输入，依此类推，直到 FFNk。每一步都明确依赖于前一步的完成。FFN Fusion 打破了这条依赖链。数学上，一个 FFN 通常涉及两个线性变换，中间有一个非线性激活函数（如 GeLU 或 SwiGLU）：FFN(x) = W_out * Activation(W_in * x)。FFN Fusion 利用了线性变换通常可以组合这一事实。
融合过程通过连接各个 FFN 层的权重来实现。具体来说，连续 FFNs 的输入权重矩阵 (W_in) 被组合（例如，以块对角线方式）成融合层的单个、更大的输入权重矩阵。类似地，输出权重矩阵 (W_out) 被连接起来形成一个单一、更宽的输出权重矩阵。激活函数在这个更大的结构内按元素应用。这种构造确保了融合后的 FFN 能够同时在对应于原始 FFNs 的并行路径上操作原始输入 x。来自这些并行路径的输出随后通过连接的输出权重的结构被隐式地聚合起来。
理论基础证实，只要原始层之间的依赖性确实很低，这种融合结构就能保持与原始 FFN 序列相同的表示能力。关键在于识别哪些序列适合融合。为了系统地做到这一点，NVIDIA 的研究人员采用了一种依赖性分析技术。他们测量了一组代表性输入 tokens 的连续 FFN 层输出隐藏状态之间的cosine distance（余弦距离）。小的 cosine distance 表明一个 FFN 的输出向量指向与序列中下一个 FFN 的输出向量非常相似的方向。这种相似性表明功能依赖性较低——第二个 FFN 并没有显著改变第一个 FFN 建立的信息表示。在各层之间表现出持续低 cosine distance 的 FFN 序列被确定为融合的主要候选对象，因为合并它们不太可能破坏模型学习到的表示和整体性能。这种数据驱动的方法允许将 FFN Fusion 有针对性地应用于模型中最有效且干扰最小的部分。
从庞然大物到短跑选手：Ultra-253B-Base 的转变
FFN Fusion 的实际威力通过将其应用于当时公开的最大模型之一 Llama-3.1-405B-Instruct 而得到了生动的展示。这个拥有 4050 亿 parameters 的模型，在推理方面代表着一项重大的计算任务。研究人员着手进行架构优化过程，将 FFN Fusion 与策略性 pruning 相结合，创建了一个新的、更高效的模型，命名为 Ultra-253B-Base。
转变过程涉及几个步骤：
- 分析 (Analysis): 使用他们的依赖性分析工具（测量 cosine distance），研究人员在 Llama-405B 架构内识别出表现出低层间依赖性的连续 FFN 层序列。
- 融合 (Fusion): 这些被识别出的 FFN 序列随后被融合成单一、更宽的 FFN 层，如前所述（连接权重）。这直接减少了网络中的顺序步骤数量。
- 剪枝 (Pruning): 同时或随后，被认为不太关键的 parameters（可能通过标准 pruning 技术识别，或受融合过程启发）从模型中移除。
这种组合方法产生了 Ultra-253B-Base，一个拥有 2530 亿 parameters 的模型。这比原始的 405B 模型减少了超过 37% 的 parameters，是一个显著的缩减。通过融合实现的架构更改是实现如此显著规模缩减同时力求保持性能的关键。目标不仅仅是创建一个更小的模型，而是一个从根本上更快、计算成本更低的模型，这得益于 FFN Fusion 解锁的增强并行性。这个案例研究作为一个关键的概念验证，表明大型模型可以通过结构重组来大幅提高效率。
衡量收益：性能、速度和资源节省
任何优化技术的真正考验在于其可衡量的影响。对于 Ultra-253B-Base 而言，将 FFN Fusion 和 pruning 应用于 Llama-405B 基础模型所获得的结果令人信服，显示出在多个维度上的显著改进，而没有对能力造成实质性妥协。
推理速度和成本 (Inference Speed and Cost): 最显著的收益体现在推理效率上。与原始的 405B parameter 模型相比，Ultra-253B-Base 实现了：
- 推理 latency 提升 1.71 倍。这意味着模型可以更快地生成响应，这对于实时应用至关重要。
- 在批量大小为 32 时，每个 token 的计算成本降低了 35 倍。每个 token 的计算操作 (FLOPs) 大幅减少，直接转化为更低的能耗和为模型提供服务所需的硬件要求降低。
模型性能基准 (Model Performance Benchmarks): 至关重要的是，这些效率提升并没有以牺牲模型的智能或能力为代价。Ultra-253B-Base 在一套标准的 LLM 基准测试中进行了严格评估，其得分与原始的、大得多的模型高度竞争，在某些情况下甚至超过了后者：
- MMLU (Massive Multitask Language Understanding): 85.17%
- MMLU-Pro (更具挑战性的版本): 72.25%
- Arena Hard (针对困难提示的人类偏好评估): 84.92%
- HumanEval (代码生成能力): 86.58%
- MT-Bench (多轮对话质量): 9.19
这些分数表明，尽管只有 2530 亿 parameters，融合和剪枝后的模型仍保持了与其 405B parameter 祖先相当的非常高的理解、推理、编码能力和对话质量水平。
内存效率 (Memory Efficiency): 除了计算速度和成本之外，FFN Fusion 也有助于节省内存。架构上的改变，可能与融合带来的其他优化相结合，使得推理过程中所需的键值 (KV) cache 的大小减少了 2 倍。KV cache 存储中间激活（attention keys 和 values），会消耗大量 GPU 内存，尤其对于长输入序列。将其需求减半使得在内存要求较低的硬件上运行模型成为可能，或者在相同的内存限制内处理更长的上下文。
这些可量化的结果强调了 FFN Fusion 的有效性。它使得创建一个不仅更小，而且在速度、计算操作和内存使用方面从根本上更高效的模型成为可能，同时在具有挑战性的基准测试中保持顶级性能。
保存知识：训练和微调的关键作用
通过像 FFN Fusion 和 pruning 这样的技术对像 Llama-405B 这样的大型、预训练语言模型进行架构修改，不可避免地会打破其学习到的 parameters 的微妙平衡。虽然数学上的等价性旨在局部保留功能，但网络的全局行为可能会发生变化。为了确保由此产生的 Ultra-253B-Base 模型不仅变得更高效，而且还保持其高水平的性能，一个精心策划的修改后训练过程至关重要。
这个过程涉及两个主要阶段：
知识蒸馏 (Knowledge Distillation): 第一步是将知识从原始的、更大的模型（或合适的教师模型）转移回修改后的架构中。这是通过蒸馏实现的，即训练 Ultra-253B-Base 模型模仿教师模型的输出或内部表示。这个阶段使用了大量的数据集，特别是 540 亿 tokens，并使用 8k 上下文窗口进行处理。蒸馏有助于融合和剪枝后的模型重新捕捉在架构更改过程中可能受到轻微干扰的细微差别和能力。
分阶段微调 (Staged Fine-Tuning): 蒸馏之后，模型经历了一系列微调阶段，专门设计用于使其适应处理逐渐增长的上下文长度。这对于现代 LLMs 至关重要，因为它们通常需要基于广泛的输入来处理和生成文本。微调分阶段进行：
- 在 16k 上下文窗口下进行微调。
- 在 32k 上下文窗口下进行进一步微调。
- 在 128k 上下文窗口下进行最终微调阶段。
这种分阶段的方法使模型能够逐渐调整其 parameters，包括新形成的融合 FFN 层和优化的 KV cache 机制，以有效地管理非常长序列上的依赖关系和信息流。每个阶段都建立在前一个阶段的基础上，确保在不同上下文大小下具有稳定性和鲁棒的性能。
这种细致的训练方案，结合了大规模蒸馏和分阶段的长上下文微调，在弥合架构效率和高保真性能之间的差距方面发挥了重要作用。它确保了 FFN Fusion 带来的速度、成本和内存优势不会损害模型在要求苛刻的基准测试中的准确性和能力。
更广阔的视野：通用性和未来方向
将 Llama-405B 成功转化为 Ultra-253B-Base 为 FFN Fusion 的潜力提供了有力证据，但其真正价值在于其更广泛的适用性以及它为未来 LLM 设计提供的见解。研究表明，这不仅仅是一个仅适用于巨大模型的一次性技巧。
跨规模验证 (Validation Across Scales): NVIDIA 的研究人员明确地在不同规模的模型上测试了 FFN Fusion 方法。他们成功地将该技术应用于 70B-parameter 模型，相对于其原始对应模型实现了类似的效率提升。他们还报告了在 49B 规模上的验证，进一步强化了 FFN 独立性和融合潜力并非最大模型的专属特性，而可能是 transformer 架构更普遍的属性，可能在 FFN 序列自然更深的更大规模上变得更加明显。这表明 FFN Fusion 可能成为 LLM 优化工具箱中的标准工具，适用于各种规模的模型。
FFN 与完整块融合 (FFN vs. Full Block Fusion): 该研究还揭示了 FFN 层与 transformer 块内 attention 层相比的具体作用。虽然连续的 FFN 层通常表现出低依赖性，使其成为融合的理想选择，但尝试并行化整个 transformer 块（包括 attention 和 FFN 层）被证明更具挑战性。分析表明涉及 attention 机制的相互依赖性更强。同时融合整个块会导致更显著的性能下降，这表明 attention 层在跨 tokens 整合信息方面扮演着更关键、顺序依赖性更强的角色。这一发现有助于界定有效并行化的边界——FFN 序列是沃土，而 attention 机制可能需要不同的优化策略。
对 LLM 架构的影响 (Implications for LLM Architecture): FFN Fusion 提供的不只是一种事后优化技术；它为设计未来的 LLMs 提供了宝贵的见解。发现 FFN 序列通常可以被视为可并行化的单元，挑战了通常支撑 transformer 设计的严格顺序假设。这可能会激发新的、从一开始就具有内在并行友好性的架构。未来的模型可能会设计成具有明确用于融合或并行执行的 FFN 结构，可能导致硬件-软件协同设计，其中 GPU 架构被进一步优化以利用这种类型的并行性。使用 cosine distance 量化层间依赖性的系统方法也为理解和重新设计神经网络结构提供了一个有价值的分析工具。通过证明通过专注于并行化现有组件的深思熟虑的架构重新设计可以实现显著的效率提升，FFN Fusion 为开发既强大又在计算上更可持续的 LLMs 铺平了道路。它突显了一条缓解前沿 AI 不断升级的资源需求的途径。