nemo-mbridge-multi-node-slurm

マルチノード Slurm

単一ノードの uv run python -m torch.distributed.run コマンドを、Enroot コンテナーに対応したマルチノード Slurm sbatch スクリプトに変換し、よくあるマルチノードの障害をデバッグします。

最初の回答で確認する項目

Bridge のマルチノードジョブを変換またはデバッグする際は、次の順序で回答してください。

1. initialize.py に到達する Bridge スクリプトでは、srun ネイティブの起動形式を優先します。#SBATCH --ntasks-per-node=8 を指定し、srun ... uv run python <script> ... で直接起動してください。これらのジョブを python -m torch.distributed.run でラップしないでください。
2. Bridge は initialize.py の分散初期化中に、SLURM 変数から RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT を設定することを明記してください。
3. リポジトリ、データ、ログ、HF_HOME、UV_CACHE_DIR、NEMO_HOME には、共有パスと一致するコンテナーマウントが必要です。
4. NCCL タイムアウトの報告では、推測する前に、最初のログで次の点を確認してください。
   • 警告やフレームのノイズを除外して、実際のエラーを grep する
   • Failures: を調べ、最初に失敗したランクとノードを特定する
   • ncclUniqueId、timeout、または crash on rank 0 を grep する

2 つのアプローチ：srun ネイティブと uv run torch.distributed

アプローチ	ntasks-per-node	プロセスの起動方法	最適な用途
srun ネイティブ（推奨）	8	Slurm がノードごとに 8 タスクを起動	変換、推論、Bridge スクリプト
uv run torch.distributed（レガシー）	1	uv run python -m torch.distributed.run がノードごとに 8 プロセスを起動	MLM pretrain_gpt.py

srun ネイティブを優先してください。よりシンプルで、TRAIN_CMD のシェルエスケープの問題を回避できます。Megatron Bridge は、initialize.py の分散初期化中に呼び出される common_utils.py のヘルパーを通じて、SLURM 環境変数（SLURM_PROCID、SLURM_NTASKS、SLURM_LOCALID、SLURM_NODELIST）から RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT を自動的に導出します。そのため、これらを手動で設定する必要はありません。

クラスター環境

リポジトリ、データ、ログ、HF_HOME、UV_CACHE_DIR、NEMO_HOME には共有ファイルシステムを使用してください。マルチノード SFT/PEFT ジョブでは、ノード 0 で準備されたパック済みシーケンスデータを他のノードからも参照できるようにする必要があるため、NEMO_HOME にコンテナー内のデフォルト（/root/.cache/nemo）を使用しないでください。

認証情報を sbatch テンプレートやログに含めないでください。HF_TOKEN、GH_TOKEN、WANDB_API_KEY は、スケジューラー環境またはアクセス制限されたシークレットファイルを通じて指定し、スクリプト本文にトークン値を直接記述しないでください。コピーして使用できる環境設定と sbatch テンプレートについては、references/templates.md を参照してください。

ログディレクトリ

<SHARED_FS>/logs/<job_name>_<suffix>

srun ネイティブのアプローチ（推奨）

Slurm がすべてのプロセスを直接起動します。torch.distributed.run は使用せず、TRAIN_CMD のエスケープも不要です。

SBATCH ヘッダー

#SBATCH --job-name=<model>-<task>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=8          # Slurm spawns 8 tasks per node
#SBATCH --gpus-per-node=8
#SBATCH --time=00:30:00
#SBATCH --account=<YOUR_ACCOUNT>
#SBATCH --partition=batch
#SBATCH --output=<SHARED_FS>/logs/<job_name>_%j.log
#SBATCH --exclusive

ビルドと起動

2 段階の srun パターンを使用してください。まず単一プロセスで uv sync を実行して共有キャッシュを作成し、その後、マルチノードジョブ全体を起動します。コピーしてそのまま使える完全版は references/templates.md にあります。

srun ネイティブの要点

• フェーズ 1 では、単一ノード／単一プロセスで uv sync を一度実行し、Lustre 上の共有キャッシュにすべての wheel をビルドします
• フェーズ 2 の uv sync は高速なノーオペレーションです（すべてキャッシュ済み）。スリープによるガードなしで全ランクから安全に実行できます
• initialize.py と common_utils.py が、SLURM 環境変数から RANK、WORLD_SIZE、LOCAL_RANK、MASTER_ADDR、MASTER_PORT を自動設定します
• sbatch レベルでエクスポートされた HF_TOKEN、HF_HOME、UV_CACHE_DIR などの環境変数は、srun タスクに引き継がれます
• 参照先：examples/models/glm/glm_45v/slurm_sft.sh、examples/models/minimax/minimax_m2/slurm_conversion.sh

---

uv run torch.distributed のアプローチ（レガシー）

