Kog Laneformer 2B：Kog推理引擎背后的延迟优先模型

Kog 今日在 Hugging Face Hub 上发布了 Laneformer 2B 的权重和模型代码。这是一款拥有 23 亿参数的指令微调编码模型，专为高速解码设计。

大多数大语言模型研究优先考虑基准测试质量，而推理速度等指标通常被视为后续的部署问题：先训练模型，然后量化、分片、批量输入、缓存输入，并编写更好的内核。Kog 采取了不同的路线，将速度作为首要目标。当模型从零开始设计时就以最大化解码速度为目标时，会发生什么变化？哪些架构选择被排除，哪些仍能保持强大的模型性能？

这篇博文讲述了 Kog 如何从零开始训练 Laneformer 2B，使其成为一款有能力的编码模型，同时兼顾 Kog 推理引擎的硬件约束和初创公司的预算限制。

关于 Kog

Kog 是一家总部位于巴黎的 AI 基础设施初创公司，致力于通过创新的底层 GPU 工程和大语言模型架构研究，为 AI 代理构建实时推理引擎。

设计理念

在低批量大小下，解码速度不仅仅是 FLOP 问题。大量时间花费在移动权重、同步内核以及逐层支付通信成本上。在多 GPU 设置中，这种开销更加严重，因为引入了 GPU 间通信。在模型架构层面，张量并行（TP）是一种众所周知的在 GPU 间分割工作的方法，但每一层都迫使设备停止并在进入下一层之前交换结果。这引出了一个问题：我们能否隐藏这些通信成本，而不是在每一层支付它们？

试图解决这个问题的简单方法可能会引入损害模型质量的临时架构变更，并且使得该方法难以应用于现有的预训练架构而不牺牲性能。快速推理不需要从头训练新模型，但为了进一步优化，架构和运行时必须共同设计。Laneformer 是我们为了探索这种协同设计而从零开始训练的第一个模型。

隐藏开销：延迟张量并行

张量并行（TP）是有效的，但它引入的同步成本在批量大小为1的解码中尤其痛苦。Kog 提出了延迟张量并行（DTP），通过将通信延迟到多个层之后来隐藏开销。在测试了多种变体后，最有效的方案是简单的：先尝试显而易见的事情，并用最少的必要架构变更来修复失败。DTP 允许模型在多个层内保持隐藏状态不同步，从而减少通信频率，同时通过精心设计的架构调整保持模型质量。

架构设计

一旦 DTP 有了可行的形状，模型的其他部分保持保守。模型采用 8 车道结构以支持 DTP，每个车道处理不同的注意力头。使用了分组查询注意力（GQA），32 个查询头和 16 个键/值头，均匀分布在 8 个车道上。15 层中的 10 层使用了滑动窗口注意力（SWA）以加速 KV 缓存流。模型大小选择为 23 亿参数，这是一个在资源、性能和速度之间的最佳平衡点。

训练过程

训练分为三个主要阶段：

• 预训练：在约 4 万亿通用 token 上训练，采用标准数据混合。
• 中期训练：在约 2 万亿代码和推理密集型 token 上继续训练，数据混合强烈转向编码能力。
• 指令微调：在约 2.1 亿 token 上进行监督微调和偏好优化。

训练基础设施包括 24 个节点，每个节点 8 个 NVIDIA H100 GPU（共 192 个 GPU），使用 TorchTitan 进行分布式训练，实现了约 17k tokens/s/GPU 的吞吐量。训练历时约 21 天。

结果与发布

Laneformer 2B 在贪婪解码中达到了 45.1% 的 HumanEval+ 和 51.6% 的 MBPP+，在其规模范围内具有竞争力。模型权重、代码和文档已在 Hugging Face 上以 kogai-laneformer-2b-it 发布。Kog 还提供了一个通过其推理引擎体验加速版本的在线游乐场。

通过延迟张量并行和车道结构架构，Kog 展示了即使在小规模上，通过协同设计模型和推理引擎也能实现显著的解码速度提升，同时保持强大的编码性能。