src/train_stage.py

import gc
import time
from pathlib import Path

import torch
import wandb
from peft import LoraConfig, PeftModel, TaskType
from transformers import AutoModelForCausalLM
from trl import GRPOConfig, GRPOTrainer

from src.evaluate import load_tokenizer
from src.logging_utils import JsonlLogCallback, aggregate_train_logs
from src.rewards import countdown_reward, format_reward, proximity_reward, valid_numbers_reward


def release_trainer(trainer):
    """Release the colocated vLLM engine before constructing another trainer."""
    generation = getattr(trainer, "vllm_generation", None)
    llm = getattr(generation, "llm", None)
    llm_engine = getattr(llm, "llm_engine", None)
    engine_core = getattr(llm_engine, "engine_core", None)
    if engine_core is not None and hasattr(engine_core, "shutdown"):
        engine_core.shutdown()
    if generation is not None:
        generation.llm = None
    trainer.vllm_generation = None


def train_stage(base_model_name, previous_adapter, dataset, config, stage_dir):
    stage_dir = Path(stage_dir)
    checkpoint_dir = stage_dir / "checkpoint"
    tokenizer = load_tokenizer(base_model_name)
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    peft_config = None
    if previous_adapter is None:
        model = base_model_name
        peft_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=config["lora_r"],
            lora_alpha=config["lora_alpha"],
            lora_dropout=config["lora_dropout"],
            target_modules="all-linear",
        )
    else:
        base = AutoModelForCausalLM.from_pretrained(base_model_name, dtype=dtype)
        model = PeftModel.from_pretrained(base, previous_adapter, is_trainable=True)

    args = GRPOConfig(
        output_dir=str(stage_dir / "trainer_output"),
        model_init_kwargs={"dtype": "bfloat16", "attn_implementation": "sdpa"},
        max_steps=config["steps_per_stage"],
        learning_rate=config["learning_rate"],
        lr_scheduler_type=config.get("lr_scheduler_type", "cosine"),
        warmup_ratio=config.get("warmup_ratio", 0.03),
        beta=config["beta"],
        temperature=config["temperature"],
        max_completion_length=config["max_completion_length"],
        num_generations=config["num_generations"],
        per_device_train_batch_size=config["per_device_train_batch_size"],
        gradient_accumulation_steps=config["gradient_accumulation_steps"],
        use_vllm=config.get("use_vllm", True),
        vllm_mode=config.get("vllm_mode", "colocate"),
        vllm_gpu_memory_utilization=config.get("vllm_gpu_memory_utilization", 0.55),
        vllm_enable_sleep_mode=config.get("vllm_enable_sleep_mode", False),
        vllm_importance_sampling_correction=config.get("vllm_importance_sampling_correction", True),
        vllm_max_model_length=config.get("vllm_max_model_length", 768),
        vllm_tensor_parallel_size=config.get("vllm_tensor_parallel_size", 1),
        gradient_checkpointing=True,
        bf16=torch.cuda.is_bf16_supported(),
        fp16=not torch.cuda.is_bf16_supported(),
        logging_steps=1,
        disable_tqdm=True,
        save_strategy="no",
        report_to="wandb",
        run_name=f"{stage_dir.parent.name}_{stage_dir.name}",
        remove_unused_columns=False,
        seed=config["seed"],
        data_seed=config["seed"],
    )
    log_path = stage_dir / "train_log.jsonl"
    trainer = GRPOTrainer(
        model=model,
        args=args,
        train_dataset=dataset,
        reward_funcs=[countdown_reward, format_reward, valid_numbers_reward, proximity_reward],
        processing_class=tokenizer,
        peft_config=peft_config,
        callbacks=[JsonlLogCallback(log_path)],
    )
    started = time.time()
    trainer.train()
    trainer.save_model(str(checkpoint_dir))
    tokenizer.save_pretrained(str(checkpoint_dir))
    metrics = aggregate_train_logs(log_path)
    metrics["wall_clock_seconds"] = time.time() - started
    if wandb.run is not None:
        wandb.finish()
    release_trainer(trainer)
    del trainer, model, tokenizer
    gc.collect()
    torch.cuda.empty_cache()
    return str(checkpoint_dir), metrics


def train_continuous(base_model_name, dataset, config, run_dir):
    run_dir = Path(run_dir)
    trainer_output = run_dir / "trainer_output"
    tokenizer = load_tokenizer(base_model_name)
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=config["lora_r"],
        lora_alpha=config["lora_alpha"],
        lora_dropout=config["lora_dropout"],
        target_modules="all-linear",
    )
    total_steps = config["num_stages"] * config["steps_per_stage"]
    args = GRPOConfig(
        output_dir=str(trainer_output),
        model_init_kwargs={"dtype": "bfloat16", "attn_implementation": "sdpa"},
        max_steps=total_steps,
        learning_rate=config["learning_rate"],
        lr_scheduler_type=config.get("lr_scheduler_type", "cosine"),
        warmup_ratio=config.get("warmup_ratio", 0.03),
        beta=config["beta"],
        temperature=config["temperature"],
        max_completion_length=config["max_completion_length"],
        num_generations=config["num_generations"],
        per_device_train_batch_size=config["per_device_train_batch_size"],
        gradient_accumulation_steps=config["gradient_accumulation_steps"],
        use_vllm=config.get("use_vllm", True),
        vllm_mode=config.get("vllm_mode", "colocate"),
        vllm_gpu_memory_utilization=config.get("vllm_gpu_memory_utilization", 0.55),
        vllm_enable_sleep_mode=config.get("vllm_enable_sleep_mode", False),
        vllm_importance_sampling_correction=config.get("vllm_importance_sampling_correction", True),
        vllm_max_model_length=config.get("vllm_max_model_length", 768),
        vllm_tensor_parallel_size=config.get("vllm_tensor_parallel_size", 1),
        gradient_checkpointing=True,
        bf16=torch.cuda.is_bf16_supported(),
        fp16=not torch.cuda.is_bf16_supported(),
        logging_steps=1,
        disable_tqdm=True,
        save_strategy="steps",
        save_steps=config["steps_per_stage"],
        save_total_limit=config["num_stages"],
        report_to="wandb",
        run_name=f"{run_dir.name}_continuous",
        remove_unused_columns=False,
        seed=config["seed"],
        data_seed=config["seed"],
    )
    log_path = run_dir / "train_log.jsonl"
    trainer = GRPOTrainer(
        model=base_model_name,
        args=args,
        train_dataset=dataset,
        reward_funcs=[countdown_reward, format_reward, valid_numbers_reward, proximity_reward],
        processing_class=tokenizer,
        peft_config=peft_config,
        callbacks=[JsonlLogCallback(log_path)],
    )
    started = time.time()
    trainer.train()
    wall_clock_seconds = time.time() - started
    tokenizer.save_pretrained(str(trainer_output))
    if wandb.run is not None:
        wandb.finish()
    checkpoints = [
        str(trainer_output / f"checkpoint-{step}")
        for step in range(config["steps_per_stage"], total_steps + 1, config["steps_per_stage"])
    ]
    missing = [path for path in checkpoints if not Path(path).exists()]
    if missing:
        raise FileNotFoundError(f"missing continuous checkpoints: {missing}")
    release_trainer(trainer)
    del trainer, tokenizer
    gc.collect()
    torch.cuda.empty_cache()
    return checkpoints, str(log_path), wall_clock_seconds