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。
  • 統一されたEPBuffer API – ディスパッチ、コンバイン、ウェイトプリフェッチ、勾配削減、バッファ割り当てを処理する単一オブジェクト。
  • 実験的なプリミティブ – Engram(RDMAベースのリモートメモリ)、パイプライン並列(PP)、コンテキスト/データ並列ワークロード用のBucketコレクティブ。
  • NVIDIA Hopper+ GPUのサポートと、別のリポジトリを介したHuawei Ascend 950 NPUのサポート。
  • 分析的なSM/QP数推定 – 自動チューニングは不要です。
  • 冗長エキスパートサポート – 動的なエキスパートレプリケーション用のlb_prefetch_weights / lb_reduce_grads。

典型的なワークフロー

  1. EPBufferを作成し、MoEハイパーパラメータ(hidden size、top-k、エキスパート数、FP8フラグ)を指定します。バッファはオールツーオールカーネルに割り当てるSM数を推定します。
  2. ディスパッチ – buffer.dispatch(...)を呼び出して、トークンを全ランク間で適切なエキスパートにルーティングします。この呼び出しはEventOverlapを返し、後で待機できるため、計算と通信のオーバーラップが可能です。
  3. エキスパートGEMMを実行します。
  4. コンバイン – 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+が必要です。

リソース


DeepEPはエキスパート並列ディープラーニングワークロード向けの専門的な通信ライブラリであり、JITコンパイルされたカーネルとAPIを提供し、研究者やエンジニアが現代のGPUクラスタで効率的にMoEモデルをスケーリングできるようにします。

関連

  • プロジェクト
  • プロジェクト
  • Dispatch
  • プロジェクト
  • プロジェクト