Kimi K3技术报告揭秘：Infra技巧如何提升大模型训练效率？

惊！Kimi K3公布技术报告

Kimi K3近期终于公布了技术报告，不仅涉及模型架构，还透露了重要的Infra层面的不少细节。Infra指的是基础设施层，若把硬件比作手机本体，AI比作应用软件，那AI Infra就是中间的操作系统，没有它硬件就无法发挥作用。

Infra：国产开源大模型商业化的关键支撑

行业普遍认为，国产开源大模型未来商业化有一个重要支撑点，即大模型厂商自己部署模型，成本可比第三方部署低好几倍，依靠的便是Infra上的技术和经验差距。OpenAI核心工程师翁家翌也曾提到，大模型训练的核心差距在于Infra，工程师主要工作是给Infra修bug。当模型架构逐渐趋同、公开数据和蒸馏技术降低部分能力复现门槛后，Infra工程能力会成为拉开训练成本、部署效率和运行稳定性的重要差距之一。

技术报告是否说透Infra秘密？

从商业角度估计，Kimi K3的技术报告对Infra的秘密一定有所保留。这意味着一百个公司部署一百个Kimi K3，可能有一百种不同的运行效率。不过，Kimi K3公开了那么多Infra细节，值得学习。刚好，Infra对很多非从业者来讲很陌生，所以将借Kimi K3技术报告中一些有趣的Infra技巧，科普一下AI Infra的核心原则。

Infra为何重要？

OpenAI核心工程师翁家翌说，大模型训练过程中工程师主要工作是给Infra修bug，这里的bug既包括导致结果错误、训练不稳定或任务中断的传统bug，也包括虽不影响数学正确性，但会造成GPU等待、内存浪费、通信阻塞和吞吐下降的系统性问题。更具体而言，是在保证计算逻辑正确的前提下，看GPU是被充分利用还是有大量闲置算力，就算GPU被充分利用，还要看是否存在大量不必要的重复计算，要看GPU上的计算负载是否均衡、是否稳定等。要满足这些要求，还需要大量的通信在GPU之间做负载协调，这时候通信又可能遇上新的瓶颈。这些bug最终影响集群吞吐量，并反映训练效率。模型浮点运算利用率 (MFU) 是评估训练效率的标准指标，它是 “观测吞吐量” 与 “理论最大吞吐量” 的比值，理论最大吞吐量是假设达到峰值浮点运算的100%。从行业角度看，大模型的训练效率仍有很大提升空间，目前公开可验证的万卡、千亿参数、完整LLM预训练案例中，字节跳动的MegaScale的55.2% MFU仍是最有代表性的最高纪录之一。今年4月甚至出现一个极端反面案例，据《The Information》报道，马斯克旗下xAI坐拥约55万张英伟达GPU，但内部备忘录显示其MFU仅11%，远低于业界35%-45%的正常水准 ，也不及Meta（43%）、谷歌（46%）等竞争对手。

高训练效率的意义

拥有更高的训练效率意义重大，它可以直接转化为更短的训练周期、更低的训练成本、更多可尝试的实验、更大的数据量或模型规模。由于没有足够的数据和参数参考，很难自主判断技术报告中哪些Infra技术是最重要的或贡献最大的。主要从Kimi K3自主提出的架构设计入手，如KDA、AttnRes、MoonEP，看看其背后的Infra技巧体现的原则。

Kimi K3的分布式系统设计

先来总体看看Kimi K3的分布式系统设计，其训练采用的并行方法都是比较经典的方法，如流水线并行、专家并行、ZeRO分片等。大模型训练的 “并行化” 简单理解就是，大模型训练过程包含大量串行过程，从前向传播到反向传播，但细节上存在大量矩阵计算，矩阵计算一般可并行化。此外，一张卡的内存很难装下超大参数规模的所有数值。不严谨地说，就像盖房子必须从底下一步一步往上盖（串行过程），但盖房子过程中可以同时搬好多水泥和钢筋（矩阵计算的并行）。为此，大模型训练通常会从模型参数、训练数据和中间激活值等不同维度进行切分，将计算和存储任务分配到多张GPU上，并通过合理的并行策略和调度机制，尽量减少GPU等待和空转。例如，把模型的不同神经网络层放到不同GPU上依次计算，叫流水线并行。由于后面的GPU必须等待前面的GPU产生中间结果，流水线中容易出现部分GPU空闲的 “气泡”。实际训练中通常会把一个batch切成多个微批次，让不同GPU同时处理不同微批次，从而提高流水线利用率。如果模型的某一层本身规模很大，单张GPU无法容纳，或者单卡计算速度不足，还可以把这一层中的权重矩阵和矩阵运算拆分到多张GPU上共同完成，这叫张量并行。张量并行会在同一层计算过程中频繁交换数据，因此对GPU之间的互联带宽和通信延迟要求很高。当通信耗时接近甚至超过计算耗时时，就需要在并行规模、计算效率和通信开销之间进行权衡。不同的大模型集群会有自己独有的 “烦恼”，特别是模型越大，越可能出现在小模型中不会出现的问题，因为某个原先不被注意的计算进程或内存存储对象会因为规模的增大，而突然变得对于GPU容量而言显著了。顺便提一个细节，技术报告显示：“MoE层使用在各EP rank之间复制的共享专家”，这意味着在专家并行（EP）设置下，每个GPU都保存一份参数完全相同的共享专家，以及不同的路由专家，这样可以让共享专家尽可能吸收数据中的共有特征（rank可以理解为一个参与分布式通信的计算进程，实际部署中通常一个rank对应一张GPU，为简单起见把rank都理解为单个GPU）。

