Reanudar Entrenamiento

Este ejemplo entrena dos agentes pequeños para dos epochs, luego continúa la misma ejecución para dos epochs adicionales a partir de un punto de control. Ambas partes se registran en una única ejecución de W&B.

Configuración#

La tarea pide al segundo agente que escriba una respuesta de 2 a 3 veces más larga que la respuesta del primer agente.

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 misma función construye trainer para ambas partes de la ejecución, por lo que los dos trainers coinciden.

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

Entrenamiento con Puntos de Control#

Con checkpoint_dir, trainer guarda un punto de control después de cada 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)

El directorio ahora contiene checkpoint-1 y checkpoint-2.

Continuar la Ejecución#

El último punto de control guarda la ejecución de W&B ID. Añadirlo a wandb_config envía los nuevos registros a la misma ejecución.

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 cuenta toda la ejecución, por lo que 4 añade dos epochs a los dos que ya están completados. El mismo código también reanuda una ejecución que se detuvo en medio.

¿Por qué las Curvas se Alinean?#

Las métricas se registran en env_step, y cada punto de control guarda env_step y los estados aleatorios.
La ejecución reanudada comienza desde el punto de control guardado env_step, en la misma ejecución de W&B.
Después de una interrupción, el trabajo desde el último punto de control se vuelve a realizar. W&B omite los pasos que no son mayores que el último registrado, por lo que la curva no tiene duplicados.
Los modelos, los datos, la función de recompensa y otras configuraciones deben permanecer iguales. Solo algunas configuraciones, como num_train_epochs, los dispositivos e intervalos de registro, pueden cambiar.