ChartRL: Improving Chart understanding in VLMs

Repo for RL post-training on Chart data. We present a small 3 billion chart understanding model which outputs answers AND rationales. Currently the model is based on Qwen2.5VL (3 billion).

The design hinges on training VLMs using GRPO with verifiable rewards. The training data is sampled from train sets of - ChartQA, PlotQA, and ChartFC (in-domain).

Huge thanks to the team at Morgan Stanley for their support!

Find our Chart-RVR-3B models and associated datasets on Huggingface.

Main Benchmark Results

Running Evals

Run inference on ChartQA with Chain-of-Thought: python main.py --mode eval --vlm-name qwen2-5-3b --dataset-name chartqa-src --cot True

Run inference on ChartQA with GRPO trained models: python main.py --mode eval --vlm-name qwen2-5-3b --dataset-name chartqa-src --cot True --grpo-lora True

Run inference on ChartQA with SFT trained models: python main.py --mode eval --vlm-name qwen2-5-3b --dataset-name chartqa-src --cot True --sft-lora True

Ensure path to models correspond to checkpoints from Chart-RVR or Chart-RVR-Hard

To ease computation, we have provided scripts which can directly be executed on SLURM in ./eval-scripts For instance to run GRPO on all datasets, use: sbatch ./eval-scripts/grpo_all_eval.sh. This runs evaluation on all datasets and stores it in ./logs.

Training Models

SFT

We utilize the great work here to fine-tune models.

Chart-RVR (GRPO-based)

To directly run training, we recommend first double checking TRL version. Next double check the num_processes in ./deepspeed_zero3.yaml is set to 4 (recommended). For any other combination, there might be OOM errors with default completions, prompt size, etc.

Then, simply run on the SLURM cluster sbatch grpo.sh. For running directly on local GPUs, run: accelerate launch --config_file=deepspeed_zero3.yaml main.py --mode grpo --vlm-name qwen2-5-3b --dataset-name chartqa-src

Note: The datset name is only placeholder for the eval-script. It is mandatory to run the dataset generation script first.