KDA：混合注意力机制的巧思

首先是混合KDA注意力机制。KDA全称Kimi Delta Attention，是一种线性注意力机制。标准注意力机制的核心在于注意力的计算，用每个输入token计算出对应的Query、Key、Value向量（简称Q、K、V）。借用搜索引擎的比喻，Q、K、V好比搜索引擎中的搜索查询词、网页索引库、网页内容，标准注意力机制会先将搜索查询词、网页索引库进行配对，即Q和K先相乘，然后再去匹配网页内容。这样做能精确匹配所有潜在关系，但计算量大，数学上是一个随着输入长度的增长而平方增长的矩阵，这也是制约当前大模型Scaling Law的核心因素。在自回归推理中，还需要保存随序列长度线性增长的KV cache。KDA做出的改变是，让网页索引库和网页内容先进行配对，即K和V先相乘，再用查询搜索词去匹配。K和V相乘好比是对可检索的网页内容先进行总结，数学上是一个长和宽固定的矩阵，无论输入长度是多少，这是KDA计算效率高的核心机制。在更精细的结构中，KDA继承Gated DeltaNet，加入并精细化了总结网页内容也就是模型记忆库的管控结构，使其可以在推进推理的过程中，合理地忘记部分内容，并记住新内容（所以严谨来说，KDA实际还包含遗忘门，并不是简单地把所有K和V相乘后累加）。KDA绕开了标准注意力机制的完整记忆矩阵，但引入了串行依赖：后一个token的状态必须等前一个token计算完成。所以，这个方法在Kimi K3这个参数体量下，会遇到算力瓶颈，它是串行更新的，而GPU又是天然偏好大规模、高带宽并行计算的，容易导致算力利用率不足和算力浪费。优化方式基本就是把其中串行过程尽可能并行化。于是，月之暗面提出了FlashKDA，FlashKDA把token进行分块（chunk）计算，chunk内部可以并行计算，chunk之间仍需依次传递递归状态。在不同设备（或者GPU）之间，Kimi K3采用更进一步的KDA上下文并行（KCP）方法，来提升并行度。对于标准的线性注意力的记忆递归计算，每个输入序列分段的状态转移可以在不知道输入状态的情况下独立求值，随后再精确合成递归链条。通俗来说，不是让每个序列分段计算自己的 “历史”，而是先让每个序列分段独立算出自己那段 “会怎样改变历史信息”，再把这些结果拼起来，就能还原每一段真正应该接到的前文状态，从而完成 “历史演化” 的计算。但对于KDA，这样并不直接可行，因为其状态转移函数具有记忆遗忘机制，在数学上导致一些额外的困难。KCP的做法是，把状态转移中可以独立计算的部分尽可能拆解出来，再分到不同的GPU进程中去独立、并行计算。

AttnRes：注意力残差机制的优化

