main_7B.py

# 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 (with PEFT/LoRA) and processor.
4. Preparing the training and evaluation datasets.
5. Setting up and running the GPRO trainer.
"""

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 transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from trl import GRPOConfig
from peft import LoraConfig, get_peft_model, TaskType  # NEW: Import PEFT modules

from config.config_7B import CONFIG
from data_utils.commom_util import collate_fn, define_task_data_func
from trainer.DyMETrainer_7B import DyMETrainer
from reward_utils.checker import RewardCalculator, RewardCalculatorLocal
from reward_utils.refiner import ContextRefiner, ContextRefinerLocal


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


# NEW: Updated function to accept peft_config
def load_model_and_processor(model_config: Dict[str, Any], peft_config: Dict[str, Any]):
    """
    Loads the pre-trained vision-language model, applies PEFT/LoRA, 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-enabled model and processor.
    """
    model_id = model_config['pretrained_model_path']

    model = Qwen2_5_VLForConditionalGeneration.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,
    )

    # Freeze the vision tower to save memory and computation
    model.model.visual.requires_grad_(False)

    # NEW: Create and apply LoRA configuration
    print("Applying LoRA configuration...")
    lora_config = peft_config

    model = get_peft_model(model, lora_config)

    # NEW: Print trainable parameters to verify LoRA is active
    # This should show a very small percentage of trainable parameters.
    model.print_trainable_parameters()

    processor = AutoProcessor.from_pretrained(model_id)
    processor.tokenizer.padding_side = "left"
    # image_token_id = processor.tokenizer.additional_special_tokens_ids[
    #     processor.tokenizer.additional_special_tokens.index("<image>")]

    return model, processor


def prepare_datasets(task: str, dataset_config: Dict[str, Any]) -> (Dataset, Dataset):
    """
    Prepares the training and evaluation datasets based on the specified task.

    Args:
        task (str): The name of the task (e.g., 'chartqa').
        dataset_config (Dict[str, Any]): Configuration for datasets.

    Returns:
        Tuple[Dataset, Dataset]: The training and evaluation datasets.
    """
    data_func = define_task_data_func(task)

    # Create training dataset
    train_data_list = data_func(json_path=dataset_config['train_dataset'])
    train_dataset = Dataset.from_list(train_data_list)

    # Create evaluation dataset
    if 'chart' in task:
        eval_dataset = load_dataset(dataset_config['eval_dataset'])['test']
        # Note: You can uncomment the line below for quick testing/debugging.
        # eval_dataset = eval_dataset.select(range(1000, 1100))
    else:
        # Extend this section for other tasks if needed in the future.
        raise NotImplementedError(f"Task '{task}' is not supported for evaluation in this script.")

    return train_dataset, eval_dataset


def main():
    """
    Main function to orchestrate the model training pipeline.
    """

    # 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=64,
        lora_alpha=128,
        lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        use_rslora=True,
        bias="none",
        task_type="CAUSAL_LM",
    )

    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
    # NEW: Pass peft_config to the loading function
    model, processor = load_model_and_processor(model_config, peft_config)

    # 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, processor=processor)
    # 7. Initialize the Trainer
    dyme_trainer = DyMETrainer(
        model=model,  # NEW: This is now a PeftModel
        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()

    output_dir = training_args.output_dir
    output_dir = os.path.join(output_dir, "final_checkpoint")

    # NEW: save_model will now save only the LoRA adapter weights
    dyme_trainer.save_model(output_dir)

    if accelerator.is_main_process:
        processor.save_pretrained(output_dir)
        # NEW: The saved model is just the adapter, not the full model.
        print(f"LoRA adapter and processor saved to {output_dir}")


if __name__ == "__main__":
    main()