EfficientMoE/MoE-Infinity

PyTorch library for cost-effective, fast and easy serving of MoE models.

📚 개요

MoE‑Infinity 는 전체 모델을 담기에 VRAM이 부족한 GPU에서 대규모 Mixture‑of‑Experts (MoE) 언어 모델을 실행할 수 있는 라이브러리입니다. 이는 GPU에 "핫" 전문가 가중치만 유지하고 나머지는 호스트 RAM 또는 SSD로 이동하여 적시에 가져오는 방식으로 작동합니다. 활성화 인식 캐시, 사전 가져오기, CUDA 그래프 트릭을 통해 추가 데이터 이동을 실시간 추론에 충분히 빠르게 유지합니다.


🚀 주요 기능

기능 사용자에게 의미
전문가 오프로딩 대규모 MoE 체크포인트(예: DeepSeek‑V2, Mixtral, Qwen3‑MoE)는 대부분의 전문가 파라미터를 CPU/SSD에 저장하여 16GB 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 양자화 전문가 오프로드.

🛠️ 설치 빠른 시작

  1. 가상 환경 생성 (Python 3.10+ 권장).
    conda create -n moe-infinity python=3.12
    conda activate moe-infinity
    
  2. 릴리스된 휠 설치 (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
    
  3. (선택 사항) 소스에서 빌드 – 사용자 정의 CUDA 컴퓨트 기능이 필요하거나 최신 코드를 원하는 경우. README의 "Install from Source" 섹션을 따르세요. 주로 CUDA, CUTLASS, 그리고 호환되는 PyTorch 휠이 필요합니다.
  4. 선택적 커널 활성화 (설치된 경우):
    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이 24GB 이하인 GPU를 보유하고 있지만 30B 파라미터 MoE 모델을 제공해야 하는 경우.
  • 지연 시간이 중요한 경우 – 캐시 인식 오프로딩은 추가 데이터 전송 오버헤드를 대화형 채팅에 충분히 낮게 유지합니다.
  • 드롭인 HuggingFace 경험을 원하면서도 프로덕션 등급 HTTP 서버를 실행할 수 있어야 하는 경우.
  • 단일 머신에서 여러 GPU를 실행하고 라이브러리가 전문가를 자동으로 샤딩하도록 하려는 경우.

📚 추가 자료

  • 아키텍처 세부 사항 – ARCHITECTURE.md
  • 모델 호환성 매트릭스 – docs/model-compatibility.md
  • 서빙 가이드 – docs/serving.md
  • DFlash 투기적 디코딩 – docs/dflash.md
  • 벤치마킹 모범 사례 – docs/benchmarking.md

이 요약은 리포지토리의 README에만 기반합니다. 최신 정보는 링크된 문서 파일을 참조하세요.

관련

  • 프로젝트
  • 프로젝트
  • Dispatch
  • 프로젝트
  • 프로젝트