其次是注意力残差AttnRes机制。标准残差连接是指，一层神经网络在处理前面层传入的输入时，不是把原来的信息完全替换掉，而是只学习 “需要修改或补充的部分”，再把这部分加回原输入，即 “输出 = 原输入 + 本层计算结果” 。这就像在原稿上做增量修改，而不是每一层都重新写一遍，使原始信息和梯度能沿网络直接传递，从而让很深的模型更容易训练，也减少信息在多层处理中被破坏或遗忘。但随着神经网络深度增加，也会越发稀释靠前的每一层的贡献。注意力残差AttnRes就是为了克服这一点而提出的，使其对原输入的继承更加建立在基于语义理解的筛选上。然而，额外引入的注意力机制在大参数模型中，也会带来新的内存负担，所以月之暗面又提出了Block AttnRes。Block AttnRes将神经网络层分为多个块，块内使用的还是标准的残差连接，块间使用注意力机制，研究表明，在约8个块的情况下，它能够达到AttnRes的大部分性能优势。当然，Block AttnRes本身引入的内存开销也是不可忽视的，仍然需要做进一步的优化。注意力机制是在块间进行的，且需要随着前向传播的演进计算多次。所以块表示只需要计算一次，就可以长期保留在GPU中，并重复使用。比如按6层为一个块，前6层的块表示可以被后6层在注意力计算使用，也可以被后7到12层使用。AttnRes相比标准的残差连接需要额外计算注意力机制相关的中间结果，如注意力分数等，这些结果在反向传播中也要用到，但如果一直保留到反向传播阶段，会占用大量内存。因此Kimi K3选择在前向传播时不在内存中保存这部分中间结果，等反向传播阶段时再重新计算这些数值，属于典型的用额外计算换取显存。此外，由于训练还采用了流水线并行，也就是模型不同的层放在不同的GPU上，而它们在计算注意力残差时需要用到其它块的块表示，这些块表示还保存在其它块所在GPU上。由于这些层之间存在串行关联，因此可以把块表示按照层顺序的轨迹，依次传递过去，并进行缓存，而不需要从源头重复传递。比如把GPU0的Block_0的表示B_0，传递给GPU1的Block_1之后，GPU1将B_0缓存，等Block_1计算出表示B_1后，将B_0、B_1一起传递给GPU2的Block_2，GPU2将B_0、B_1缓存。最后，这些块表示仅限于在一个微批次的训练数据中可重复使用，因此在微批次训练结束后可以释放。综合以上手段，Block AttnRes的优化方案做到了：不重复计算同一个块；不保存可以重算的中间激活；不重复传输已经缓存的块；不保留已经失效的微批次块；因此显存里只留下 “当前时刻确实仍有用的数据”，最终达到Kimi K3所声称的 “内存占用理论下限”。

MoonEP：解决负载均衡难题

最后，看看MoonEP的机制原理。在MoE架构的大模型中，每一个输入token可能激活不同的专家，那么反过来，对于任意一段输入token，输入每一个专家的token数量可能都是不同的，进而造成激活形状的动态变化。然后，在传统专家并行中，每个专家固定放在某个GPU上，而路由器不会保证token被平均分给这些专家。于是，有些GPU要处理大量token，有些GPU很快就做完并等待。由于整个训练步骤必须等待最忙的GPU，整体训练吞吐量被最慢的GPU “拖后腿”。同时，路由专家每轮收到的token数不断变化，使中间张量时大时小，GPU需要频繁申请和释放不同大小的显存块，留下许多难以复用的不连续空洞，从而浪费显存、增加分配开销。在这种内存碎片化的情况下，虽然零散的剩余显存总量可能不少，但单个新的张量或通信缓冲区需要一块足够大的可用空间，分散的小块无法直接拼成一次有效分配。为此，Kimi K3提出了MoonEP方法，也就是直接给每个GPU分配相同的任务量，接收S×K个token，其中S是序列长度，K是每个token选择的专家数量。而一个GPU里不同的专家还是可能分配到不同的任务量，所以需要在GPU中配置冗余专家，以应对随时增加的计算负载。冗余专家本身是一个高负载专家也就是被分配了较多token的专家的复制，作用是为高负载专家分摊计算任务，从而缩短最忙碌GPU的工作时间。月之暗面证明了，每个GPU只需要固定数量的冗余专家（比如专家总数为100，参与专家并行的GPU数量是20，每个GPU最多需要100/20 = 5个冗余专家），便在理论上确保可以实现负载均衡。过去的方案一般会设置固定的冗余专家数量，或设置单个GPU的token上限，容易因为无法满足条件使得训练被迫终止。但要精确应用这个算法也会导致计算成本过高，因此实际落地时还是使用了一些近似优化手段，并确保冗余专家数量不超过上限。

总结：Infra技巧的工程思维

可以看到，以上三个Infra技巧关注不同的方面，KDA关注串行过程的并行化，AttnRes关注计算的重复性，MoonEP关注负载均衡。但把这些技巧放在一起看，是一套贯穿系统的工程思维。凡是串行链条，就寻找可并行拆解的部分，凡是重复计算、存储和通信，就判断能否重算、缓存或增量传输，凡是动态负载，就尽可能把它转化为可预测、可调度的静态形状。AI Infra的竞争并不是简单堆更多GPU，而是不断识别那些被默认接受的等待、复制、同步和闲置，再将它们逐一消除。单个优化看起来可能只节省几个百分点，但当模型扩大到万亿参数、训练持续数月时，每个百分点都会被放大成巨额算力、时间和资金。需要提醒的是，关注的这三个点，相对于技术报告的整个Infra章节也只是冰山一角。Kimi K3公开了方法，但这不是一套可以直接复制的标准答案。不同芯片、网络、模型结构和业务场景，都会产生不同的最优解。技术报告能帮助后来者少踩一些已知的坑，但真正决定部署效率的，仍然是团队能否在长期运行中发现并修复自己遇到的Infra bug。总之，模型架构决定了能力的可能性，而Infra决定了这种可能性能否以可承受的成本，稳定地变成现实。