Reprendre l'entraînement

Cet exemple entraîne deux petits agents pour deux epochs, puis poursuit la même exécution pour deux autres epochs à partir d’un point de contrôle. Les deux parties sont enregistrées dans une seule exécution W&B.

Configuration#

La tâche demande au deuxième agent d’écrire une réponse de 2 à 3 fois plus longue que la réponse du premier agent.

import math
from pathlib import Path

from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer
from comlrl.utils import read_checkpoint_metadata

model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
train_dataset = Dataset.from_dict({
    "prompt": [
        "Describe a city in the clouds:",
        "Invent a new holiday and explain it:",
        "Write a bedtime story for a dragon:",
        "Design a playground on Mars:",
    ]
})


def length_ratio_reward(completions1, completions2):
    rewards = []
    for first, second in zip(completions1, completions2):
        ratio = len(second) / max(len(first), 1)
        distance = 0.0 if 2 <= ratio <= 3 else min(abs(ratio - 2), abs(ratio - 3))
        rewards.append(math.exp(-distance))
    return rewards

La même fonction construit le trainer pour les deux parties de l’exécution, de sorte que les deux trainers correspondent.

def build_trainer(num_train_epochs, wandb_config):
    config = MAGRPOConfig(
        num_turns=1,
        num_train_epochs=num_train_epochs,
        num_generations=4,
        max_new_tokens=64,
        logging_steps=1,
    )
    return MAGRPOTrainer(
        agents=[AutoModelForCausalLM.from_pretrained(model_name) for _ in range(2)],
        tokenizer=tokenizer,
        reward_func=length_ratio_reward,
        train_dataset=train_dataset,
        args=config,
        wandb_config=wandb_config,
    )

Entraînement avec des points de contrôle#

Avec checkpoint_dir, le trainer sauvegarde un point de contrôle après chaque epoch.

wandb_config = {"project": "comlrl-examples", "name": "resume-demo"}
checkpoint_dir = "./checkpoints/resume-demo"

trainer = build_trainer(num_train_epochs=2, wandb_config=wandb_config)
trainer.train(checkpoint_dir=checkpoint_dir, checkpoint_interval=1)

Le répertoire contient maintenant checkpoint-1 et checkpoint-2.

Continuer l’exécution#

Le dernier point de contrôle sauvegarde l’exécution W&B ID. L’ajout à wandb_config envoie les nouveaux journaux à la même exécution.

checkpoint = max(
    Path(checkpoint_dir).glob("checkpoint-*"),
    key=lambda path: int(path.name.rsplit("-", 1)[1]),
)
identity = read_checkpoint_metadata(checkpoint)["wandb"]
if identity:
    wandb_config = {**wandb_config, **identity, "resume": "allow"}

trainer = build_trainer(num_train_epochs=4, wandb_config=wandb_config)
trainer.train(
    resume_from_checkpoint=str(checkpoint),
    checkpoint_dir=checkpoint_dir,
    checkpoint_interval=1,
)

num_train_epochs compte toute l’exécution, de sorte que 4 ajoute deux epochs aux deux qui ont été terminés. Le même code reprend également une exécution qui s’est arrêtée au milieu.

Pourquoi les courbes sont alignées#

Les métriques sont enregistrées dans env_step, et chaque point de contrôle sauvegarde env_step et les états aléatoires.
La reprise démarre à partir du point de contrôle sauvegardé env_step, dans le même exécution W&B.
Après une interruption, le travail effectué depuis le dernier point de contrôle est refait. W&B ignore les étapes qui ne sont pas plus importantes que la dernière enregistrée, de sorte que la courbe n’a pas de doublons.
Les modèles, les données, la fonction de récompense et autres paramètres doivent rester inchangés. Seuls quelques paramètres, tels que num_train_epochs, les périphériques et les intervalles d’enregistrement, peuvent changer.