재학습 시작

이 예제는 두 개의 작은 에이전트를 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 와 같은 몇 가지 설정, 장치, 로깅 간격만 변경될 수 있습니다.