mcore-run-on-slurm

SLURM で Megatron-LM を実行する

先に回答する定数

テキストのみの SLURM セットアップに関する質問には、完全なスクリプトを示す前に、以下の定数を回答してください。

• すべてのノードからアクセスできる共有ワークツリーのパスからジョブを送信し、スクリプト内でトレーニングを開始する前にそこへ cd します。
• ノードごとに srun タスクを 1 つ使用し、ワーカーの起動には torchrun を直接使わず、uv run python -m torch.distributed.run を使用します。
• MASTER_ADDR は scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1 から設定し、MASTER_PORT、NNODES=${SLURM_NNODES}、GPUS_PER_NODE=<GPUS_PER_NODE>、WORLD_SIZE=$((NNODES * GPUS_PER_NODE)) を設定します。
• torch.distributed.run に --nnodes、--nproc-per-node、--node-rank、--master-addr、--master-port を渡します。
• CUDA_DEVICE_MAX_CONNECTIONS：Blackwell より前の Hopper／Ampere で TP>1 または CP>1 を使用し、FSDP を使用しない場合は 1 にします。Blackwell／GB200 では不要です。Torch-FSDP2 または Megatron-FSDP では 1 を使用してはいけません。overlap_moe_expert_parallel_comm を使用する場合は 32 にします。

前提条件

• GPU パーティションへのジョブ送信権限がある SLURM クラスターのログイン環境。
• 割り当て内のすべてのノードからアクセスできるファイルシステム（NFS、Lustre など）にチェックアウトされた Megatron-LM。すべてのノードから、コード、データ、チェックポイント、出力用の同じパスにアクセスできる必要があります。
• uv がインストールされていること。ジョブの送信前に、ワークツリー上で一度 uv sync --extra training --extra dev（または --extra lts）を実行し、.venv を作成して、すべてのノードからアクセスできるようにします。

最小限の sbatch スクリプト

ワークツリーに run_megatron.slurm として保存します。

#!/bin/bash
#SBATCH --job-name=megatron
#SBATCH --account=<SLURM_ACCOUNT>
#SBATCH --partition=<SLURM_PARTITION>
#SBATCH --nodes=<NODES>
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=<GPUS_PER_NODE>
#SBATCH --time=<HH:MM:SS>
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

set -euo pipefail
cd <MEGATRON_WORKTREE>

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=${MASTER_PORT:-29500}
export NNODES=${SLURM_NNODES}
export GPUS_PER_NODE=<GPUS_PER_NODE>
export WORLD_SIZE=$((NNODES * GPUS_PER_NODE))

# Set CUDA_DEVICE_MAX_CONNECTIONS only when your configuration requires it
# (see the section below). Example for pre-Blackwell with TP>1 or CP>1
# (non-FSDP):
#   export CUDA_DEVICE_MAX_CONNECTIONS=1

srun --ntasks=${NNODES} --ntasks-per-node=1 bash -c '
  # NODE_RANK comes from SLURM_NODEID with one task per node.
  NODE_RANK=${SLURM_NODEID}
  uv run python -m torch.distributed.run \
    --nnodes='"${NNODES}"' \
    --nproc-per-node='"${GPUS_PER_NODE}"' \
    --node-rank=${NODE_RANK} \
    --master-addr='"${MASTER_ADDR}"' \
    --master-port='"${MASTER_PORT}"' \
    pretrain_gpt.py \
      <MEGATRON_ARGS>
'

ジョブを送信します。

mkdir -p logs && JOB_ID=$(sbatch --parsable run_megatron.slurm)
echo "Submitted ${JOB_ID}"

マルチノードのルール

• 実行するワークツリーからジョブを送信するか、スクリプト内でそこへ cd します。すべてのノードから共有ファイルシステム（NFS、Lustre など）上の同じパスにアクセスできる必要があります。ノードローカルのパスは、他のランクからは見えません。
• すべてのノードで 1 つの torchrun ワーカーグループを使用します。ノードごとに独立した単一ノードジョブを起動しないでください。
• --nproc-per-node は、ノードごとに認識される GPU の数と一致させます。
• チェックポイント、TensorBoard データ、構造化ログは共有ストレージに書き込みます。

