EfficientMoE/MoE-Infinity

📚 概要

MoE‑Infinity は、モデル全体を保持するのに十分なVRAMがないGPU上で、大規模なMixture‑of‑Experts（MoE）言語モデルを実行できるライブラリです。これは、GPU上に「ホットな」エキスパート重みのみを保持し、残りをホストRAMまたはSSDに移動して、ジャストインタイムでフェッチすることで実現します。アクティベーション認識キャッシュ、プリフェッチ、CUDAグラフのトリックにより、リアルタイム推論に十分な速度で追加のデータ移動を維持します。

🚀 主な機能

機能	ユーザーへのメリット
エキスパートオフロード	大規模なMoEチェックポイント（例：DeepSeek‑V2、Mixtral、Qwen3‑MoE）は、ほとんどのエキスパートパラメータをCPU/SSDに保存することで、16 GB GPUで提供できます。
アクティベーション認識キャッシュ	頻繁に使用されるエキスパートはGPU上に常駐します。システムは次に必要になるエキスパートを予測し、事前にロードします。
速度最適化	融合CUDAカーネル、CUDAグラフキャプチャ、Marlin INT4 GEMM、FP4/MXFP4パス、オプションのFlashAttention/FlashInferカーネルにより、レイテンシを低く抑えます。
HuggingFace互換API	`MoE`クラスは通常の`transformers`インターフェース（`generate`、`encode`など）を模倣します。推奨される方法は`MoE.serve()`で、連続バッチ処理用の非同期OpenAI互換HTTPサーバーを起動します。
本番環境対応のサービス提供	ストリーミング（SSE）、ページングKVキャッシュ、リクエストのプリエンプション、ホットリロード、ヘルスチェック、クラッシュリカバリログを備えたOpenAIスタイルのエンドポイント。
マルチGPUサポート	単一マシン上で、エキスパートは表示されているすべてのGPUにラウンドロビンで分散されます。各GPUには独自のキャッシュがあり、可能な場合はピアツーピアアクセスを使用できます。
実験的DFlash	ドラフトパスと検証パスを統合し、リッチなバッチ処理をサポートする投機的デコーディングフレームワーク。
オプションのアクセラレータ	FlashAttentionとFlashInferのプラグインサポート。DeepSeek‑V4‑Flash用のFP4量子化エキスパートオフロード。

🛠️ インストールのクイックスタート

仮想環境を作成（Python 3.10以上を推奨）。

conda create -n moe-infinity python=3.12
conda activate moe-infinity

リリース済みのホイールをインストール（v0.2.1を希望のタグに置き換えてください）。

TAG=v0.2.1
gh release download "$TAG" --repo EfficientMoE/moe-store --pattern 'moe_store-*.whl' -O /tmp/moe-store.whl
gh release download "$TAG" --repo EfficientMoE/MoE-Infinity --pattern 'moe_infinity-*.whl' -O /tmp/moe-infinity.whl
pip install /tmp/moe-store.whl /tmp/moe-infinity.whl

（オプション）ソースからビルド – カスタムCUDAコンピュートケーパビリティが必要な場合や、最新のコードが必要な場合は、READMEの「Install from Source」セクションに従ってください。主にCUDA、CUTLASS、および対応するPyTorchホイールが必要です。

オプションのカーネルを有効化（インストールされている場合）。

pip install -e '.[flash_attn]'   # FlashAttention
pip install -e '.[flashinfer]'   # FlashInfer

📦 使用方法

プロセス内（簡単な）例

from transformers import AutoTokenizer
from moe_infinity import MoE
import os, torch

ckpt = "deepseek-ai/DeepSeek-V2-Lite-Chat"
tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)

config = {
    "offload_path": os.path.expanduser("~/moe-infinity"),
    "device_memory_ratio": 0.75,   # キャッシュ用に予約するGPUメモリの割合
}
model = MoE(ckpt, config)

inp = tokenizer("translate English to German: How old are you?", return_tensors="pt").input_ids.cuda()
out_ids = model.generate(inp)          # 非推奨の同期パス。互換性のために維持
print(tokenizer.decode(out_ids[0], skip_special_tokens=True))

本番環境でのサービス提供（連続バッチ処理）

python -m moe_infinity.entrypoints.openai.api_server_v2 \
    --host 127.0.0.1 \
    --model deepseek-ai/DeepSeek-V2-Lite-Chat \
    --offload-dir ./offload_dir

サーバーはOpenAIスタイルのPOST /v1/completionsリクエストを受け付け、Server‑Sent Eventsを介して結果をストリーミングします。

📊 対応モデル（抜粋）

ファミリー	チェックポイントの例
DeepSeek‑V2 / V3	`deepseek-ai/DeepSeek-V2-Lite-Chat`
DeepSeek‑V4‑Flash (FP4)	`deepseek-ai/DeepSeek-V4-Flash`
Mixtral	`mistralai/Mixtral-8x7B-Instruct-v0.1`
Qwen3‑MoE	`Qwen/Qwen3-30B-A3B`
Qwen3.5‑MoE	`Qwen/Qwen3.5-35B-A3B`
GLM‑5.2 / 5.3 (FP8)	`zai‑org/GLM-5.2‑FP8`
GPT‑OSS	`openai/gpt-oss‑*`
DBRX	`databricks/dbrx-instruct`
Jamba	`ai21labs/Jamba‑*`
OLMoE	`allenai/OLMoE‑*`
Meta NLLB‑MoE	`facebook/nllb‑moe‑54b`

ライブラリは、インストールされているtransformersバージョンに対応するクラスが存在する場合にのみ、モデルを自動的に登録します。

🎯 MoE‑Infinityを使用すべき場合

VRAMが24 GB以下のGPUを持っているが、30BパラメータのMoEモデルを提供する必要がある場合。
レイテンシが重要 – キャッシュ認識オフロードにより、追加のデータ転送オーバーヘッドをインタラクティブなチャットに十分な低さに保ちます。
ドロップインのHuggingFaceエクスペリエンスを望みながら、本番環境対応のHTTPサーバーも実行したい場合。
単一マシンで複数のGPUを実行しており、ライブラリにエキスパートを自動的にシャーディングさせたい場合。

📚 さらに読む

アーキテクチャの詳細 – ARCHITECTURE.md
モデル互換性マトリックス – docs/model-compatibility.md
サービングガイド – docs/serving.md
DFlash投機的デコーディング – docs/dflash.md
ベンチマークのベストプラクティス – docs/benchmarking.md

この概要はリポジトリのREADMEのみに基づいています。最新の情報については、リンクされたドキュメントを参照してください。