Weights & Biaises sur Jean Zay

Avant-propos

Weights & Biaises (W&B ou wandb) fonctionne de pair avec son service cloud et nécessite la possession d'une licence (gratuite ou non). Vous devez donc posséder un compte sur leur site, activer une des licences proposées et envoyer vos résultats sur leur cloud afin de les visualiser.

Instance locale

La possibilité d'hégerber une instance W&B locale à l'IDRIS a été étudiée mais n'a pas été fructueuse.

Données sensibles - Usage alternatif de MLFLow

Si l'envoi de vos données sur une plateforme externe comme wandb.ai est rédhibitoire, vous pouvez vous orienter vers la solution MLFlow que nous recommandons.

MLflow n'a pas de contrainte particulière et une instance est hébergée sur la plateforme JupyterHub de l'IDRIS. C'est une solution idéale si vous souhaitez conserver à l’abri vos données sensibles puisque vos données restent sur le supercalculateur. Les logs peuvent ensuite être visualisés depuis le serveur MLFLow de JupyterHub.

Présentation et fonctionnalités

Weights & Biaises (W&B ou wandb) est une plateforme collaborative conçue pour optimiser et suivre les projets de machine learning et d’intelligence artificielle. Elle fournit des outils puissants pour le suivi des expériences, la gestion des hyperparamètres, l’analyse des résultats et la collaboration entre équipes.

Grâce à W&B, les chercheur·ses, ingénieur·es et data scientists peuvent :

suivre facilement leurs expériences : enregistrer chaque run, métrique, et paramètre d'entraînement pour mieux comprendre les performances des modèles ;
visualiser les résultats en temps réel : générer des tableaux de bord interactifs pour analyser les expériences et identifier les tendances ;
partager et collaborer : travailler en équipe sur des projets complexes grâce à des rapports partagés et des intégrations avec des outils comme GitHub ou Slack ;
garantir une reproductibilité de leurs résultats : garder un historique clair et organisé pour recréer vos résultats ou partager vos travaux avec d'autres.

W&B sur Jean Zay

Création d'un compte W&B

Comme indiqué en avant-propos, il vous faut un compte W&B pour pouvoir visualiser vos logs. Vous pouvez vous inscrire ici.

Lors de la création de votre compte vous devrez sélectionner une licence. Il en existe plusieurs qui procurent différents niveaux de service. Les licences gratuites sont :

la licence Free qui est relativement limitée : équipe de 5 personnes max et 5 Go de stockage ;
la licence Free W&B academic account qui offre plus de services : équipe de 10 personnes max, pas de limite du nombre d'équipes et 100 Go de stockage. Il faudra en faire la demande une fois votre compte créé.

statut académique

Dans tous les cas, n'oubliez pas de déclarer votre statut d'académique si vous en avez un.

Statut académique

Par la suite vous devrez soit créer une “équipe”, soit en rejoindre une déjà existante.

Sur la page de votre équipe, vous pourrez récupérer votre token d'accès qui vous permettra de vous connecter sur wandb.ai depuis Jean Zay.

Token wandb

Installation

W&B est disponible dans la plupart de nos modules PyTorch (s'il est absent n'hésitez pas à contacter le support IDRIS).

Si vous souhaitez utiliser une version spécifique X.Y.Z, vous pouvez l'installer localement sur votre compte :

pip install --no-cache-dir --user wandb==X.Y.Z

En cas d'installation locale, il vous faudra aussi exporter le chemin vers le fichier binaire dans votre PATH :

export PATH=$PATH:$HOME/.local/bin
#export PATH=$PATH:$PYTHONUSERBASE/bin

Tutoriel

Vous pouvez commencer par vous connecter sur W&B depuis Jean Zay avec votre token d'accès :

wandb login
#wandb: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)
#wandb: You can find your API key in your browser here: https://wandb.ai/authorize
#wandb: Paste an API key from your profile and hit enter, or press ctrl+c to quit: 9acc5adc39....
#wandb: Appending key for api.wandb.ai to your netrc file: $HOME/.netrc

Il existe des commandes CLI W&B pour gérer votre projet depuis le terminal :

Usage: wandb [OPTIONS] COMMAND [ARGS]...

Options:
  --version  Show the version and exit.
  --help     Show this message and exit.

Commands:
  agent         Run the W&B agent
  artifact      Commands for interacting with artifacts
  beta          Beta versions of wandb CLI commands.
  controller    Run the W&B local sweep controller
  disabled      Disable W&B.
  docker        Run your code in a docker container.
  docker-run    Wrap `docker run` and adds WANDB_API_KEY and WANDB_DOCKER...
  enabled       Enable W&B.
  init          Configure a directory with Weights & Biases
  job           Commands for managing and viewing W&B jobs
  launch        Launch or queue a W&B Job.
  launch-agent  Run a W&B launch agent.
  launch-sweep  Run a W&B launch sweep (Experimental).
  login         Login to Weights & Biases
  offline       Disable W&B sync
  online        Enable W&B sync
  pull          Pull files from Weights & Biases
  restore       Restore code, config and docker state for a run
  scheduler     Run a W&B launch sweep scheduler (Experimental)
  server        Commands for operating a local W&B server
  status        Show configuration settings
  sweep         Initialize a hyperparameter sweep.
  sync          Upload an offline training directory to W&B
  verify        Verify your local instance

