PyTorch torchtune - CoreWeave Docs

torchtune est une bibliothèque basée sur PyTorch qui simplifie la création, le fine-tuning et l’expérimentation de LLM. torchtune prend également en charge nativement la journalisation avec W&B, ce qui facilite le suivi et la visualisation des processus d’entraînement. Ce guide explique comment activer la journalisation W&B dans les recettes torchtune, configurer le logger de métriques WandBLogger, identifier les métriques que torchtune suit par défaut et enregistrer des points de contrôle de modèle dans W&B Artifacts.

Tableau de bord d’entraînement torchtune

Consultez l’article de blog W&B consacré au fine-tuning de Mistral 7B avec torchtune.

Activer la journalisation W&B

Vous pouvez activer la journalisation W&B de deux façons : redéfinir des arguments au lancement depuis la ligne de commande, ou modifier le fichier de configuration de la recette. Choisissez la méthode la mieux adaptée à votre flux de travail.

Ligne de commande
Recette

Redéfinissez les arguments de ligne de commande au lancement :

tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
  metric_logger._component_=torchtune.utils.metric_logging.WandBLogger \
  metric_logger.project="llama3_lora" \
  log_every_n_steps=5

Activez la journalisation W&B dans la configuration de la recette :

# dans llama3/8B_lora_single_device.yaml
metric_logger:
  _component_: torchtune.utils.metric_logging.WandBLogger
  project: llama3_lora
log_every_n_steps: 5

Utiliser le logger de métriques W&B

Activez la journalisation W&B dans le fichier de configuration de la recette en modifiant la section metric_logger. Remplacez la valeur de _component_ par la classe torchtune.utils.metric_logging.WandBLogger. Vous pouvez également transmettre un nom de projet (project) et log_every_n_steps pour personnaliser le comportement de la journalisation. Vous pouvez aussi transmettre tout autre argument kwargs, comme vous le feriez avec la méthode wandb.init(). Par exemple, si vous travaillez au sein d’une équipe, vous pouvez transmettre l’argument entity à la classe WandBLogger pour indiquer le nom de l’équipe.

Recipe
Ligne de commande

# dans llama3/8B_lora_single_device.yaml
metric_logger:
  _component_: torchtune.utils.metric_logging.WandBLogger
  project: llama3_lora
  entity: my_project
  job_type: lora_finetune_single_device
  group: my_awesome_experiments
log_every_n_steps: 5

tune run lora_finetune_single_device --config llama3/8B_lora_single_device \
  metric_logger._component_=torchtune.utils.metric_logging.WandBLogger \
  metric_logger.project="llama3_lora" \
  metric_logger.entity="my_project" \
  metric_logger.job_type="lora_finetune_single_device" \
  metric_logger.group="my_awesome_experiments" \
  log_every_n_steps=5

Données journalisées

Une fois la journalisation W&B activée, vous pouvez explorer le tableau de bord W&B pour consulter les métriques journalisées. Par défaut, W&B journalise tous les hyperparamètres du fichier de configuration ainsi que les redéfinitions spécifiées au lancement, ce qui vous permet de conserver la configuration de chaque run en regard de ses métriques. W&B enregistre la configuration résolue dans l’onglet Overview. W&B stocke également la configuration au format YAML dans l’onglet Files.

Métriques journalisées

Chaque recette possède sa propre boucle d’entraînement. Consultez chaque recette pour connaître les métriques qu’elle journalise, dont celles-ci par défaut :

Métrique	Description
`loss`	La perte du modèle.
`lr`	Le taux d’apprentissage.
`tokens_per_second`	Le nombre de jetons par seconde traités par le modèle.
`grad_norm`	La norme du gradient du modèle.
`global_step`	Correspond à l’étape en cours dans la boucle d’entraînement. Tient compte de l’accumulation de gradients. À chaque étape de l’optimiseur, le modèle est mis à jour, les gradients s’accumulent et le modèle est mis à jour une fois toutes les `gradient_accumulation_steps` étapes.

global_step ne correspond pas au nombre d’étapes d’entraînement. Il correspond à l’étape en cours dans la boucle d’entraînement et tient compte de l’accumulation de gradients. À chaque étape de l’optimiseur, global_step est incrémenté de 1. Par exemple, si le dataloader contient 10 lots, que le nombre d’étapes d’accumulation de gradients est de 2 et que vous lancez l’entraînement sur 3 époques, l’optimiseur effectue 15 étapes : global_step varie donc de 1 à 15.

La conception de torchtune vous permet d’ajouter des métriques personnalisées ou de modifier les métriques existantes. Pour ce faire, modifiez le fichier de recette correspondant. Par exemple, vous pouvez journaliser current_epoch sous forme de pourcentage du nombre total d’époques, comme suit :

# dans la fonction `train.py` du fichier de recette
self._metric_logger.log_dict(
    {"current_epoch": self.epochs * self.global_step / self._steps_per_epoch},
    step=self.global_step,
)

L’ensemble des métriques journalisées peut varier d’une version de torchtune à l’autre. Pour ajouter une métrique personnalisée, modifiez la recette et appelez la fonction self._metric_logger.* correspondante.

Enregistrer et charger des points de contrôle

Enregistrez les points de contrôle dans W&B Artifacts pour versionner les poids du modèle avec les métriques et la configuration de chaque run. Vous pourrez ainsi reproduire les résultats et comparer les versions du modèle par la suite. La bibliothèque torchtune prend en charge plusieurs formats de points de contrôle. Selon l’origine du modèle que vous utilisez, vous devez utiliser la classe de checkpointer appropriée. Pour enregistrer les points de contrôle du modèle dans W&B Artifacts, il est recommandé de redéfinir les fonctions save_checkpoint dans la recette correspondante. L’exemple suivant montre comment redéfinir la fonction save_checkpoint pour enregistrer les points de contrôle du modèle dans W&B Artifacts.

def save_checkpoint(self, epoch: int) -> None:
    ...
    ## Enregistrer le point de contrôle dans W&B.
    ## Le nom du fichier dépend de la classe Checkpointer.
    ## Voici un exemple pour le cas full_finetune.
    checkpoint_file = Path.joinpath(
        self._checkpointer._output_dir, f"torchtune_model_{epoch}"
    ).with_suffix(".pt")
    wandb_artifact = wandb.Artifact(
        name=f"torchtune_model_{epoch}",
        type="model",
        # description du point de contrôle du modèle
        description="Model checkpoint",
        # vous pouvez ajouter toutes les métadonnées souhaitées sous forme de dict
        metadata={
            utils.SEED_KEY: self.seed,
            utils.EPOCHS_KEY: self.epochs_run,
            utils.TOTAL_EPOCHS_KEY: self.total_epochs,
            utils.MAX_STEPS_KEY: self.max_steps_per_epoch,
        },
    )
    wandb_artifact.add_file(checkpoint_file)
    wandb.log_artifact(wandb_artifact)

​Activer la journalisation W&B

​Utiliser le logger de métriques W&B

​Données journalisées

​Métriques journalisées

​Enregistrer et charger des points de contrôle

Activer la journalisation W&B

Utiliser le logger de métriques W&B

Données journalisées

Métriques journalisées

Enregistrer et charger des points de contrôle