# train_grpo.py """ Main script for training a Llava-based model using the custom MyGRPOTrainer. This script handles: 1. Configuration loading. 2. Initialization of Weights & Biases (wandb) and Hugging Face Accelerate. 3. Loading the model and processor. 4. Preparing the training and evaluation datasets. 5. Setting up and running the GRPO trainer. """ import argparse import os from functools import partial from typing import Dict, Any import torch import wandb from accelerate import Accelerator from datasets import Dataset, load_dataset from peft import LoraConfig, get_peft_model, TaskType from transformers import AutoProcessor, AutoModelForCausalLM from trl import GRPOConfig from config.config_llm import CONFIG from data_utils.commom_util import collate_fn, define_task_data_func, collate_fn_woI from trainer.DyMETrainer_llm import DyMETrainer from reward_utils.checker import RewardCalculator, RewardCalculatorLocal from reward_utils.refiner import ContextRefiner, ContextRefinerLocal def print_trainable_parameters(model): """ Prints the number of trainable parameters in the model. """ trainable_params = 0 all_param = 0 for _, param in model.named_parameters(): all_param += param.numel() if param.requires_grad: trainable_params += param.numel() print( f"trainable params: {trainable_params} || all params: {all_param} || " f"trainable%: {100 * trainable_params / all_param:.2f}" ) def setup_accelerator_and_wandb(bf16) -> Accelerator: """ Initializes Weights & Biases and the Hugging Face Accelerator. Returns: Accelerator: The configured accelerator instance. """ wandb_key = os.environ.get("WANDB_API_KEY") if wandb_key: wandb.login(key=wandb_key) if bf16: accelerator = Accelerator(mixed_precision="bf16", log_with="wandb") else: accelerator = Accelerator(log_with="wandb") return accelerator def load_model_and_processor(model_config: Dict[str, Any], peft_config: Dict[str, Any]): """ Loads the base model, applies LoRA configuration, and loads its processor. Args: model_config (Dict[str, Any]): Configuration dictionary for the model. peft_config (Dict[str, Any]): Configuration dictionary for PEFT (LoRA). Returns: Tuple[PeftModel, PreTrainedProcessor]: The loaded PEFT model and processor. """ model_id = model_config['pretrained_model_path'] # Load base model base_model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=getattr(torch, model_config['torch_dtype']), attn_implementation='flash_attention_2' if model_config['use_flash_attention_2'] else 'sdpa', low_cpu_mem_usage=True, ) processor = AutoProcessor.from_pretrained(model_id, padding_side='left') processor._tokenizer.padding_side = "left" lora_config = peft_config model = get_peft_model(base_model, lora_config) print("LoRA model created:") print_trainable_parameters(model) return model, processor # ## --- LoRA modification End --- ## def prepare_datasets(task: str, dataset_config: Dict[str, Any]) -> (Dataset, Dataset): """ Prepares the training and evaluation datasets based on the specified task. """ data_func = define_task_data_func(task) train_data_list = data_func(json_path=dataset_config['train_dataset']) train_dataset = Dataset.from_list(train_data_list) if 'chart' in task: eval_dataset = load_dataset(dataset_config['eval_dataset'])['test'] else: eval_dataset = None return train_dataset, eval_dataset def main(): """ Main function to orchestrate the model training pipeline. """ parser = argparse.ArgumentParser(description="Train a model using GRPO with LoRA.") parser.add_argument( '--config', type=str, default='norm', help="config file to use: 'norm' or 'llavacot'..." ) args = parser.parse_args() config_select = args.config if config_select == 'norm': from config_llm import CONFIG # 1. Load Configurations model_config = CONFIG['model'] training_config = CONFIG['training'] rl_config = CONFIG['rl'] client_config = CONFIG['client'] dataset_config = CONFIG['dataset'] peft_config = LoraConfig( r=16, lora_alpha=64, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"], task_type="CAUSAL_LM", lora_dropout=0.05, ) task = training_config['task'] # 2. Setup Environment accelerator = setup_accelerator_and_wandb(bf16=training_config['dyme_args']['bf16']) device_id = accelerator.process_index # 3. Initialize Model and Processor # ## --- LoRA modification Start --- ## # Pass peft_config to the model loading function model, processor = load_model_and_processor(model_config, peft_config) # ## --- LoRA modification End --- ## # 4. Prepare Datasets train_dataset, eval_dataset = prepare_datasets(task, dataset_config) # 5. Initialize Reward Calculator checker = RewardCalculatorLocal(rl_config, client_config.copy(), gpu_id=device_id) refiner = ContextRefinerLocal(rl_config, client_config.copy(), gpu_id=device_id) # 6. Define Training Arguments training_args = GRPOConfig(**training_config['dyme_args']) collate_fn_with_processor = partial(collate_fn_woI, processor=processor) # 7. Initialize the Trainer # Trainer handles PeftModel automatically dyme_trainer = DyMETrainer( model=model, checker=checker, refiner=refiner, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=processor, processing_func=collate_fn_with_processor, task_name=task, end_flag=rl_config['end_flag'], ) # 8. Start Training dyme_trainer.train() # When saving, the Trainer automatically saves only the LoRA adapter weights output_dir = training_args.output_dir output_dir = os.path.join(output_dir, "final_checkpoint") dyme_trainer.save_model(output_dir) if accelerator.is_main_process: # Non-model files like the processor still need to be saved manually processor.save_pretrained(output_dir) print(f"LoRA adapters and processor saved to {output_dir}") if __name__ == "__main__": main()