train.py

"""GRPO-Zero training script.

Minimal version close to the original GRPO:Zero codebase.
For experiment-oriented training with hooks (LR scheduling, entropy
decay, timing, early stop, safety warnings), use train_v2.py instead.
"""
import os
from argparse import ArgumentParser
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import yaml

# RTX 3090 24GB 显存优化
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
from torch.utils.data import DataLoader
from torch.utils.tensorboard.writer import SummaryWriter

from countdown_task import CountdownTasksDataset, reward_function
from grpo import rollout, update_policy
from optimizer import MemoryEfficientAdamW
from qwen2_model import Transformer
from tokenizer import Tokenizer


def evaluate(model, tokenizer, device, dtype, config):
    val_dataset = CountdownTasksDataset(
        data_path=config["data"]["path"],
        tokenizer=tokenizer,
        split="val",
        test_size=config["data"]["test_size"],
        val_size=config["data"].get("val_size", 128),
        split_manifest_path=config["data"].get("split_manifest_path"),
    )
    generator = torch.Generator(device=device)
    dataloader = DataLoader(
        val_dataset,
        shuffle=False,
        collate_fn=CountdownTasksDataset.collate_fn,
        generator=generator,
        batch_size=config["training"]["batch_size"] // 2,
        drop_last=False,
    )
    success = []
    for batch in dataloader:
        episodes = rollout(
            model=model,
            tokenizer=tokenizer,
            batch=batch,
            max_gen_len=config["training"]["max_gen_len"],
            num_answer_per_question=1,
            reward_function=reward_function,
            device=device,
            dtype=dtype,
        )
        success.extend([episode.reward_info["answer_reward"] for episode in episodes])
    return np.mean(success)


def main(config_path: str):
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    pretrained_model_path = Path(config["model"]["pretrained_model_path"])
    device = torch.device(config["model"]["device"])
    dtype_map = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }
    dtype = dtype_map.get(config["model"]["dtype"], torch.bfloat16)
    torch.set_default_device(device)
    torch.random.manual_seed(config["training"]["random_seed"])
    BATCH_SIZE = config["training"]["batch_size"]
    NUM_QUESTIONS_PER_BATCH = config["training"]["num_questions_per_batch"]
    NUM_ANSWERS_PER_QUESTION = BATCH_SIZE // NUM_QUESTIONS_PER_BATCH

    current_time = datetime.now().strftime(r"%Y%m%d-%H%M%S")
    tb_writer = SummaryWriter(log_dir=f"{config['training']['log_dir']}/{current_time}")
    tokenizer = Tokenizer(str(pretrained_model_path / "tokenizer.json"))

    train_dataset = CountdownTasksDataset(
        data_path=config["data"]["path"],
        tokenizer=tokenizer,
        split="train",
        test_size=config["data"]["test_size"],
        val_size=config["data"].get("val_size", 128),
        split_manifest_path=config["data"].get("split_manifest_path"),
    )
    generator = torch.Generator(device=device)
    train_dataloader = DataLoader(
        train_dataset,
        shuffle=True,
        collate_fn=CountdownTasksDataset.collate_fn,
        generator=generator,
        batch_size=NUM_QUESTIONS_PER_BATCH,
    )

    model = Transformer.from_pretrained(pretrained_model_path, device=device).train()

    # KL penalty (optional)
    kl_coeff = config["training"].get("kl_coeff", 0.0)
    if kl_coeff > 0:
        ref_model = Transformer.from_pretrained(pretrained_model_path, device=device).eval()
        ref_model.requires_grad_(False)
        print(f"Loaded reference model for KL penalty (kl_coeff={kl_coeff})")
    else:
        ref_model = None

    optimizer = MemoryEfficientAdamW(
        model.parameters(),
        lr=config["training"]["learning_rate"],
        weight_decay=config["training"]["weight_decay"],
        betas=config["training"]["betas"],
        enabled=config["training"]["memory_efficient_adamw"],
    )

    ckpt_dir = Path(config["training"]["ckpt_dir"])
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    max_steps = config["training"].get("max_steps")

    train_iter = iter(train_dataloader)
    step = 0
    for batch in train_iter:
        step += 1

        episodes = rollout(
            model=model,
            tokenizer=tokenizer,
            batch=batch,
            max_gen_len=config["training"]["max_gen_len"],
            num_answer_per_question=NUM_ANSWERS_PER_QUESTION,
            reward_function=reward_function,
            device=device,
            dtype=dtype,
        )
        if config["training"]["skip_unfinished_episodes"]:
            episodes = [ep for ep in episodes if ep.is_finished]

        results = update_policy(
            model=model,
            optimizer=optimizer,
            episodes=episodes,
            micro_batch_size=config["training"]["micro_batch_size"],
            pad_token_id=tokenizer.pad_token_id,
            max_grad_norm=config["training"]["max_grad_norm"],
            device=device,
            dtype=dtype,
            entropy_coeff=config["training"].get("entropy_coeff", 0.0),
            kl_coeff=config["training"].get("kl_coeff", 0.0),
            ref_model=ref_model,
        )

        reward = [episode.reward for episode in episodes]
        answer_reward = [episode.reward_info["answer_reward"] for episode in episodes]
        formatted_reward = [episode.reward_info["format_reward"] for episode in episodes]

        print(
            f"\rStep {step}, mean_reward: {np.mean(reward):.2f}, "
            f"train success_rate: {np.mean(answer_reward):.2f}, "
            f"format_reward: {np.mean(formatted_reward):.2f}, "
            f"grad_norm: {results['grad_norm']:.2f}, "
            f"num_finished_episodes: {sum(ep.is_finished for ep in episodes)}, "
            f"entropy: {results['entropy']:.2f}, "
            f"kl: {results['kl']:.4f}"
        )

        if step % config["training"]["eval_interval"] == 0:
            eval_success_rate = evaluate(model, tokenizer, device, dtype, config)
            print(f"\rEval success rate: {eval_success_rate:.2f}")
            tb_writer.add_scalar("success_rate/eval", eval_success_rate, step)

        tb_writer.add_scalar("loss", results["loss"], step)
        tb_writer.add_scalar("mean_reward", np.mean(reward), step)
        tb_writer.add_scalar("success_rate/train", np.mean(answer_reward), step)
        tb_writer.add_scalar("format_reward", np.mean(formatted_reward), step)
        tb_writer.add_scalar("grad_norm", results["grad_norm"], step)
        tb_writer.add_scalar("entropy", results["entropy"], step)

        if step % config["training"]["ckpt_save_interval"] == 0:
            output_file = ckpt_dir / f"ckpt_{step:06d}.pt"
            torch.save(model.state_dict(), output_file)
            print(f"Saved checkpoint to {output_file}")

        if max_steps is not None and step >= max_steps:
            print(f"Reached max_steps={max_steps}, stopping.")
            break


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--config",
        type=str,
        default="CONFIG/grpo_full_qwen25_lr1e-5_legacy.yaml",
    )
    args = parser.parse_args()
    main(args.config)