nmoe:NVIDIA B200向け Mixture-of-Experts トレーナー

AI/ML

概要

本リポジトリ「nmoe」は、NVIDIA Blackwell B200(アーキテクチャタグ:sm_100a)に最適化したMixture-of-Experts(MoE)トレーナーです。標準的なエキスパート経路で使われるNCCLベースのall-to-allを避け、RDEP(direct dispatch/return)という方式でエキスパートへの入力配布と戻り値を直接転送します。ノード内はCUDA IPC、ノード間はNVSHMEMを使うことで通信レイテンシと帯域のボトルネックを低減し、テンソル並列を使わずにB200の特性を活かす設計になっています。セットアップはハードウェアとドライバ要件が厳しく、B200専用のビルド/実行を前提としています。

GitHub

リポジトリの統計情報

  • スター数: 36
  • フォーク数: 4
  • ウォッチャー数: 36
  • コミット数: 6
  • ファイル数: 15
  • メインの言語: Python

主な特徴

  • B200(sm_100a)に特化したMoEトレーニング実装(B200-only)
  • RDEP(direct dispatch/return)でエキスパート経路の通信を最適化(CUDA IPC + NVSHMEM)
  • NCCLのall-to-allやテンソル並列を使わない設計でエキスパートパスのオーバーヘッドを削減
  • 小規模で意見志向(opinionated)なコードベース、実験的アプローチに適合

技術的なポイント

nmoe最大の技術的特徴は「エキスパート並列の実装手法」にあります。従来の大規模MoEでは、シーケンスごとに全GPU間でデータを交換するall-to-allが通信ボトルネックとなりやすく、さらにテンソル並列を併用すると複雑性と通信量が増します。本プロジェクトはこれを避けるためにRDEP(direct dispatch/return)を採用し、エキスパートへの入力配布と出力回収を直接的なGPU間通信で行います。ノード内ではCUDA IPCを使ってプロセス間でGPUメモリを直接参照し、コピーとオーバーヘッドを最小化します。ノード間通信はNVSHMEMを利用し、RDMAネイティブな共有メモリモデルで低レイテンシのデータ移動を行います。

このアプローチの利点は主に二つあります。第一に、all-to-allに伴う大規模な集合通信を排除することで通信待ち時間とメモリスパイクを抑えられる点。第二に、テンソル並列を導入しないため実装やデバッグが単純化され、B200のアーキテクチャ上で効率的に動作するようにチューニングしやすい点です。一方でトレードオフも存在します。CUDA IPCやNVSHMEMを前提にするためノード・ドライバ・OSの要件が厳しく、プラットフォームの汎用性は低下します。またRDEPはプロセス管理や同期の実装が難しく、スケーリング特性は実環境で検証が必要です。さらに「B200-only」という方針により、他GPUアーキテクチャへの移植には追加の開発工数が必要です。コードベースは小さく(ファイル数15)、実験的・研究目的での改造や検証に向いています。セットアップの前提条件(ドライバ、NVSHMEMのサポート、CUDA IPCの許可など)を満たせば、ブラックウェル世代のGPUでMoEの通信効率を大幅に改善できる可能性があります。

プロジェクトの構成

主要なファイルとディレクトリ:

  • .dockerignore: file
  • .gitignore: file
  • AGENTS.md: file
  • LICENSE: file
  • NOTICE: file

…他 10 ファイル

(リポジトリはコンパクトで、実行や実験に必要な最小限の実装が揃っています。READMEにセットアップ手順や依存要件、B200固有のビルドターゲットが記載されています。)

まとめ

B200に最適化した実験的MoEトレーナー。通信ボトルネックを避ける独自路線が魅力。

リポジトリ情報:

READMEの抜粋:

nmoe

   _ __   _ __ ___   ___   ___
  | '_ \ | '_ ` _ \ / _ \ / _ \
  | | | || | | | | | (_) |  __/
  |_| |_||_| |_| |_|\___/ \___|

No all-to-all. No tensor parallel. B200-only.

This repo is an opinionated Mixture-of-Experts trainer hard-targeted to NVIDIA Blackwell B200 (sm_100a). MoE expert parallelism is implemented via RDEP: direct dispatch/return using CUDA IPC (intra-node) and NVSHMEM (inter-node), instead of NCCL all-to-all collectives on the expert path.

Quick st…