CUDA_DEVICE_MAX_CONNECTIONS

適切な値は、ハードウェアと並列処理モードによって異なります。無条件にエクスポートしないでください。

• Blackwell より前（Hopper、Ampere）で TP>1 または CP>1 を使用し、FSDP を使用しない場合： 1 に設定します。このコードパスではこの値がアサートされます。1 でない場合、デッドロックが発生するのではなくアサーションエラーになります。
• Blackwell： 設定は不要です。設定しても効果はありません。
• Torch-FSDP2 または Megatron-FSDP： 1 にしてはいけません。環境変数を未設定のままにするか、1 より大きい値に設定します。
• overlap_moe_expert_parallel_comm を有効にする場合： 32 に設定します。

設定で必要な場合は、sbatch スクリプト内で明示的に設定してください。

コンテナー

多くのサイトでは、Megatron-LM をコンテナー内で実行します（クラスターによっては enroot／pyxis、または singularity を使用します）。その場合、uv が管理する .venv はコンテナー内からアクセスできるパスに配置し、コンテナーイメージにはリポジトリが想定するバージョンの CUDA／NCCL／torch を用意する必要があります（docker/.ngc_version.dev と .ngc_version.lts を参照）。上記のひな形はそのまま使用できます。srun の呼び出しを、スケジューラーのコンテナーフラグ（--container-image=…、--container-mounts=… など）で囲んでください。

モニタリングと結果の収集

squeue -j "$JOB_ID" -o "%.10i %.8T %.10M %.6D %R"
sacct -j "$JOB_ID" --format=JobID,State,ExitCode,Elapsed
scancel "$JOB_ID"

トレーニングスクリプトが結果の成果物（ランク 0 が出力する JSON メトリクスファイルや最終チェックポイントなど）を書き込む場合は、squeue の状態だけを待つのではなく、成果物をポーリングしてください。通常、SLURM がジョブの完了を示す前に出力が生成されます。成果物をポーリングすれば、割り当てがタイムアウトまで保持されるのを待たず、成果物ができた時点ですぐにジョブをキャンセルできます。

障害の診断

ランク 0 だけでなく、すべてのランクの stderr を確認してください。最初に発生した NCCL 以外の Python トレースバックが根本原因であることがほとんどです。他のランクで後から発生する NCCL タイムアウトは、最初のクラッシュに続いて起きる症状です。

すばやく分類します。

• OOM：調整する前に、ランク、フェーズ（順伝播／逆伝播／オプティマイザー）、バッチサイズ、シーケンス長、並列化（TP／DP／CP／PP）、ピークメモリを記録します。
• 形状／除算可能性エラー：WORLD_SIZE = TP × DP × CP × PP と、ヘッド数の除算可能性（num_attention_heads % TP == 0）を確認します。
• インポートエラー：作業ツリーの誤り、uv sync の実行漏れ、または古い PYTHONPATH が原因です。起動前に cd <MEGATRON_WORKTREE> を実行したことを確認します。
• Python トレースバックを伴わない NCCL エラー：割り当て、ポートへの到達性、MASTER_ADDR の名前解決、全ランク間でのコマンドの一貫性を確認します。

よくある落とし穴

• 初回のジョブ投入前に uv sync を実行し忘れること。仮想環境がない場合、各ジョブで srun の内部から仮想環境が再構築されるため、ジョブごとに数分かかります。
• ジョブ終了時に消えるノードローカルパスにログを出力すること。必ず共有ファイルシステムに書き込んでください。
• CUDA_DEVICE_MAX_CONNECTIONS=1 を何も考えずに設定すること。適切な値はハードウェアと並列化モードによって異なります（上記の専用セクションを参照）。FSDP で 1 に設定すると別の問題が発生します。Blackwell では効果がありません。Blackwell 以前の環境で TP>1 または CP>1（非 FSDP）の場合、コードはデッドロックせずにアサートします。
• uv run python -m torch.distributed.run の代わりに torchrun を直接実行すること。仮想環境の設定によっては、直接実行した torchrun は仮想環境のパッケージを認識しない Python インタープリター経由でディスパッチされる場合があります。