再学習の開始

この例では、2つの小さなエージェントを2つの epochs で学習し、その後、チェックポイントから同じ実行をさらに2つの epochs で継続します。両方の部分のログは1つの W&B 実行に記録されます。

設定#

このタスクでは、2番目のエージェントに、最初のエージェントの応答よりも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 を構築するため、2つの 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 は完了した2つの epochs を追加します。同じコードは、途中で停止した実行を再開することもできます。

なぜ曲線が一致するのか#

指標は env_step に記録され、各チェックポイントで env_step とランダムな状態が保存されます。
再開された実行は、同じ W&B 実行内の env_step から開始されます。
中断後に、最後のチェックポイント以降の作業が再度実行されます。 W&B は、最後の記録されたものよりも小さいステップをスキップするため、曲線には重複はありません。
モデル、データ、報酬関数、およびその他の設定は同じである必要があります。 num_train_epochs、デバイス、およびログ間隔などのいくつかの設定のみを変更できます。