Kog 在 @huggingface 上开源了他们用于展示模型每秒 3000+ token 速度的 2B 模型。非常酷的工作！https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model…

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

来源：https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model 返回文章页 (https://huggingface.co/blog)

• 摘要 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#tldr)
• 核心思路 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#the-idea)
• 研发历程 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#the-story)
  • 隐藏开销 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#hiding-overhead)
  • 架构设计 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#designing-the-architecture)
• 预训练与中期训练 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#pre-training-and-mid-training)
  • 数据混合与方案 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#data-mixtures-and-recipes)
  • 基础设施与性能 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#infrastructure-and-performance)
  • 训练过程 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#runs)
  • 评估追踪 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#evaluation-tracking)
• 后训练 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#post-training)
  • 完整训练总结 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#full-training-summary)
• 评估 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#evaluation)
  • 第二阶段验证 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#second-phase-validation)
  • 与同规模模型对比 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#comparison-with-similar-size-models)
• 推理速度 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#inference-speed)
• 为何开源我们的 2B 模型？ (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#why-open-source-our-2b-model)
  • 我们发布的内容 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#what-we-are-releasing)
• 局限性 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#limitations)
• 结论 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#conclusion)
• 致谢 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#acknowledgments)
• 深入探索的链接 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#explore-these-links-to-dig-deeper)
• 引用 (https://huggingface.co/blog/kogai/kog-laneformer-2b-the-latency-first-model#citation)

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

大多数 LLM 研究首先优化基准质量，而推理速度等指标通常被视为后续的部署问题：先训练模型，然后量化、分片、批处理输入、缓存输入，并编写更好的内核。

Kog 采用了不同的路线，将速度作为首要目标。当一个模型从零开始设计就旨在最大化解码速度时，会发生什么变化？哪些架构选择会被排除，哪些仍能保持强大的模型性能？

这篇博文讲述了 Kog 如何从头训练 Laneformer 2B，使其成为一个有能力的编码模型，同时尊重我们的 Kog 推理引擎所需的硬件约束，以及初创公司的预算限制。

关于 Kog Kog 是一家总部位于巴黎的 AI 基础设施初创公司，正在为 AI 智能体构建一个实时推理引擎，采用创新的底层 GPU 工程和 LLM 架构研究。更多背景信息，请参见 Kog 的网站和介绍性博文：

• 标准数据中心 GPU 上的实时 LLM 推理（每请求 3,000 个 token/s）(https://blog.kog.ai/real-time-llm-inference-on-standard-gpus-3-000-tokens-s-per-request/)

TL;DR

• Kog 设计了一种通道结构 Transformer 架构，用于在我们的推理栈上实现高速单请求解码。
• Kog 在小规模上验证了定制的架构更改，然后从头训练最终的 2.3B 模型，使用约 4T 预训练 token，接着使用约 2T 代码/推理密集型 token 继续训练，并使用约 210M token 进行指令微调。
• Kog 表明，即使资源有限，也可以构建和部署一个在编码基准测试中具有竞争力的自定义小型语言模型。Laneformer 2B 在贪婪解码中达到了 45.1% 的 HumanEval+ 和 51.6% 的 MBPP+。
• Kog 以 kogai/laneformer-2b-it (https://huggingface.co/kogai/laneformer-2b-it) 发布权重、Hugging Face 模型代码实现和文档。
• 您可以在我们的 playground (https://playground.kog.ai/) 体验加速版本。

核心思路

在低批量大小下，解码速度不仅仅是 FLOPs 问题 (https://blog.kog.ai/real-time-llm-inference-on-standard-gpus-3-000-tokens-s-per-request/)。大量时间花费在移动权重、同步内核以及逐层支付通信成本上。

在多 GPU 设置中，这种开销会进一步增加，因为引入了 GPU 间通信。在模型架构层面，张量并行 (TP) (https://danielvegamyhre.github.io/ml/performance/2025/03/30/illustrated-megatron.html) 是一种众所周知的跨 GPU 分配工作的方法，但每一层都迫使设备在进入下一层之前停止并交换结果。

这让我们想到了一个简单的问题：我们能否隐藏这些通信成本，而不是在每一层都支付它们？

简单粗暴的尝试解决这个问题可能会引入损害模型质量的临时架构更改，并且使得该方法难以应用于现有的预训练架构，而不会牺牲性能。

快速推理并不需要从头训练新模型，Kog 的推理引擎已经通过底层 GPU 优化在标准预训练架构上实现了非常高的解码速度。但要更进一步，运行时不能再被视为一个独立的部署层：模型架构本身必须暴露出正确的结构供引擎利用。

这些观察让我们得出一个结论：对于最快的单请求推理，架构和运行时应该一起设计。Laneformer 是我们为探索这种协同设计点而从头训练的第一个模型。

作为一家小型初创公司，我们目前还不能通过无限扩展来解决这个问题。目标必须有意识地设定：

设计并训练一个小型模型，具有强大的编码能力和极致的推理解码速度。

研发历程

隐藏开销

张量并行 (TP) 是有效的，因为它将大型矩阵运算分割到多个 GPU 上，但需要支付 GPU 间同步的代价。在单批量解码时，这个代价尤其痛苦。

一个显而易见的想法是延迟 TP 引入的通信。在实践中，简单地这样做会导致模型质量不佳：一旦隐藏状态不再在常规边界同步，模型质量就会急剧下降，需要寻找架构上的思路来保证训练稳定性和维持模型质量。

我们在这个阶段测试了小规模的各种变体。有趣的是，我们许多更复杂的想法要么降低了质量，要么让实现过于脆弱。有用的教训简单得令人尴尬：先尝试显而易见的事情！理解它为什么失败，然后只做最小的必要架构更改。

这条路径引出了我们现在称之为延迟张量并行 (DTP) (https://blog.kog.ai/delayed-tensor-parallelism-for-faster-transformer-inference) 的机制。完整的机制请参见我们的 DTP 深度解读。

架构设计

一旦 DTP 有了可行的形态，模型其余部分的设计必须保持保守。DTP 已经改变了标准张量并行 Transformer 的一个重要假设，因此我们不想同时引入不相关的架构创新。

还有一个速度预算。在理想情况下，当通信和其他开销被隐藏时，单批量吞吐量会受到 GPU 内存带宽的限制。每个生成的 token 都需要将模型权重和 KV 缓存从 GPU 内存流式传输到计算单元，因此每秒最大理论 token 数强烈取决于每次前向传播需要读取多少数据。

由于本次发布不针对非常长的上下文，主要的实际问题变成：模型可以有多大？

理论上，这个问题可以仅从硬件带宽得到答案。但实际上，另一个约束同样重要：我们实际能承受的预算和算力是多少，以及我们可以训练多长时间？

因此，最终大小是在三个约束的交集中选择的：

• 小到可以用我们的资源从头训练，
• 大到足以使编码基准测试和后训练有意义，
• 能够支持 DTP 和我们的推理引擎以达到尽可能高的速度。

2B 规模最终对我们来说是完美的。

下面是我们的架构卡片，展示了 DTP 和通道如何融入一个基本标准的解码器 Transformer。

一些说明：

• 最大的架构变化是新的 8 通道系统以支持 DTP。为了完美隐藏 TP 通信开销，我们将延迟设为 2。
• 我们使用因果分组查询注意力 (GQA)，有 32 个查询头和 16 个键/值头，这样可以将头均匀分布到我们的 8 个通道中。
• 我们在模型的 15 层中使用了滑动窗口注意力 (SWA) 10 层，以确保在上下文大小下流式传输 KV 缓存不会引入显著的延迟。

model-overview (https://cdn-uploads.huggingface.co/production/uploads/6746f7f283941c4e436b7538/VZZu5wK6xRAiIROIEz6ge.png)

Laneformer 2B 架构概览，展示了用于支持 DTP 的八通道结构。

预训练与中期训练

即使是 2B 规模的预训练，也是一座需要攀登的高山！幸运的是，如今有许多有用的资源，例如 Smol 训练手册 (https://huggingface.co/spaces/HuggingFaceTB/smol-training-playbook) 可以帮助你入门。

那么，首先让我们设计数据方案！

数据混合与方案

我们的目标是产生一个有能力的编码模型。由于我们无法访问专有数据，我们在开源约束下处理这个问题。

得益于当前的开源生态系统，特别是 NVIDIA 在 Nemotron 上的工作，我们能够收集超过 20TB 的高质量过滤开源数据，并避免了漫长且昂贵的数据混合搜索。

备注 从 Hugging Face Hub 获取 20TB 数据听起来简单，但在这个规模下工作本身就是一个系统问题。移动、过滤、验证和重新整理数据可能需要数天时间，因此数据管道中的小错误可能很快变得代价高昂。

遵循 NVIDIA Nemotron 论文 (https://arxiv.org/abs/2604.12374)，我们使用了一个多阶段方案作为起点：

• 从一个广泛的通用数据混合开始，
• 继续使用更聚焦的数据混合，
• 可选地以短的最后阶段结束，以改进特定能力。

也就是说，我们的 2B 模型是有意设计得小的。在这个规模下，它无法吸收可用数万亿 token 中的所有有用信号。因此，即使我们可以从 Nemotron 数据混合中汲取灵感，最终的混合必须更具选择性。对于 Laneformer 2B，我们引入了一个变化：我们将大部分专业化集中在第二阶段。

第一阶段将建立广泛的基础。第二阶段将模型重点转向编码和推理能力。

备注 这是一个明确的权衡。但是，阶段之间强烈的数据混合转变并非标准做法，我们预计这会损害一些通用能力。对于本次发布，我们接受了这个风险，期望换回一个强大的编码聚焦模型。

data-mixtures (https://cdn-uploads.huggingface.co/production/uploads/6746f7f283941c4e436b7538/NJm4u71rHZOmaZEmwXLpe.png)

第一阶段和第二阶段数据混合。

基础设施与性能

训练栈围绕可重复性和原始吞吐量构建：带有独立快照的启动脚本、检查点转换、Hugging Face 导出、验证、lm-eval 支持、CORE 追踪和仪表板都成为项目的一部分，以确保可靠的预训练运行和实时监控。

无需重复花费数十万美元进行一次无法事后修复的运行是多么可怕。为了避免任何潜在的训练停机，我们与合作伙伴确保可以通过备用机器有效交换节点，并将重启可靠性作为训练系统的一部分。

我们还花了相当多的时间精确调整我们的模型参数预算、架构细节、数据分片数量以及 TorchTitan 的分布式训练和编译参数，以在预训练期间达到约 17k tokens/s/GPU 的速度。

我们的预训练设置最终是：

• 24 个节点，每个节点 8 块 NVIDIA H100 GPU，共 192 块 H100 GPU。
• 通过 Scaleway 和 ADASTRA 实现的高效欧洲基础设施。
• 大约 21 天的训练时间。
• 混合 FP32/BF16 精度。
• 全分片数据并行 (FSDP)。
• AdamW 优化器，WSD 学习率调度。
• 最终步检查点选择。

训练过程

对于预训练和中期训练，我们使用 TorchTitan (https://github.com/pytorch/torchtitan) 作为我们的分布式训练栈。TorchTitan 为大规模 FSDP 训练提供了一个简单可靠的基础，同时保留了足够的灵活性以支持自定义的 Laneformer 架构。

Laneformer 在许多方面接近标准的仅解码器 Transformer，但其通道结构和面向 DTP 的布局意味着我们需要仔细控制模型实现、分布式配置和编译。

预训练阶段使用的序列长度为 4,096，全局批量大小为 1,536，即每个优化器步约 6.29M token。

losses (https://cdn-uploads.huggingface.co/production/uploads/6746f7f283941c4e436b7538/FZaEtgbkZGJaraQC_5vAE.png)

预训练（上图）和中期训练（下图）的损失曲线。

我们可以看到，进入第二阶段大大降低了 token 多样性，模型稳定在一个比第一阶段低得多的交叉熵上。