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