Training fortsetzen

Dieses Beispiel trainiert zwei kleine Agenten für zwei epochs, setzt dann die gleiche Ausführung für zwei weitere epochs von einem Checkpoint fort. Beide Teile werden in einem W&B-Run protokolliert.

Einrichtung#

Die Aufgabe verlangt, dass der zweite Agent eine Antwort verfasst, die 2 bis 3 Mal so lang ist wie die Antwort des ersten Agenten.

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

Die gleiche Funktion erstellt trainer für beide Teile des Runs, sodass die beiden trainers übereinstimmen.

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

Training mit Checkpoints#

Mit checkpoint_dir speichert trainer einen Checkpoint nach jedem 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)

Das Verzeichnis enthält jetzt checkpoint-1 und checkpoint-2.

Fortsetzen der Ausführung#

Der neueste Checkpoint speichert den W&B-Run ID. Das Hinzufügen zu wandb_config sendet die neuen Protokolle an denselben Run.

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 zählt den gesamten Run, sodass 4 zwei epochs zu den beiden hinzufügt, die abgeschlossen sind. Der gleiche Code setzt auch einen Run fort, der mitten in der Ausführung angehalten wurde.

Warum die Kurven übereinstimmen#

Metriken werden bei env_step protokolliert, und jeder Checkpoint speichert env_step und die zufälligen Zustände.
Die fortgesetzte Ausführung beginnt von dem gespeicherten env_step, im selben W&B-Durchlauf.
Nach einer Unterbrechung werden die seit dem letzten Checkpoint durchgeführten Schritte erneut ausgeführt. W&B überspringt Schritte, die nicht größer sind als der zuletzt aufgezeichnete, sodass die Kurve keine Duplikate hat.
Die Modelle, Daten, die Belohnungsfunktion und andere Einstellungen müssen gleich bleiben. Nur einige Einstellungen, wie num_train_epochs, Geräte und Protokollintervalle, können sich ändern.