Kog Laneformer 2B：Kog 推論エンジンの背後にあるレイテンシ最優先モデル
Kog は、高速シングルリクエスト推論のために一から設計された23億パラメータの命令チューニング済みコーディングモデル「Laneformer 2B」をリリースしました。モデルアーキテクチャと推論エンジンを共同設計することで、Kog は Delayed Tensor Parallelism とレーン構造の Transformer を導入し、通信オーバーヘッドを隠蔽します。このモデルは競争力のあるコーディングベンチマーク（HumanEval+ 45.1%、MBPP+ 51.6%）を達成し、Hugging Face でオープンソース化されています。
Kog は本日、Hugging Face Hub で Laneformer 2B の重みとモデルコードを公開しました。これは23億パラメータの命令チューニング済みコーディングモデルで、高速デコード向けに設計されています。
ほとんどの大規模言語モデルの研究はベンチマーク品質を優先し、推論速度などの指標は後から対処されるサーバー問題として扱われます。しかし Kog は異なるアプローチを取り、速度を第一の目標としました。モデルをゼロからデコード速度最大化を念頭に設計すると、何が変わるのでしょうか？どのようなアーキテクチャ選択が排除され、どれが強力なモデル性能を維持するのでしょうか？
このブログ記事は、Kog が Laneformer 2B をゼロからトレーニングし、Kog 推論エンジンのハードウェア制約とスタートアップの予算制約の両方を尊重しながら、有能なコーディングモデルに仕上げた物語です。
Kog について
Kog はパリに拠点を置く AI インフラストラクチャスタートアップで、革新的な低レベル GPU エンジニアリングと LLM アーキテクチャ研究を通じて、AI エージェント向けのリアルタイム推論エンジンを構築しています。
設計思想
低バッチサイズでは、デコード速度は FLOP 問題だけではありません。重みの移動、カーネルの同期、レイヤーごとの通信コストに多くの時間が費やされます。マルチ GPU セットアップでは、GPU 間通信が導入されるため、このオーバーヘッドはさらに悪化します。モデルアーキテクチャレベルでは、テンソル並列化（TP）が GPU 間で作業を分割するよく知られた方法ですが、各レイヤーはデバイスを停止させ、次のレイヤーに進む前に結果を交換する必要があります。
これにより、単純な疑問が生じます：これらの通信コストを毎レイヤー支払うのではなく、隠すことができるでしょうか？
この問題を解決しようとする単純な試みは、モデル品質を損なうアドホックなアーキテクチャ変更を導入し、既存の事前トレーニング済みアーキテクチャに適用することを難しくします。高速推論には新しいモデルをゼロからトレーニングする必要はありませんが、さらなる最適化のためには、アーキテクチャとランタイムを共同設計する必要があります。Laneformer は、この共同設計点を探求するためにゼロからトレーニングした最初のモデルです。
オーバーヘッドの隠蔽：Delayed Tensor Parallelism
テンソル並列化（TP）は効果的ですが、バッチサイズ1のデコードではその同期コストが特に痛烈です。Kog は Delayed Tensor Parallelism（DTP）を提案し、通信を複数レイヤー後に遅延させることでオーバーヘッドを隠します。さまざまな変種をテストした後、最も効果的なアプローチは単純でした：明白なことを最初に試し、失敗した理由を理解し、必要最小限のアーキテクチャ変更で修正する。DTP により、モデルは複数レイヤーにわたって隠れ状態を同期させずに済み、通信頻度を減らしながら、注意深く設計されたアーキテクチャ調整によりモデル品質を維持します。
アーキテクチャ設計
DTP が実現可能な形になった後、モデルの他の部分は保守的に保たれました。モデルは DTP をサポートするために8レーン構造を採用し、各レーンが異なるアテンションヘッドを処理します。グループ化クエリアテンション（GQA）が使用され、32のクエリヘッドと16のキー/バリューヘッドが8レーンに均等に分散されています。15レイヤーのうち10レイヤーはスライディングウィンドウアテンション（SWA）を使用してKVキャッシュのストリーミングを高速化しています。モデルサイズは23億パラメータに設定され、これはリソース、性能、速度の間の最適なバランスです。
トレーニングプロセス
トレーニングは3つの主要フェーズに分かれています：
- 事前トレーニング：約4兆の汎用トークンでトレーニング。標準的なデータ混合を使用。
- 中期トレーニング：約2兆のコードおよび推論集中トークンで継続トレーニング。データ混合はコーディング能力に強くシフト。
- 命令チューニング：約2.1億トークンで教師ありファインチューニングと選好最適化を実施。
トレーニングインフラストラクチャは24ノード、各ノード8基のNVIDIA H100 GPU（合計192 GPU）で構成され、TorchTitan を使用した分散トレーニングにより、約17kトークン/秒/GPUのスループットを達成しました。トレーニング期間は約21日間です。
結果と公開
Laneformer 2B は greedy decoding で HumanEval+ 45.1%、MBPP+ 51.6% を達成し、同サイズのモデルの中で競争力のある結果を示しました。モデルの重み、コード、ドキュメントは Hugging Face で kogai-laneformer-2b-it として公開されています。Kog はまた、推論エンジンを介して高速化バージョンを体験できるオンラインプレイグラウンドを提供しています。
Delayed Tensor Parallelism とレーン構造アーキテクチャにより、Kog は小規模でもモデルと推論エンジンの共同設計によって顕著なデコード速度向上を達成し、同時に強力なコーディング性能を維持できることを示しました。