We present Workforce, a hierarchical multi-agent framework that decouples planning from execution through a modular architecture with a domain-agnostic Planner, Coordinator, and specialized Workers. This enables cross-domain transfer by allowing worker modification without full system retraining. Our OPTIMIZED WORKFORCE LEARNING (OWL) approach further improves generalization through reinforcement learning. On the GAIA benchmark, Workforce achieves state-of-the-art 69.70% accuracy, outperforming commercial systems. Our OWL-trained 32B model reaches 52.73% accuracy, comparable to GPT-4o. This work provides a foundation for general-purpose AI assistants with scalable domain adaptation.
This repository contains code for the OWL framework, including inference part (Workforce) and training part (OWL).
To reproduce Workforce inference performance (69.70% - Claude-3.7 accuracy on GAIA benchmark and 60.61% - GPT-4o accuracy on GAIA benchmark) shown in the paper, follow the steps below:
- Create a Python 3.11 Conda environment:
conda create -n owl python=3.11- Install the required packages:
pip install -r requirements.txt- Set up envionment variables:
copy .env.example to .env and set the environment variables, and set the keys in .env file.
- Run the inference:
- For reproducing results using GPT-4o, run:
python run_gaia_workforce.py- For reproducing results using Claude-3.7, run:
python run_gaia_workforce_claude.pyWe use LLaMA-Factory to train the Qwen2.5-32B-Instruct model. For GPU resources, 8xH100 GPUs are recommended.
pip install openrlhf
pip install liger-kernel
cd train/LLaMA-Factorypython train/make_sft_data.py
python train/make_dpo_data.py --best_of_n 4 --level 1 --dataset gaia--datasetcan be gaia, imm, wtq-xlsx, wtq-csv, math and hotpotqa
Download the dataset from here.
- For SFT stage, run:
llamafactory-cli train examples/train_full/all_qwen2_32B_planner_scored.yaml- For DPO stage, run:
deepspeed --include=localhost train_planner_dpo.py \
--save_path ./checkpoint/DPO_Qwen2.5-32B-Instruct_workforce_config2 \
--save_steps -1 \
--logging_steps 1 \
--eval_steps 10 \
--train_batch_size 32 \
--micro_train_batch_size 1 \
--pretrain Qwen/Qwen2.5-32B-Instruct \
--bf16 \
--max_epochs 2 \
--max_len 16384 \
--zero_stage 3 \
--beta 0.1 \
--learning_rate 5e-7 \
--dataset dataset/dpo_wtq-csv_workforce_1_config2_w_final_answer.json,dataset/dpo_wtq-xlsx_workforce_1_config2_w_final_answer.json \
--apply_chat_template \
--chosen_key chosen \
--rejected_key rejected \
--flash_attn \
--gradient_checkpointing \
--packing_samples \
--use_liger_kernel \
--adam_offloadHere we use openrlhf to train the model. The --dataset argument should point to the dataset you prepared in the
previous step.
The --pretrain argument should point to the pre-trained model you want to use.
We use vLLM for LLM inference.
Here is an example command for running the inference:
vllm serve [YOUR_MODEL_PATH] \
--model-name [YOUR_MODEL_NAME] \
--dtype=bfloat16 \
--tensor-parallel-size=4 \
--port 25001 \
--enable-auto-tool-choice \
--tool-call-parser hermes
python run_gaia_workforce_vllm_planner.py \
--model_name [YOUR_MODEL_NAME] \
--port 25001