KimiK3训练技术揭秘,896专家架构与LatentMoE如何平衡规模与效率
近日,AI研究员苏剑林在其文章《简单谈谈 K3 的 MoE 和 Attention》中详细拆解了Kimi K3在专家架构、注意力设计及训练优化方面的关键取舍。作为一款拥有2.8万亿参数的超大模型,K3...
近日,AI研究员苏剑林在其文章《简单谈谈 K3 的 MoE 和 Attention》中详细拆解了Kimi K3在专家架构、注意力设计及训练优化方面的关键取舍。作为一款拥有2.8万亿参数的超大模型,K3不仅在规模上实现了突破,更通过一系列技术创新,在计算效率和稳定性之间找到了新的平衡点。
在专家架构方面,K3采用了896个路由专家的设计,但并非每个Token都会调用全部专家。相反,模型只会激活其中16个专家,这种设计既保证了模型的表达能力,又有效控制了计算开销。苏剑林指出,随着专家数量的增加,模型需要解决两个核心问题:一是避免少数专家持续过载,二是降低跨设备通信带来的额外开销。为此,K3引入了LatentMoE技术,将Token的隐藏状态先压缩到更低维度的空间再进行专家计算,从而显著降低了计算和通信成本。
LatentMoE的核心在于重新分配专家预算。以K3为例,其主隐藏维度为7168,而路由专家则在3584维空间中计算,完成后再恢复到完整维度。这种降维策略使得模型可以在不增加单个专家计算量的情况下,将专家池从448个扩展到896个。更重要的是,这种设计让多个较窄的专家可以共同处理一个Token,从而实现更加细致的能力分工。例如,路由专家主要负责学习细分任务,而2个始终参与计算的共享专家则专注于处理通用的语言结构和基础语义。
除了专家架构的创新,K3还在注意力机制上进行了优化。不同于传统的单一注意力方案,K3采用了KDA(Key-Value Decomposition Attention)与Gated MLA(Mixture of Low-rank Attention)交替排列的设计。这种混合注意力机制不仅能够以更低的成本处理长上下文,还通过NoPE MLA重新分配了长上下文中的记忆和检索任务,进一步提升了模型的效率。
然而,模型规模的扩张也带来了新的挑战。当Token被发送到不同GPU上的专家时,跨设备通信会迅速增加;同时,更长的上下文会推高KV Cache和注意力计算的成本。为了解决这些问题,K3还引入了Quantile Balancing技术来协调近千个专家之间的负载,以及RMSNorm与SiTU-GLU来稳定专家分支的数值。这些配套机制共同确保了模型在扩大容量的同时,每一步计算都能得到有效控制。
总的来说,Kimi K3的技术取舍展示了一个超大规模模型在追求更高性能的同时,如何通过技术创新来平衡计算效率、资源消耗和稳定性。这一系列设计不仅为K3的成功奠定了基础,也为未来更大规模模型的开发提供了重要的参考经验。