スクリプトが torch.distributed.run を必要とする場合（例：MLM pretrain_gpt.py）、または Bridge の initialize.py が呼び出し経路に含まれない場合に使用します。

1. SBATCH ヘッダーを追加する

#SBATCH --job-name=<model>-<framework>
#SBATCH --nodes=<NNODES>
#SBATCH --ntasks-per-node=1          # ALWAYS 1 — torchrun handles per-node spawning
#SBATCH --gpus-per-node=8
#SBATCH --time=00:30:00
#SBATCH --account=<YOUR_ACCOUNT>
#SBATCH --partition=batch
#SBATCH --output=<SHARED_FS>/logs/<job_name>_%j.log
#SBATCH --exclusive

重要：--ntasks-per-node=8 ではなく、--ntasks-per-node=1 を指定してください。uv run python -m torch.distributed.run --nproc_per_node=8 は、ノードごとに 8 プロセスを起動します。ntasks-per-node=8 を指定すると、EADDRINUSE によるポート衝突が発生します（ノードごとに 8 タスク × 8 プロセス = 64 プロセス）。

2. マルチノードに変換する

単一ノードの設定を次の内容に置き換えます。

uv run python -m torch.distributed.run --nproc_per_node=8 \
  <script> <args>

TRAIN_CMD 文字列内で、マルチノードの設定を使用します。

uv run python -m torch.distributed.run \
  --nproc_per_node=8 \
  --nnodes=\${SLURM_JOB_NUM_NODES} \
  --node_rank=\${SLURM_NODEID} \
  <script> <args>

MASTER_ADDR と MASTER_PORT は、initialize.py／common_utils.py によって SLURM 環境変数から自動的に導出されるため、設定する必要はありません。

3. TRAIN_CMD と2段階の srun でラップする

同じ2段階のパターンを使用します。まず単一プロセスの srun で uv キャッシュをウォームアップし、その後に全体を実行します。

コンテナ内でランタイム変数を設定します。ただし、トークン値を長い bash -c 文字列に埋め込まないでください。スケジューラー経由で認証情報をエクスポートするか、ジョブの開始前にアクセス制限されたシークレットファイルを読み込みます。HF_HOME、UV_CACHE_DIR、NEMO_HOME は共有ストレージ上に配置してください。

4. 起動する（2段階）

references/templates.md の2段階起動テンプレートを使用し、このレガシー方式では #SBATCH --ntasks-per-node=1 を維持してください。

5. （任意）損失抽出のフッターを追加する

echo "======================================"
echo "Done. Losses:"
echo "======================================"
grep -E "iteration\s+" "$LOGDIR/<prefix>_${SLURM_JOB_ID}.log" | grep -iE "lm loss|reduced_train_loss" | head -25

---

対話型 GPU 割り当て（salloc + srun）

その場でのテスト（推論、変換のデバッグ）では、必ず次の3つの手順に従ってください。

手順 1：ノードを割り当てる

salloc --account <YOUR_ACCOUNT> -N 1 \
  -J <YOUR_ACCOUNT>-debug \
  -p interactive --gpus-per-node=8 -t 240

手順 2：コンテナのシェルを起動する

srun --mpi=pmix --no-kill \
  --container-image $CONTAINER_IMAGE \
  --container-mounts $CONTAINER_MOUNTS \
  --account <YOUR_ACCOUNT> -N 1 \
  -J <YOUR_ACCOUNT>-debug \
  --no-container-mount-home --gpus-per-node=8 \
  -p interactive --pty bash

手順 3：コンテナ内で環境をセットアップする

export GH_TOKEN=<YOUR_GITHUB_TOKEN>
wandb login <YOUR_WANDB_KEY>
export HF_TOKEN=<YOUR_HF_TOKEN>
export HF_HOME=<SHARED_FS>/HF_HOME
export UV_CACHE_DIR="<SHARED_FS>/uv_cache"
export NEMO_HOME="<SHARED_FS>/cache/nemo"
uv sync

その後、uv run を使ってコマンドを実行します（同期済みの仮想環境を使用します）。

uv run python -m torch.distributed.run --nproc_per_node=8 \
  examples/conversion/hf_to_megatron_generate_text.py \
  --hf_model_path <org>/<model> --prompt "What is AI?" --max_new_tokens 50 --ep 8

対話型割り当てでの落とし穴：

