Table des matières

Weights & Biaises sur Jean Zay

Avant-propos

Weights & Biaises fonctionne de paire avec son service cloud wandb.ai.
Si le code en lui même est open-source, son usage nécessite une licence (gratuite ou non).

L'IDRIS a déjà contacté l'entreprise pour avoir la possibilité d'une licence pour la plateforme sur Jean Zay.
Le prix de la licence (hébergeur) proposée étant très (très) élevée, il n'y aura pas d'hébergement de solution / serveur chez nous. L'usage de wandb.ai est donc obligatoire.

L'IDRIS recommande l'usage de MLFlow.
MLflow n'a pas de contrainte particulière et une instance peut être hébergée sur Jean Zay.
C'est une solution idéale si vous souhaitez conserver à l’abri vos données sensibles (vos données restent sur le supercalculateur). Vous pouvez lancer une instance pour visualiser vos logs depuis JupyterHub.

Présentation et Fonctionnalités

Weights & Biases (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 chercheurs, ingénieurs et data scientists peuvent :

  • Suivre facilement leurs expériences : enregistrez chaque run, métrique, et paramètre d'entraînement pour mieux comprendre les performances de vos modèles.
  • Visualiser les résultats en temps réel : générez des tableaux de bord interactifs pour analyser vos expériences et identifier les tendances.
  • Partager et collaborer : travaillez en équipe sur des projets complexes grâce à des rapports partagés et des intégrations avec des outils comme GitHub ou Slack.
  • Reproductibilité garantie : gardez un historique clair et organisé pour recréer vos résultats ou partager vos travaux avec d'autres.

Avec ses nombreuses intégrations (TensorFlow, PyTorch, Hugging Face, etc.) et son interface intuitive, Weights & Biases est un allié incontournable pour accélérer le développement de projets d’apprentissage machine et favoriser l’innovation.

Wandb sur Jean Zay

Création d'un compte W&B

Comme indiqué en avant-propos, il vous faut un compte W&B pour pouvoir faire la visualisation de vos logs. Vous pouvez vous incrire ici.

Lors de la création de votre compte vous devrez sélectionner une licence. Il en existe des différentes qui procurent différents niveaux de service. Celles qui nous intéressent (celles qui sont gratuites parce que la recherche n'a pas d'argent…) :

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

Dans tous les cas, n'oubliez pas d'indiquer votre statut d'académique (sauf si ce n'est pas le cas évidemment).

Par la suite vous devrez soit crée une “équipe” ou 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.

Installation

Wandb est disponible dans la plupart de nos modules Pytorch (si il est absent n'hésitez pas à contacter le support).

Si le package est absent ou que vous souhaitez utiliser une version spécifique, vous pouvez l'installer localement sur votre compte : 

pip install --no-cache-dir --user wandb

En cas d'installation local, il vous faudra aussi export le chemin du /bin dans votre $PATH : 

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

Utilisation

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 wandb CLI 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 faire les logs de ce toy problem:

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()

Attention : 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 noeuds de calcul de Jean Zay non 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 inactivent 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 pour une erreur réseau après vous avoir consommé quelques heures de calcul. 

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

Le mode offline, va venir écrire dans un dossier wandb/ vos logs d'apprentissage. Vous pouvez modifier ce comportement avec l'arguments 'dir' de wandb.init(). Par la suite, depuis la frontale, on enverra les logs et métadata sur le serveur W&B.

Nous allons lancer notre code exemple sur un noeud 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
#SBATCH -A 
 
# 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
export WANDB_MODE=offline
 
# execution
srun python train.py
sbatch run.slurm
#Submitted batch job 49185

Avec ce code exemple, les logs ressemblent à cela :

log49258.out
log49258.out 
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. Ceux nous indiquent même la commande à réaliser (sympa) pour faire l'envoie / la synchronisation avec wandb.ai. andb.ai) 

# depuis mon dossier de projet
wandb sync 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 :