# Copyright (c) Microsoft. All rights reserved. """Training helper for ChartQA modeled VERL workflow. Example usage: ```bash python train_chartqa_agent.py debug --n-runners 2 ``` or: ```bash AGL_MANAGED_STORE=0 python train_chartqa_agent.py qwen --external-store-address http://localhost:9999 ``` Make sure to run `python prepare_data.py` so the parquet files referenced here exist. """ from __future__ import annotations import argparse import os import uuid from copy import deepcopy from datetime import datetime from typing import Any, Dict, Optional, cast import env_var as chartqa_env_var import pandas as pd from chartqa_agent import ChartQAAgent import agentlightning as agl from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var RL_CONFIG: Dict[str, Any] = { "algorithm": {"adv_estimator": "grpo", "use_kl_in_reward": False}, "data": { "image_base_dir": chartqa_env_var.CHARTQA_IMAGES_DIR, "train_batch_size": 32, "max_prompt_length": 4096, "max_response_length": 1024, "truncation": "error", }, "actor_rollout_ref": { "rollout": { "tensor_model_parallel_size": 1, "n": 4, "log_prob_micro_batch_size_per_gpu": 1, "name": "vllm", "gpu_memory_utilization": 0.8, "enable_prefix_caching": True, "engine_kwargs": {"vllm": {"allowed_local_media_path": chartqa_env_var.CHARTQA_IMAGES_DIR}}, }, "actor": { "ppo_mini_batch_size": 32, "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, "optimizer_offload": True}, }, "ref": {"log_prob_micro_batch_size_per_gpu": 1, "fsdp_config": {"param_offload": True}}, "model": { "path": "Qwen/Qwen2-VL-2B-Instruct", "use_remove_padding": True, "enable_gradient_checkpointing": True, }, }, "trainer": { "n_gpus_per_node": 1, "val_before_train": False, "critic_warmup": 0, "logger": ["console", "wandb"], "project_name": "AgentLightning", "experiment_name": "chartqa", "nnodes": 1, }, } def config_ci() -> Dict[str, Any]: """Return a CI-friendly RL config for ChartQA.""" # For CI testing, we need to set the experiment name and project name so that # they are available to subsequent steps. timestamp = datetime.now().strftime("%Y%m%d%H%M%S") random_suffix = uuid.uuid4().hex[:8] EXPERIMENT_NAME = f"chartqa_ci_{timestamp}_{random_suffix}" PROJECT_NAME = "AgentLightningCI" 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") config = deepcopy(RL_CONFIG) config["data"]["train_batch_size"] = 16 config["trainer"]["n_gpus_per_node"] = 1 config["trainer"]["total_training_steps"] = 4 config["trainer"]["val_before_train"] = True config["trainer"]["test_freq"] = 2 config["trainer"]["experiment_name"] = EXPERIMENT_NAME config["trainer"]["project_name"] = PROJECT_NAME return config def config_debug() -> Dict[str, Any]: """Return a short debugging config for smoke testing ChartQA training.""" config = deepcopy(RL_CONFIG) config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.5 config["trainer"]["total_training_steps"] = 10 config["trainer"]["test_freq"] = 2 return config def config_qwen() -> Dict[str, Any]: """Return a Qwen-focused config with validation before each epoch.""" config = deepcopy(RL_CONFIG) config["trainer"]["val_before_train"] = True config["trainer"]["n_gpus_per_node"] = 2 config["trainer"]["total_epochs"] = 2 config["trainer"]["test_freq"] = 32 return config def train( config: Dict[str, Any], train_data: agl.Dataset[Any], val_data: agl.Dataset[Any], external_store_address: str, n_runners: int, debug: bool, ) -> None: """Run VERL training for ChartQA. Args: config: VERL configuration produced by one of the helpers above. train_data: Training dataset of ChartQA samples. val_data: Validation dataset for periodic evaluation. external_store_address: Optional address of an existing LightningStore to reuse. n_runners: Number of runners passed to [`Trainer.fit`][agentlightning.Trainer.fit]. debug: Enables verbose logging tied to `--debug`. """ agl.setup_logging(level="DEBUG" if debug else "INFO", apply_to=["agentlightning", __name__]) agent = ChartQAAgent() algorithm = agl.VERL(config) if external_store_address: store: Optional[agl.LightningStore] = agl.LightningStoreClient(external_store_address) else: store = None trainer = agl.Trainer( n_runners=n_runners, algorithm=algorithm, store=store, ) trainer.fit(agent, train_dataset=train_data, val_dataset=val_data) # type: ignore def main(): """Parse CLI arguments and kick off ChartQA training.""" agl.setup_logging(apply_to=["chartqa_agent"]) parser = argparse.ArgumentParser(description="Train ChartQA agent") parser.add_argument("config", choices=["debug", "qwen", "ci"], help="Training configuration") parser.add_argument("--n-runners", type=int, default=10, help="Number of runners for Trainer") parser.add_argument( "--external-store-address", type=str, default=None, help="Connect to an external store instead of creating a new one in memory (e.g., http://localhost:4747)", ) parser.add_argument("--debug", action="store_true", help="Enable debug logging") args = parser.parse_args() if args.external_store_address: print(f"Connecting to external store at: {args.external_store_address}") if resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=True): raise ValueError( "When using an external store, please set the environment variable AGL_MANAGED_STORE=0. " "Otherwise the trainer will still try to manage the store lifecycle for you!" ) CONFIGS = { "debug": config_debug, "qwen": config_qwen, "ci": config_ci, } train_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "train_chartqa.parquet") val_data_path = os.path.join(chartqa_env_var.CHARTQA_DATA_DIR, "test_chartqa.parquet") train_data = pd.read_parquet(train_data_path).to_dict(orient="records") # type: ignore if args.config in ["debug", "ci"]: val_data = pd.read_parquet(val_data_path).sample(n=100, random_state=42).to_dict(orient="records") # type: ignore else: val_data = pd.read_parquet(val_data_path).to_dict(orient="records") # type: ignore train( config=CONFIGS[args.config](), train_data=cast(agl.Dataset[Any], train_data), val_data=cast(agl.Dataset[Any], val_data), external_store_address=args.external_store_address, n_runners=args.n_runners, debug=args.debug, ) if __name__ == "__main__": main()