エラー	原因	対処方法
Cannot find GPU specification	--gpus-per-node が指定されていない	salloc と srun の両方に --gpus-per-node=8 を必ず指定します
invalid partition specified: pool0	パーティション名が正しくない	対話型には interactive、sbatch には batch を使用します。確認するには sinfo --summarize を実行します
Invalid account or account/partition combination	アカウントでそのパーティションを利用できない	sacctmgr -nP show assoc where user=$USER format=account,partition で組み合わせを確認します
Unable to create step for job... Requested node configuration is not available	-w <node> が割り当てと競合している	-w フラグを削除します。HF キャッシュは共有ファイルシステム上にあり、どのノードからもアクセスできます
コンテナ内で uv: command not found が発生する	コンテナに uv がプリインストールされていない	uv がプリインストールされたコンテナを使用するか、pip install uv を実行します
uv または pip の実行中に No space left on device が発生する	コンテナの /root/.cache/ がいっぱいになっている	export UV_CACHE_DIR=<SHARED_FS>/uv_cache を設定して保存先を変更します
ModuleNotFoundError: No module named 'megatron.core.activations'	コンテナにプリインストールされた megatron-core が、ローカルの 3rdparty/Megatron-LM と競合している	ローカル版をインストールします：pip install -e 3rdparty/Megatron-LM --no-deps --no-build-isolation

---

マルチノードの障害をデバッグする

すばやく診断する

ログで次のパターンを順に確認します。

# 1. Find the actual error (filter noise)
grep -a 'Error\|OOM\|CUDA out of memory\|FAILED\|Killed' job.log \
  | grep -v 'UserWarning\|AllocatorConfig\|transformer_engine\|frame\|srun: error'

# 2. Check which rank crashed first
grep -a 'Failures:' -A 20 job.log | head -25

# 3. Check for NCCL timeout
grep -a 'ncclUniqueId\|timeout\|crash on rank 0' job.log | head -5

デバッグチェックリスト

マルチノードジョブが失敗した場合：

1. 終了コードを確認する：1 = Python エラー、9 = OOM による強制終了、143 = SIGTERM（タイムアウトまたは連鎖的な終了）
2. 最初の障害を特定する：最初にクラッシュしたタスク／ノードはどれかを確認します。他のタスクは連鎖的に SIGTERM（143）を受け取ります
3. 実際のエラーを grep する：UserWarnings と NCCL のフレームダンプを除外します
4. 特に rank 0 を確認する：保存／エクスポートのエラーの多くは rank 0 で発生します
5. EP のサイズを確認する：MoE モデルでは、余裕を持って GPU メモリに収まるよう、num_experts / EP を確認します
6. まず対話型で試す：sbatch キューよりもすばやく試行できるよう、salloc -N 2 -p interactive を使用します

dist.barrier() での NCCL タイムアウト — 「rank 0 でクラッシュ」

症状: ノード 2 以降のすべての rank で以下が表示されます:

[rank8] is setting up NCCL communicator and retrieving ncclUniqueId from [0]
... wait timeout after 600000ms
This may indicate a possible application crash on rank 0

根本原因 (この順に確認してください):

原因	確認方法	修正
save_artifacts が rank 0 でハングする	エラーは save_hf_weights → dist.barrier() にある	タイムアウトを増やす: init_process_group("nccl", timeout=timedelta(minutes=60))
カスタムモデルコード内の ImportError	grep ImportError job.log	save_artifacts で ImportError を捕捉する (以下を参照)
エクスポート中の rank 0 の OOM	grep 'OutOfMemory' job.log	EP またはノード数を増やす
ノード間のネットワーク問題	クロスノードの rank でのみエラーが発生する	sinfo を確認し、別のノードを試す

save_artifacts の問題: trust_remote_code=True の場合、rank 0 は save_artifacts() (トークナイザー、config、カスタムモデリングコードをダウンロード) を実行し、他のすべての rank は直接 dist.barrier() に進みます。save_artifacts が遅いかクラッシュすると、他の rank がタイムアウトします。

save_artifacts の ImportError の修正 (hf_pretrained/base.py):

# Change:
except OSError:
    pass
# To:
except (OSError, ImportError):
    pass

MoE モデルの OOM

症状: モデルのロード中またはフォワードパス中に torch.OutOfMemoryError: CUDA out of memory。

重要なポイント: TP はエキスパートのメモリを削減しません。エキスパートを GPU に分割するのは EP だけです。

サイジングの計算式:

experts_per_gpu = num_experts / EP
expert_memory_gb ≈ experts_per_gpu * expert_params * 2 / 1e9  (bf16)
total_per_gpu ≈ expert_memory_gb + attention_memory_gb + kv_cache_gb

MiniMax-M2 の例 (256 エキスパート、~230GB fp8 → ~460GB bf16):

構成	ノード	GPU	エキスパート/GPU	結果
TP=2, EP=4	1	8	64	OOM (エキスパートが多すぎる)
TP=2, EP=8	2	16	32	ラウンドトリップ (重みのみ) では動作、推論では OOM
TP=1, EP=16	2	16	16	推論で動作
TP=2, EP=32	8	64	8	トレーニングに余裕あり

