Skip to main content
SLURM scheduling system で W&B sweeps を実行する場合は、スケジュールされた各ジョブで wandb agent --count 1 SWEEP_ID を実行してください。エージェントは 1 つのトレーニング run を実行すると終了するため、リソースの見積もりが簡単になり、SLURM スケジューラーでジョブ間のハイパーパラメーター探索を並列化しやすくなります。 マルチ GPU または分散トレーニング ジョブでは、追加のルールが 1 つあります。各 SLURM ジョブで wandb.agent() を呼び出すプロセスは 1 つだけにしてください。すべての rank がエージェントを呼び出すと、各 rank が sweep コントローラーからそれぞれ別のハイパーパラメーター設定を取得してしまい、同じモデルをトレーニングしなくなります。 以下のセクションでは、単一プロセス ジョブのサブミッション コマンド、分散トレーニングを伴うマルチ GPU ジョブでの rank-gating パターン、グローバル rank とローカル rank のどちらでゲートするべきか、そして rank 0 以外からログする方法を説明します。

​
単一プロセスまたは単一 GPU ジョブ

スケジュールされた各ジョブで次を送信してください: 
wandb agent --count 1 SWEEP_ID

​
分散トレーニングでのマルチ GPU ジョブ

分散トレーニングで W&B sweep を実行する場合 (たとえば --gpus-per-node を使用するマルチ GPU ジョブ) 、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
        # 設定を他のランクにブロードキャストし、分散トレーニングを実行します。
        ...

if os.environ.get("SLURM_PROCID", "0") == "0":
    wandb.agent(sweep_id, function=train, count=1)
else:
    # ランク0以外のランクはランク0から設定を受け取り、
    # wandb.agent() や wandb.init() を呼び出さずにトレーニングを実行します。
    run_worker()
このパターンにより、次のことが保証されます。
  • 各 SLURM ジョブは、sweep コントローラーに対して 1 つの run だけを登録します。
  • rank 0 以外の rank は、重複した孤立 run を作成することなく、分散トレーニングの担当部分を実行します。
  • sweep コントローラーは進捗を正しくトラッキングし、新しいハイパーパラメーター設定をスケジュールします。
submitit または同様のランチャーを使用する場合も、wandb.agent() を呼び出す前に、トレーニングのエントリポイントで同じチェックを行ってください。

​
SLURM_PROCID と SLURM_LOCALID の使い分け

SLURM_PROCID は、SLURM ジョブ全体でのグローバルrankです。分散グループ全体に対してエージェントを 1 つだけ起動したい場合は、SLURM_PROCID == 0 を条件にします。これは、Sweeps で一般的なケースです。 SLURM_LOCALID は、単一ノード内でのrankです。ノードごとに 1 つのエージェントを起動したい場合は、SLURM_LOCALID == 0 を条件にします。たとえば、独立した sweep run をノード間で並列に起動する場合です。

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

上記の例は、rank 0 のみが W&B にログする 単一のプロセスをトラッキングする パターンに従っています。rank ごとのシステムメトリクスも取得したいなど、rank 0 以外の rank も同じ run にログする場合は、shared mode を使用し、rank 0 の Run ID を各ワーカーに渡してください。詳しくは、すべてのプロセスを単一の run にトラッキングする を参照してください。
Sweeps Experiments