train_rag.py

# Copyright (c) Microsoft. All rights reserved.

"""Train a RAG agent using Agent-lightning.

Usage:
    python train_rag.py fast        # Fast training for CI/testing
    python train_rag.py single_gpu  # Optimized for Single GPU (1.5B/7B models)
"""

from __future__ import annotations

import argparse
import os
import uuid
from copy import deepcopy
from datetime import datetime
from typing import Any, Dict, List, Optional

import pandas as pd
from rag_agent import RAGAgent  # Make sure to import your RAGAgent class

import agentlightning as agl

# Base configuration (default configuration, can be overridden)
RL_TRAINING_CONFIG: Dict[str, Any] = {
    "algorithm": {
        "adv_estimator": "grpo",  # Use GRPO algorithm
        "use_kl_in_reward": False,
    },
    "data": {
        "train_batch_size": 16,  # Default configuration for multi-GPU
        "max_prompt_length": 8192,
        "max_response_length": 2048,
        "truncation": "error",
    },
    "actor_rollout_ref": {
        "rollout": {
            "tensor_model_parallel_size": 1,
            "n": 4,  # Generate 4 responses per sampling
            "log_prob_micro_batch_size_per_gpu": 4,
            "multi_turn": {"format": "hermes"},  # Ensure using template format matching the model
            "name": "vllm",
            "gpu_memory_utilization": 0.6,  # vLLM GPU memory utilization
            "engine_kwargs": {
                "vllm": {
                    "enable_auto_tool_choice": True,
                    "tool_call_parser": "hermes",
                }
            },
        },
        "actor": {
            "ppo_mini_batch_size": 16,
            "ppo_micro_batch_size_per_gpu": 4,
            "optim": {"lr": 1e-6},
            "use_kl_loss": False,
            "kl_loss_coef": 0.0,
            "entropy_coeff": 0,
            "clip_ratio_low": 0.2,
            "clip_ratio_high": 0.3,
            "fsdp_config": {
                "param_offload": True,  # Enable parameter offloading to save GPU memory
                "optimizer_offload": True,
            },
        },
        "ref": {
            "log_prob_micro_batch_size_per_gpu": 8,
            "fsdp_config": {"param_offload": True},
        },
        "model": {
            "path": "Qwen/Qwen2.5-1.5B-Instruct",  # Default model
            "use_remove_padding": True,
            "enable_gradient_checkpointing": True,
        },
    },
    "trainer": {
        "n_gpus_per_node": 1,
        "val_before_train": True,
        "critic_warmup": 0,
        "logger": ["console"],  # Disable wandb for easier local debugging, add back when needed
        "project_name": "AgentLightning",
        "experiment_name": "rag_agent",
        "nnodes": 1,
        "test_freq": 10,
        "total_epochs": 200,
    },
}


def config_train_fast() -> Dict[str, Any]:
    """Fast training configuration for CI/testing"""
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    random_suffix = uuid.uuid4().hex[:8]
    EXPERIMENT_NAME = f"rag_fast_{timestamp}_{random_suffix}"

    PROJECT_NAME = "AgentLightningCI"

    # Simulate writing to $GITHUB_OUTPUT if it’s set
    github_output = os.getenv("GITHUB_OUTPUT")
    if github_output:
        with open(github_output, "a") as f:
            f.write(f"project_name={PROJECT_NAME}\n")
            f.write(f"run_name={EXPERIMENT_NAME}\n")

        print("Set environment variables:")
        print(f"PROJECT_NAME={PROJECT_NAME}")
        print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")

    config = deepcopy(RL_TRAINING_CONFIG)

    # Keep it tiny/light without adding new knobs
    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.8
    config["trainer"]["total_epochs"] = 2
    config["trainer"]["test_freq"] = 5
    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
    config["trainer"]["project_name"] = PROJECT_NAME
    config["trainer"]["logger"] = ["console", "wandb"]
    return config


def config_train_single_gpu() -> Dict[str, Any]:
    """Single GPU training optimized configuration (optimized for 24GB GPU memory)"""

    config = deepcopy(RL_TRAINING_CONFIG)

    # 1. Reduce vLLM memory usage to leave space for training
    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.4

    # 2. Reduce Batch Size to prevent OOM
    config["data"]["train_batch_size"] = 4
    config["actor_rollout_ref"]["actor"]["ppo_mini_batch_size"] = 4
    config["actor_rollout_ref"]["actor"]["ppo_micro_batch_size_per_gpu"] = 1
    config["actor_rollout_ref"]["rollout"]["log_prob_micro_batch_size_per_gpu"] = 2

    # 3. Ensure Offload is enabled
    config["actor_rollout_ref"]["actor"]["fsdp_config"]["param_offload"] = True
    config["actor_rollout_ref"]["actor"]["fsdp_config"]["optimizer_offload"] = True

    return config


def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
    """Train the RAG agent with the given configuration."""

    # 1. Instantiate your Agent
    agent = RAGAgent()

    # 2. Initialize algorithm (VERL)
    algorithm = agl.VERL(config)

    # 3. Initialize Trainer
    # n_runners=4 means 4 concurrent rollout runners (can be reduced if insufficient memory, or managed internally by VERL)
    trainer = agl.Trainer(n_runners=4, algorithm=algorithm, adapter={"agent_match": active_agent})

    # 4. Load data
    # NOTE: Fill in the path to your previously converted parquet file here
    # For demo purposes, we use the same dataset for training and validation,
    # which should be avoided in production.
    train_df: pd.DataFrame = pd.read_parquet("data/dataset_tiny.parquet")  # type: ignore
    val_df: pd.DataFrame = pd.read_parquet("data/dataset_tiny.parquet")  # type: ignore

    # Keep the rest of the code unchanged
    train_data: List[Dict[str, Any]] = train_df.to_dict(orient="records")  # type: ignore
    val_data: List[Dict[str, Any]] = val_df.to_dict(orient="records")  # type: ignore

    # 5. Start training
    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)


def main() -> None:
    parser = argparse.ArgumentParser(description="Train a RAG agent using different configurations")

    parser.add_argument(
        "config",
        choices=["fast", "single_gpu"],
        default="single_gpu",
        nargs="?",
        help="Training configuration name",
    )

    parser.add_argument("--active-agent", type=str, help="Override the active agent name")

    args = parser.parse_args()

    config_functions = {
        "fast": config_train_fast,
        "single_gpu": config_train_single_gpu,
    }
    config = config_functions[args.config]()

    # Print key information for confirmation
    print(f"Starting training with '{args.config}' configuration...")
    print(f"Model: {config['actor_rollout_ref']['model']['path']}")
    print(f"Batch Size: {config['data']['train_batch_size']}")
    print(f"GPU Mem Util: {config['actor_rollout_ref']['rollout']['gpu_memory_utilization']}")

    train(config, args.active_agent)


if __name__ == "__main__":
    main()