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