DeepSeek、MoEの学習・推論向けオープンソースDeepEPライブラリを公開

 (github.com/deepseek-ai)
1 ポイント 投稿者 GN⁺ 2025-02-26 | まだコメントはありません。 | WhatsAppで共有
  • DeepEPは、現代のML学習・推論における**Expert Parallelism（EP）**に焦点を当てた高性能通信ライブラリで、MoE dispatch/combine向けのall-to-all GPUカーネルと、FP8などの低精度サポートを提供する
  • V2リリースではEPを全面的にリファクタリングし、V1よりはるかに少ないSMリソースで同等以上の性能を実現し、バックエンドをNVSHMEMからより軽量なNCCL Ginへ切り替えた
  • V3構成基準のテストでは、8K tokens per batch、hidden 7168、top 8 experts、FP8 dispatch、BF16 combineの条件で測定され、V2はV1比で最大1.3倍のピーク性能と最大4倍のSM削減を記録した
  • すべてのカーネルは軽量なJITモジュールとしてランタイムコンパイルされ、インストール時にCUDAコンパイルは不要で、V2では高スループット・低レイテンシAPIを単一のElasticBufferインターフェースへ統合している
  • Hopper SM90 GPU、Python 3.8+、CUDA 12.3+、PyTorch 2.10+、NCCL 2.30.4+、NVLink、ノード間RDMAネットワークが必要で、Engram・PP・CPは実験的機能である

DeepEPが提供する範囲

  • **DeepEP（DeepEveryParallel）**は、現代の機械学習の学習と推論のための高性能通信ライブラリである
  • 現在の中心機能は**Expert Parallelism（EP）**であり、MoE dispatchとcombineのための高スループット・低レイテンシなall-to-all GPUカーネルを提供する
  • FP8を含む低精度通信をサポートする
  • パイプライン並列化（PP）、コンテキスト並列化（CP）、リモートメモリアクセス（Engram）向けの実験的プリミティブも含む
  • すべてのカーネルは軽量な**JIT（Just-In-Time）**モジュールとしてランタイムコンパイルされ、インストール時にCUDAコンパイルは不要
  • 軽量設計でありながら、複数の構成においてハードウェア帯域幅の限界に匹敵、またはそれを超える性能を目指している

V2リリースの主な変更点

  • V2はExpert Parallelismを完全にリファクタリングしたリリースである
    • V1より数倍少ないSMリソースで極限の性能を達成するよう設計されている
    • より大きなscale-upおよびscale-outドメインをサポートする
    • バックエンドはNVSHMEMから、より軽量なNCCL Gin backendへ移行した
  • 新機能は以下の通り
    • 全面的なJITコンパイル
    • ヘッダオンリーで軽量なNCCL Gin backend
    • 既存のNCCL communicatorを再利用可能
    • EPv2で高スループットAPIと低レイテンシAPIを単一のElasticBufferインターフェースへ統合
    • 新しいGEMMレイアウトを提供
    • 最大EP2048までの、より大きなscale-upおよびscale-outドメインをサポート
    • SM数とQP数を解析的に計算し、自動チューニングが不要
    • ハイブリッドモードとダイレクトモードを引き続きサポート
    • V3相当のレガシー学習において、SM使用量を24個から4〜6個へ削減しつつ同等以上の性能を維持
    • RDMAベースの0 SM Engram
    • RDMAベースの0 SM PP
    • Copy Engineベースの0 SM CP

制約と進行中の機能

  • V2はV1よりバッファサイズ消費量が大きい
  • 0 SM RDMA低レイテンシEPは今後サポートされない
  • Engram、PP、CPは実験的機能である
  • 開発中の機能は以下の通り
    • GPUとCPUの物理メモリを混在マッピングする連続仮想アドレス空間であるElastic GPU & CPU buffers
      • 完全自動・透過的なEngramや不均衡EPを可能にする方向
    • EP replayを活用して負荷不均衡を処理し、中間バッファサイズを削減する作業
    • DPおよびTP向けall-gather更新とreduce-scatterの実装
  • NVSHMEMベースのV1ドキュメントはdocs/legacy.mdで提供される

