恢复训练 | CoMLRL

此示例训练两个小型智能体，用于两个 epochs，然后从检查点继续相同的运行，用于两个 epochs。两个部分都记录到同一个 W&B 运行。

设置#

任务要求第二个智能体生成一个比第一个智能体响应 2 到 3 倍长的响应。

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

相同的函数为两个运行部分构建 trainer，因此两个 trainers 匹配。

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

使用检查点的训练#

使用 checkpoint_dir，trainer 在每个 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)

该目录现在包含 checkpoint-1 和 checkpoint-2。

继续运行#

最新的检查点存储 W&B 运行 ID。将其添加到 wandb_config 将新的日志发送到相同的运行。

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 统计整个运行，因此 4 将两个已完成的 epochs 添加到两个已完成的。相同的代码也恢复了在中间停止的运行。

为什么曲线对齐#

指标记录在 env_step，并且每个检查点保存 env_step 和随机状态。
恢复运行从保存的 env_step 开始，在相同的 W&B 运行中。
在中断后，从上次检查点开始的工作会被重新执行。 W&B 跳过步骤，如果步骤大小小于上次记录的步骤，因此曲线没有重复。
模型、数据、奖励函数和其他设置必须保持不变。只有少数设置，如 num_train_epochs、设备和日志间隔，可能发生变化。