経験則:

• ラウンドトリップ (重みのみ): GPU あたりより多くのエキスパートを使用できます (モデルパラメータ ~60GB まで OK)
• 推論 (フォワードパス + KV キャッシュ): 余裕が必要です (モデルパラメータ最大 ~40GB)
• トレーニング (アクティベーション + オプティマイザー): さらに多くの余裕が必要です (モデルパラメータ最大 ~30GB)

ModuleNotFoundError: No module named 'megatron.core.tensor_parallel'

原因: コンテナにプリインストールされた megatron-core がローカルの 3rdparty/Megatron-LM と競合しています。

修正: 実行前に uv sync を追加してください:

CMD="if [ \"\$SLURM_LOCALID\" -eq 0 ]; then uv sync; else sleep 10; fi && "
CMD="${CMD}uv run --no-sync python <script> <args>"

ラウンドトリップでの FP8 重みの不一致

症状: ラウンドトリップは完了しますが、すべてのエキスパート重みで ❌ が表示され、ValueError: Weight mismatch detected が発生します。

原因: 元の HF 重みは FP8 で、Megatron は BF16 で保存します。エクスポートされた重みは BF16 です。元の FP8 との比較は atol=1e-1 を超えます。

これは FP8 モデルでは想定内です。 変換は正しく、比較の許容誤差が FP8→BF16 の精度差に対して不十分なだけです。

srun で WORLD_SIZE が設定されない

症状: スクリプトが「must be launched with torchrun」で終了します。

原因: スクリプトは os.environ.get("WORLD_SIZE") をチェックしますが、これは torchrun は設定するものの srun は設定しません。

修正: SLURM_NTASKS もチェックしてください:

if os.environ.get("WORLD_SIZE") is None and os.environ.get("SLURM_NTASKS") is None:
    sys.exit(1)

Bridge の common_utils.py ヘルパー (initialize.py から呼び出される) は SLURM から環境変数を設定します:

if "RANK" not in os.environ:
    os.environ["RANK"] = str(get_rank_safe())          # uses SLURM_PROCID
if "WORLD_SIZE" not in os.environ:
    os.environ["WORLD_SIZE"] = str(get_world_size_safe())  # uses SLURM_NTASKS
if "MASTER_ADDR" not in os.environ:
    os.environ["MASTER_ADDR"] = get_master_addr_safe()     # parses SLURM_NODELIST
if "MASTER_PORT" not in os.environ:
    os.environ["MASTER_PORT"] = str(get_master_port_safe()) # derives from SLURM_JOB_ID

---

重要な注意点

1. uv sync のための2段階の srun：まず単一プロセスの srun を実行してキャッシュを温め、その後、マルチノードの srun をすべて実行します。共有ファイルシステムにすでにすべてがキャッシュされているため、2回目の uv sync はすぐに終了し、何も行いません。

2. --no-container-mount-home は srun のフラグであり、#SBATCH ディレクティブではありません。

3. TRAIN_CMD 内でのエスケープ：TRAIN_CMD は二重引用符で囲まれた文字列なので、実行時（sbatch 時ではなく）に展開する必要がある Slurm 変数の内側の $ はエスケープします：

   • \${SLURM_PROCID}、\${SLURM_JOB_NUM_NODES}、\${SLURM_NODEID}
   • $GH_TOKEN、$LOGDIR、$WORKDIR などのホスト側の変数は sbatch 時に展開されるため、エスケープは不要です。

4. Bridge の rm -rf nemo_experiments：古いチェックポイントが自動再開されないよう、トレーニングの前に追加します。

5. MLM では PYTHONPATH が必要です：pretrain_gpt.py スクリプトの場合は、TRAIN_CMD 内に追加します：

   PYTHONPATH=${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} \
   

6. ノード数の目安：GPU の合計数は NNODES * 8 です。次の条件を満たす必要があります：TP * PP * EP * DP >= total_GPUs。ここで、DP = total_GPUs / (TP * PP * EP) です。

7. マルチノード SFT では共有ファイルシステム上に NEMO_HOME を設定：デフォルトの nemo キャッシュ（/root/.cache/nemo） はコンテナー内に限られます。パック済みシーケンスを使うマルチノード SFT では、あるノードで .npy ファイルが作成されますが、 ほかのノードからは参照できません。パック済みデータを共有するには、export NEMO_HOME=<SHARED_FS>/cache/nemo を設定します。 これを設定しないと、ほかのノードのランクで TypeError: 'NoneType' object is not an iterator が発生します。

完全なテンプレートとコマンド本文

コピーして使える sbatch のひな形と、Bridge／MLM 固有の TRAIN_CMD 本文については、 references/templates.md を参照してください。