性能測定結果

  • V3構成に合わせて、以下の条件でテストした
    • バッチあたり8K tokens
    • hidden dimension 7168
    • top 8 experts
    • FP8 dispatch
    • BF16 combine
  • 主な結果は以下の通り
    • SM90、CX7、EP 8 x 2: dispatch 90 GB/s RDMA、combine 81 GB/s RDMA、12 SM
    • SM90、CX7、EP 8 x 4: dispatch 61 GB/s RDMA、combine 61 GB/s RDMA、6 SM
    • SM100、CX7、EP 8 x 2: dispatch 90 GB/s RDMA、combine 91 GB/s RDMA、12 SM
    • SM100、EP 8: dispatch 726 GB/s NVLink、combine 740 GB/s NVLink、64 SM
    • SM100、EP 8: dispatch 643 GB/s NVLink、combine 675 GB/s NVLink、24 SM
  • 測定値は論理帯域幅であり、たとえばEP 8 x 2の90 GB/sにはlocal rank trafficが含まれる
  • V2はV1比で最大1.3倍のピーク性能を達成し、最大4倍のSM数を削減した
  • より大きなEP構成の結果は現時点では省略されており、ユーザー自身でベンチマークすることが推奨される
  • 内部検証では、規模が大きくなってもカーネルは引き続きハードウェア帯域幅を飽和させられると見込まれている
  • V1の性能データはdocs/legacy.mdにある

インストールと要件

  • 要件は以下の通り
    • Hopper SM90 GPU、またはSM90 PTX ISAをサポートするアーキテクチャ
    • Python 3.8以上
    • SM90 GPU向けCUDA 12.3以上
    • PyTorch 2.10以上
    • NCCL 2.30.4以上
    • ノード内通信向けNVLink
    • ノード間通信向けRDMAネットワーク
  • NCCLは、Python環境からDeepEPが自動検出できるよう、pipでのインストールが推奨される
pip install "nvidia-nccl-cu13>=2.30.4" --no-deps
  • レガシーメソッドのサポートのため、NVSHMEMにも依存しており、インストール手順はNVSHMEM Installation Guideを参照
  • 開発時のビルドとテスト実行例は以下の通り
python setup.py build
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so

python tests/elastic/test_ep.py
python tests/elastic/test_agrs.py
python tests/elastic/test_engram.py
python tests/elastic/test_pp.py
  • インストールは以下のコマンドで行う
python setup.py install
  • インストール後は、Pythonプロジェクトでdeep_epをimportして利用できる

ElasticBuffer中心のインターフェース

  • V2では、すべてのEP演算が単一の**ElasticBuffer**インターフェースの下に統合されている
    • 高スループットAPIと低レイテンシAPIを同一インターフェースで扱う
    • バッファはMoE設定を直接指定して初期化できる
    • 最適なSM数とQP数を解析的に計算する
  • バッファ初期化の例では、ElasticBuffer.get_buffer_size_hint()で必要サイズを計算し、既存バッファを再利用できるか確認する
  • 新しいバッファ作成時にはnum_max_tokens_per_rank、hidden、num_topk、use_fp8_dispatchなどを指定する
  • _buffer.get_theoretical_num_sms(num_experts, num_topk)で、通信カーネルに使用する理論上のSM数を取得する
  • dispatchおよびcombine呼び出しでnum_smsを直接指定すると、計算値を上書きできる

学習・prefill・decodingの使用パターン

  • 学習または推論prefillにおいて、MoE dispatchはトークンを全rank上の対応するexpertへルーティングする
    • BF16とFP8入力をサポートする
    • handleには後続のcombine呼び出しに必要なルーティングメタデータが格納される
    • handle.num_recv_tokens_per_expert_listはGEMMに必要なexpertごとのトークン数を提供する
  • MoE dispatchのbackward passは実際にはcombineで処理される
  • MoE combineはexpert出力を元のrankへreduceする
  • MoE combineのbackward passは実際にはdispatchで処理される
  • 通信と計算のオーバーラップは**EventOverlap**インターフェースで管理する
    • 通信進行中に独立した計算を実行できる
    • 結果を使う前にevent.current_stream_wait()でcompute streamを同期する
  • 推論decodingでも同じElasticBufferを利用する
    • gating decisionが変わらない場合はcached_handleでルーティングメタデータを再利用する
    • このパターンによりレイアウト再計算とCPU同期を回避できる

