Comment exécuter des sweeps sur SLURM ?

Pour exécuter des sweeps W&B sur le système d’ordonnancement SLURM, lancez wandb agent --count 1 SWEEP_ID dans chaque job planifié. L’agent exécute un seul run d’entraînement puis s’arrête, ce qui simplifie l’estimation des ressources et permet à l’ordonnanceur SLURM de paralléliser les recherches d’hyperparamètres sur plusieurs jobs. Pour les jobs multi-GPU ou d’entraînement distribué, une règle supplémentaire s’applique : un seul processus par job SLURM doit appeler wandb.agent(). Si chaque rank appelle l’agent, chacun récupère sa propre configuration d’hyperparamètres auprès du contrôleur du sweep, et les ranks n’entraînent plus le même modèle. Les sections ci-dessous présentent la commande de soumission pour les jobs à processus unique, le modèle de gating par rank pour les jobs multi-GPU en entraînement distribué, les cas où le gate doit porter sur le rank global plutôt que sur le rank local, ainsi que la manière de journaliser depuis les ranks autres que le rank 0.

Jobs à processus unique ou à GPU unique

Soumettez chaque job planifié avec la commande suivante :

wandb agent --count 1 SWEEP_ID

Jobs multi-GPU avec entraînement distribué

Lorsque vous exécutez un sweep W&B avec un entraînement distribué (par exemple, des jobs multi-GPU avec --gpus-per-node), réservez l’appel à wandb.agent() à un seul rank à l’aide de la variable d’environnement SLURM_PROCID. Les autres ranks doivent exécuter l’entraînement distribué sans appeler l’agent de sweep ni wandb.init(). Pour en savoir plus sur la prise en charge de l’entraînement distribué côté W&B, voir Journaliser des expériences d’entraînement distribué.

import os
import wandb

def train():
    with wandb.init() as run:
        config = run.config
        # Diffuser la configuration aux autres ranks, puis lancer l’entraînement distribué.
        ...

if os.environ.get("SLURM_PROCID", "0") == "0":
    wandb.agent(sweep_id, function=train, count=1)
else:
    # Les ranks autres que le rank 0 reçoivent de celui-ci la configuration et lancent l’entraînement
    # sans appeler wandb.agent() ni wandb.init().
    run_worker()

Ce modèle garantit que :

Chaque job SLURM enregistre exactement un run auprès du contrôleur de sweep.
Les ranks autres que le rank 0 exécutent leur part de l’entraînement distribué sans créer de runs en double et déconnectés.
Le contrôleur de sweep suit correctement la progression et planifie de nouvelles configurations d’hyperparamètres.

Si vous utilisez submitit ou un lanceur similaire, appliquez la même vérification dans le point d’entrée de votre entraînement avant d’appeler wandb.agent().

Choisir entre `SLURM_PROCID` et `SLURM_LOCALID`

SLURM_PROCID correspond au rank global sur l’ensemble du job SLURM. Conditionnez l’exécution à SLURM_PROCID == 0 si vous voulez un seul agent pour l’ensemble du groupe distribué, ce qui est généralement le cas pour les sweeps. SLURM_LOCALID correspond au rank au sein d’un même nœud. Conditionnez l’exécution à SLURM_LOCALID == 0 si vous voulez un agent par nœud, par exemple pour lancer en parallèle des runs de sweep indépendants sur plusieurs nœuds.

Journaliser depuis les ranks autres que le rank 0

L’exemple ci-dessus suit le modèle Suivre un seul processus, dans lequel seul le rank 0 journalise dans W&B. Si vous souhaitez que les autres ranks journalisent également dans le même run, par exemple pour capturer les métriques système de chaque rank, utilisez le mode partagé et transmettez l’ID du run du rank 0 à chaque worker. Voir Suivre tous les processus dans un seul run.

Sweeps Experiments

​Jobs à processus unique ou à GPU unique

​Jobs multi-GPU avec entraînement distribué

​Choisir entre SLURM_PROCID et SLURM_LOCALID

​Journaliser depuis les ranks autres que le rank 0

Jobs à processus unique ou à GPU unique

Jobs multi-GPU avec entraînement distribué

Choisir entre `SLURM_PROCID` et `SLURM_LOCALID`

Journaliser depuis les ranks autres que le rank 0