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

安装发布的 wheel（将 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 wheel。

启用可选内核（如果已安装）：

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

您的 GPU 具有 ≤24 GB VRAM，但需要提供 30B 参数 MoE 模型。
延迟很重要 – 缓存感知卸载将额外的数据传输开销保持在足够低的水平，以支持交互式聊天。
您想要即插即用的 HuggingFace 体验，同时仍能运行生产级 HTTP 服务器。
您在单台机器上运行多个 GPU，并希望库自动分片专家。

📚 进一步阅读

架构细节 – ARCHITECTURE.md
模型兼容性矩阵 – docs/model-compatibility.md
服务指南 – docs/serving.md
DFlash 投机解码 – docs/dflash.md
基准测试最佳实践 – docs/benchmarking.md

此摘要仅基于存储库的 README；有关最新信息，请参阅链接的文档文件。