chore: import upstream snapshot with attribution

examples/rag/train_rag.py

@@ -0,0 +1,200 @@
+# 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()