deepseek-ai/DeepEP
DeepEP: an efficient expert-parallel communication library
DeepEP – 高性能なMixture-of-Experts(MoE)モデル向け通信ライブラリ
概要 – DeepEP(DeepEveryParallel)は、大規模なMixture-of-Experts(MoE)トレーニングと推論に必要な超高速なGPU間通信プリミティブを提供するC++/Pythonライブラリです。エキスパート並列(EP)のオールツーオールカーネルに加え、パイプライン、コンテキスト、データ並列、リモートメモリアクセス(Engram)の実験的サポートを実装しています。カーネルはランタイム時にコンパニオンツールであるDeepJITを使用してJITコンパイルされ、使用中のGPUアーキテクチャに正確にチューニングされます。
重要性 – MoEモデルはニューラルネットワーク層を多数のエキスパートに分割し、各トークンをそのうちのいくつかにルーティングします。これによりモデル容量は大きくなりますが、各層でGPU間のテンソル移動が必要になります。従来のNCCLコレクティブはボトルネックになります。DeepEPのEPカーネルは高スループット、低レイテンシのオールツーオールディスパッチ/コンバインを実現し、FP8ディスパッチをサポートし、NVLink上でエキスパートウェイトのプリフェッチと冗長勾配の削減を行うことができ、MoonEPやUltraEPで説明されている「エキスパートレプリケーション」技術を可能にします。
主な特徴
- JITコンパイルされたEPカーネル(DeepJIT経由) – ランタイム時に正確なSM数とGPUモデル用にコンパイルされます。
- NCCL-Ginバックエンド – 既存のNCCLコミュニケータを再利用できる軽量なデバイスサイドAPI。
- 統一された
EPBufferAPI – ディスパッチ、コンバイン、ウェイトプリフェッチ、勾配削減、バッファ割り当てを処理する単一オブジェクト。 - 実験的なプリミティブ – Engram(RDMAベースのリモートメモリ)、パイプライン並列(PP)、コンテキスト/データ並列ワークロード用のBucketコレクティブ。
- NVIDIA Hopper+ GPUのサポートと、別のリポジトリを介したHuawei Ascend 950 NPUのサポート。
- 分析的なSM/QP数推定 – 自動チューニングは不要です。
- 冗長エキスパートサポート – 動的なエキスパートレプリケーション用の
lb_prefetch_weights/lb_reduce_grads。
典型的なワークフロー
EPBufferを作成し、MoEハイパーパラメータ(hidden size、top-k、エキスパート数、FP8フラグ)を指定します。バッファはオールツーオールカーネルに割り当てるSM数を推定します。- ディスパッチ –
buffer.dispatch(...)を呼び出して、トークンを全ランク間で適切なエキスパートにルーティングします。この呼び出しはEventOverlapを返し、後で待機できるため、計算と通信のオーバーラップが可能です。 - エキスパートGEMMを実行します。
- コンバイン –
buffer.combine(...)(またはバックワードヘルパー)を呼び出して、エキスパート出力を元のランクに還元します。 5.(オプション)負荷分散 –lb_prefetch_weightsとlb_reduce_gradsを使用して、負荷の高いエキスパートをNVLink接続されたGPU間でレプリケートします。
インストールスナップショット
# サブモジュール(DeepJIT)のクローンと初期化
git submodule update --init --recursive
# 対応NCCLパッケージのインストール(CUDA 13用の例)
python -m pip install "nvidia-nccl-cu13>=2.32.3" --no-deps
# ホイールのビルドとインストール
bash install.sh
このビルドは、ローカルのCUDA、NCCL、C++20ツールチェーンに対してC++拡張をコンパイルします。インストール時にGPUは必要ありません。カーネルは初回使用時にコンパイルされます。
DeepEPを使用するタイミング
- 大規模MoEモデル(例:数千のエキスパートを持つGPTスタイル)のトレーニングまたは推論。
- GPU間の負荷分散にエキスパートレプリケーションが必要なシナリオ。
- NVLinkノード内リンクと、オプションでRDMAノード間ネットワークを備えたシステム。
- PyTorch 2.xを使用しており、オールツーオール通信層のドロップイン置換を望むプロジェクト。
制限・実験的な部分
- Engram、PP、Bucketコレクティブはまだ実験的とマークされています。
- EPカーネルはGPU SMを必要とするため、純粋なRDMA専用モードはサポートされていません。
- 最近のハードウェア(NVIDIA Hopper以降)とCUDA 13.1+が必要です。
リソース
- メインリポジトリ: https://github.com/deepseek-ai/DeepEP
- ランタイムカーネルコンパイル用のDeepJITサブモジュール。
- Ascend互換フォーク: https://github.com/deepseek-ai/DeepEP-Ascend
- 関連する負荷分散プロジェクト: MoonshotAI/MoonEP、Dots-Infra/UltraEP。
DeepEPはエキスパート並列ディープラーニングワークロード向けの専門的な通信ライブラリであり、JITコンパイルされたカーネルとAPIを提供し、研究者やエンジニアが現代のGPUクラスタで効率的にMoEモデルをスケーリングできるようにします。
関連
- プロジェクト
- プロジェクト
- Dispatch
- プロジェクト
- プロジェクト