Vous pouvez soit initialiser un projet avec la CLI ou le faire en python depuis votre code d'apprentissage.

Pour notre tutoriel, on va générer les logs de ce code exemple :

train.py
import wandb
import random

# start a new wandb run to track this script
wandb.init(
    # set the wandb project where this run will be logged
    project="my-jean-zay-training",

    # set experiment name
    name="Llama-training",
    # Add your jobid in your experiment name is a good pratice
    # name="my-jean-zay-training"+str(os.environ["SLURM_JOBID"])

    # track hyperparameters and run metadata
    config={
    "learning_rate": 0.02,
    "architecture": "LLama5",
    "dataset": "CommonCorpus",
    "epochs": 10,
    }

    # You can set W&B offline mode inside you training script !
    #mode="offline",
)

# simulate training
epochs = 10
offset = random.random() / 5
for epoch in range(2, epochs):
    acc = 1 - 2 ** -epoch - random.random() / epoch - offset
    loss = 2 ** -epoch + random.random() / epoch + offset

    # log metrics to wandb
    wandb.log({"acc": acc, "loss": loss})

# Mark the run as finished
wandb.finish()

Mode OFFLINE

Contrairement à TensorBoard ou encore MLFlow, W&B a été conçu avec l'hypothèse d'un accès à internet du serveur de calcul qui réalise l'apprentissage.

Il s'avère que les nœuds de calcul de Jean Zay n'ont pas d'accès à internet pour des raisons de sécurité mais aussi de bonnes pratiques (on ne veut pas avoir des ressources de calcul GPU/CPU inactives parce que le réseau internet est trop lent).

Il faut mettre W&B en mode OFFLINE. Si vous ne le faites pas, votre code va planter avec une erreur réseau après avoir consommé quelques heures de calcul inutilement.

# Depuis la CLI
wandb offline
# Ou avec une variable d'environnement
export WANDB_MODE=offline

Le mode OFFLINE va écrire vos logs d'apprentissage dans un dossier wandb/. Vous pouvez modifier ce comportement avec l'argument dir de wandb.init(). Par la suite, depuis la frontale, vous pourrez envoyer les logs et les metadata sur le serveur W&B.

Nous allons lancer notre code exemple sur un nœud de calcul avec un script Slurm :

run.slurm
#!/bin/bash
#SBATCH --job-name=Wandb
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --hint=nomultithread
#SBATCH --time=00:01:00
#SBATCH --output=log%j.out
#SBATCH --error=log%j.out

# on se place dans le repertoire de soumission
cd ${SLURM_SUBMIT_DIR}

# nettoyage des modules charges en interactif et herites par defaut
module purge

# chargement des modules
module load pytorch-gpu/py3/2.5.0

# activation du mode offline (necessaire sur jean-zay)
export WANDB_MODE=offline

# execution
srun python train.py

sbatch run.slurm
#Submitted batch job 123456

Avec ce code exemple, les logs ressemblent à cela :

Loading pytorch-gpu/py3/2.5.0
  Loading requirement: cuda/12.2.0 nccl/2.21.5-1-cuda cudnn/8.9.7.29-cuda
    gcc/10.1.0 openmpi/4.1.5-cuda intel-mkl/2020.4 magma/2.7.2-cuda sox/14.4.2
    hdf5/1.12.0-mpi-cuda libjpeg-turbo/2.1.3 ffmpeg/6.1.1 graphviz/2.49.0
    llvm/15.0.6
wandb: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.
wandb: Tracking run with wandb version 0.18.6
wandb: W&B syncing is set to `offline` in this directory.
wandb: Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
wandb:
wandb: You can sync this run to the cloud by running:
wandb: wandb sync $WORK/exemple_wandb/wandb/offline-run-20250127_160933-88q14mmo
wandb: Find logs at: wandb/offline-run-20250127_160933-88q14mmo/logs

Notez que nos logs indiquent clairement que nous avons utilisé le mode OFFLINE. Ils nous indiquent également la commande à réaliser pour faire l'envoi/la synchronisation avec wandb.ai.

# depuis mon dossier de projet
wandb sync $WORK/exemple_wandb/wandb/offline-run-20250127_160933-88q14mmo

Si vous avez exécuté plusieurs travaux, vous pouvez tout synchroniser en même temps :

wandb sync
#wandb: Number of runs to be synced: 2
#wandb:   wandb/offline-run-20250127_155830-02n127qe
#wandb:   wandb/offline-run-20250127_160933-88q14mmo
#wandb: NOTE: use wandb sync --sync-all to sync 2 unsynced runs from local directory.

Il ne vous reste plus qu'à profiter de la visualisation de vos logs sur votre espace de projet :

Preview wandb

Avant-propos​

Présentation et fonctionnalités​

W&B sur Jean Zay​

Création d'un compte W&B​

Installation​

Tutoriel​