環境変数とビルド時固定値

  • 一般設定
    • EP_BUFFER_DEBUG: バッファ初期化、SM推定、バックエンドのデバッグ情報を出力
    • EP_SUPPRESS_NCCL_CHECK: NCCLバージョン不一致チェックを抑制
    • EP_AVOID_RECORD_STREAM: 出力テンソルのrecord_streamを回避
    • EP_NUM_TOPK_IDX_BITS: top-k indexエンコードのビット数を上書き
  • ネットワーク設定
    • EP_NIC_NAME: NICプロパティ取得に使うデフォルトNIC名。既定値はmlx5_0
    • EP_OVERRIDE_RDMA_SL: RDMA service level indexを上書き
    • EP_DISABLE_GIN: NCCL Gin backendを無効化
  • JIT設定
    • EP_JIT_CACHE_DIR: コンパイル済みカーネルのキャッシュディレクトリ。既定値は$HOME/.deep_ep
    • EP_JIT_NVCC_COMPILER: NVCCコンパイラのパス
    • EP_JIT_CPP_STANDARD: C++標準バージョン。既定値は20
    • EP_JIT_DUMP_PTX、EP_JIT_DUMP_SASS、EP_JIT_DUMP_ASM: PTX・SASS出力ダンプ関連設定
  • 一部の環境変数はpersistentとして動作する
    • ビルド時にキャプチャされ、インストール済みパッケージのデフォルト値として組み込まれる
    • import時に現在の環境変数で上書きしない限り、このデフォルト値が自動適用される
    • 対象変数はEP_JIT_CACHE_DIR、EP_JIT_PRINT_COMPILER_COMMAND、EP_NUM_TOPK_IDX_BITS、EP_NCCL_ROOT_DIRである
  • 追加の詳細はtest_ep.pyまたはPythonドキュメントを参照

ネットワーク構成の推奨事項

  • DeepEPはInfiniBandネットワーク上で完全にテストされている
  • 理論上はRDMA over Converged Ethernet、すなわちRoCEとも互換性がある
  • トラフィック分離
    • InfiniBandのVirtual Lanesでサポートされる
    • expert-parallel workloadと他のworkloadを別々のvirtual laneに分離することが推奨される
    • V2ではsl_idx引数またはEP_OVERRIDE_RDMA_SL環境変数でvirtual lane割り当てを制御できる
  • Adaptive routing
    • InfiniBandスイッチが複数経路へトラフィックを均等分散する高度なルーティング機能
    • 追加レイテンシが発生しても、あらゆるネットワーク負荷条件で有効化が推奨される
  • Congestion control
    • 最大帯域幅に悪影響があるため無効化する
    • 混雑が避けられない場合は、対象workloadを低優先度virtual laneへ割り当てることが推奨される
  • PCI atomic mode
    • ハードウェアが対応していれば、NICのPCI_ATOMIC_MODEを設定してRDMA atomic operation性能を改善することが推奨される
sudo mlxconfig -y -d mlx5_$i set PCI_ATOMIC_MODE=4

実験ブランチとコミュニティフォーク

  • 実験ブランチ
    • Zero-copy: PyTorchテンソルと通信バッファ間のコピーをなくし、通常カーネルのSM使用量を大幅に削減
    • Eager: 低レイテンシプロトコルにより、RDMA atomic OPが追加するextra RTT latencyを除去
    • Hybrid-EP: TMA instructionsを用いた新バックエンド実装で、最小SM使用量、より大きいNVLinkドメイン、single-batchのfine-grainedな通信・計算オーバーラップ、PCIeカーネル、NVFP4サポートを含む
    • AntGroup-Opt: AntGroup Network Platform Departmentによる最適化シリーズ
    • Mori-EP: MORIバックエンドに基づくROCm/AMD GPU低レイテンシモードのサポート
    • nvDev: Compute Fabric Transportなど最新CUDA機能を含むV2ベースのブランチ
  • コミュニティフォーク
    • uccl/uccl-ep: Nvidia・AMDなど異種GPUと、EFA・Broadcom・CX7などのNICでDeepEPを動作可能にするサポート
    • Infrawaves/DeepEP_ibrc_dual-ports_multiQP: IBRC transportにmulti-QPソリューションとdual-port NICサポートを追加
    • antgroup/DeepXTrace: slow rankを効率的かつ高精度に特定する診断アナライザ
    • ROCm/mori: Wide EP、KVCache transfer、Collectivesなど性能重視のAI workload向けAMD次世代通信ライブラリ

ライセンスと引用

  • DeepEP V2はNCCL Gin backend上に構築されている
  • リポジトリのコードはMIT Licenseで公開されている
  • 引用項目はDeepEP: an efficient expert-parallel communication libraryで、年は2025年とされている

関連記事

まだコメントはありません。

まだコメントはありません。