SLURM で sweep を実行するにはどうすればよいですか？

W&B の sweep を SLURM scheduling system で実行する場合、各スケジュールされたジョブで wandb agent --count 1 SWEEP_ID を実行します。エージェントは 1 つのトレーニング run を実行して終了するため、リソースの見積もりが簡単になり、SLURM スケジューラーがジョブ間でハイパーパラメーター探索を並列化できます。マルチ GPU または分散トレーニングのジョブの場合、追加のルールが適用されます。SLURM ジョブごとに 1 つのプロセスだけが wandb.agent() を呼び出す必要があります。すべての rank がエージェントを呼び出すと、各 rank が sweep コントローラーから独自のハイパーパラメーター設定をプルし、rank は同じモデルをトレーニングしなくなります。以下のセクションでは、シングルプロセスジョブのサブミットコマンド、分散トレーニングを伴うマルチ GPU ジョブの rank ゲーティングパターン、global rank と local rank のどちらでゲートするか、および non-rank-0 の rank からログする方法を示します。

シングルプロセスまたはシングル GPU のジョブ

スケジュールされた各ジョブは、次のコマンドで送信します。

wandb agent --count 1 SWEEP_ID

分散トレーニングを使用するマルチ GPU ジョブ

分散トレーニング (たとえば --gpus-per-node を使用するマルチ GPU ジョブ) で W&B sweep を実行する場合は、SLURM_PROCID 環境変数を使用して、wandb.agent() が 1 つの rank でのみ実行されるように制御してください。その他の rank では、sweep エージェントや wandb.init() を呼び出さずに分散トレーニングを実行します。分散トレーニングにおける W&B 側の設定全般については、分散トレーニングの実験をログするを参照してください。

import os
import wandb

def train():
    with wandb.init() as run:
        config = run.config
        # 設定を他の rank にブロードキャストしてから、分散トレーニングを実行します。
        ...

if os.environ.get("SLURM_PROCID", "0") == "0":
    wandb.agent(sweep_id, function=train, count=1)
else:
    # rank 0 以外の rank は rank 0 から設定を受け取り、
    # wandb.agent() や wandb.init() を呼び出さずにトレーニングを実行します。
    run_worker()

このパターンにより、次のことが保証されます。

各 SLURM ジョブは、sweep コントローラーに run をちょうど 1 つ登録します。
rank 0 以外の rank は、重複した、関連付けられていない run を作成せずに、分散トレーニングの担当部分を実行します。
sweep コントローラーは進捗を正しくトラッキングし、新しいハイパーパラメーター設定をスケジュールします。

submitit または同様のランチャーを使用する場合は、トレーニングのエントリポイントで wandb.agent() を呼び出す前に、同じチェックを適用してください。

`SLURM_PROCID` と `SLURM_LOCALID` の使い分け

SLURM_PROCID は、SLURM ジョブ全体での global rank です。分散グループ全体でエージェントを 1 つだけ実行する場合は、SLURM_PROCID == 0 を条件にします。これは sweep の一般的な構成です。 SLURM_LOCALID は、単一ノード内の rank です。ノードごとに 1 つのエージェントを実行する場合は、SLURM_LOCALID == 0 を条件にします。たとえば、複数のノードで独立した sweep の run を並列に起動する場合に使用します。

rank 0 以外の rank からログする

上記の例は単一のプロセスをトラッキングするパターンに従っており、rank 0 のみが W&B にログします。rank 0 以外の rank からも同じ run にログしたい場合 (たとえば rank ごとのシステムメトリクスを取得する場合) は、共有モードを使用し、rank 0 の run ID を各ワーカーに渡します。詳しくはすべてのプロセスを単一の run にトラッキングするを参照してください。

Sweeps Experiments

​シングルプロセスまたはシングル GPU のジョブ

​分散トレーニングを使用するマルチ GPU ジョブ

​SLURM_PROCID と SLURM_LOCALID の使い分け

​rank 0 以外の rank からログする

シングルプロセスまたはシングル GPU のジョブ

分散トレーニングを使用するマルチ GPU ジョブ

`SLURM_PROCID` と `SLURM_LOCALID` の使い分け

rank 0 以外の rank からログする