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 展示了即使在小規模上，透過協同設計模型和推理引擎也能實現顯著的解碼速度提升，同時保持強大的編碼效能。