Agentic Training/Inference Benchmarking for Tool-Calling LLMs on M3

This project benchmarks large language models (LLMs) on their ability to use tools effectively in a multi-hop, multi-turn, multi-source/tool question answering agentic setting. The agent operates in a gym-style MDP environment, solving tasks by invoking tools at each step and producing a final answer through a FINAL action.

Releases

Existing

We currently support SFT and DPO training. Specifically, for DPO we have step-level v/s trajectory-level preference data generation and training support. Additionally, by looking at step-level reward categories, we further have an option to mask suboptimal traces while creating trajectory-level preferences during training.

Future

📚 Table of Contents

• Environment Setup
• Agentic Inference Environment
• Environment Scenarios
• Supported Language Models
• Data Generation
• Training
• Evaluation

---

🧱 Python Environment Setup

To get started, create a Python environment and install the necessary dependencies.

🔧 Step 1: Create Conda Environment

conda create -n AgenticAI python=3.11
conda activate AgenticAI

📦 Step 2: Install Dependencies

pip install -r requirements.txt

This will install all required packages for running the agentic pipeline, including support for LoRA training, multi-turn inference, DeepSpeed, and more.

If you intend to tune granite 4, you will also need to install the following packages from source (not PyPi)

• causal-conv1d>=1.4.0
• mamba-ssm[causal-conv1d]==2.2.4

git clone [email protected]:Dao-AILab/causal-conv1d.git
cd causal-conv1d
pip install .
cd ..
git clone [email protected]:state-spaces/mamba.git
cd mamba
pip install -e .

---

🧭 Agentic Inference Environment

The agent acts as a decision-making policy in an MDP, invoking one tool per step, and emitting a FINAL action with the answer when complete.

📚 Amongst the tools available which are mostly API endpoints in M3, there is a Special Tool: retrieve_documents (name defined at constants.py)

A core tool that allows the agent to query a knowledge base. Required arguments:

• collection: Name of the document collection (e.g., "clapnq-company_docs")
• query: The query string to retrieve relevant documents
• [Future Work] top_k: We can consider top_k as a predicted argument to have the agent even reason over document collection. What combination of query and top_k fetches the ground-truth document in order to optimise the query. Currently, "top_k" is kept constant and provided in the config file.

Example tool call:

{
  "tool_name": "retrieve_documents",
  "args": {
    "collection": "clapnq-company_docs",
    "query": "What is the employee leave policy for 2025?"
  }
}

---

🔍 Inference Setup: Pre-requisites

• [Setup retriever] To use database index, you need to acquire the username, password, host_name, certificate. Save the certificate in es_cert at the root directory. Set username, password as ES_USERNAME , ES_PASSWORD and ES_HOSTNAME in the .env
• [Setup API Endpoint] To use and query APIs, gain access to api endpoint
• [Setup LLMs] Set the following parameters to use the Language Models and save them in the .env at the root directory
  • RITS_API_KEY: API key to use language models via RITS
  • HF_TOKEN: Huggingface token

---

🛠️ Environment Configuration

The parameters for governing the structure and behavior of the agent and its task environment are defined in: 📄 configs/infer_agent.json

---

🧠 Expert Intervention/Assistance Modes

Here's how the expert intervention mechanism works:

The framework supports various expert-assisted settings for supervising or augmenting agent decisions. Controls if and how the expert (which has access to ground-truth trajectories) is used. Set via expert_assist_mode:

Additional expert parameters for the case when expert_mode is set to informed:

Note: Collecting preference data by enabling expert intervention only collects the chosen and rejected samples at step-level.

Future Work: Expand get_alternate_action_trace() function in run.py to collect complete alternate trajectories starting at given state by recursively expanding akin to traversing possible paths in a binary tree

---

🧩 Environment Scenarios

The agent operates within configurable environment scenarios that define both tool availability and tool usage constraints. These scenarios determine how the agent interacts with tools while attempting to solve a task.

Where to specify? They must be provided for each sample/environment instance within the path_to_env_data. See env.reset() for further information on how it is loaded.

🔌 Tool Availability (tool_availability)

Specifies which tools are available to the agent during inference:

🚫 Tool Usage Policy (tool_usage)

A textual constraint that defines policies on tool invocation for specific subdomains or document collections. For example, it can restrict the agent from using the retriever on certain collections or the API in specific subdomains.

• Default: "" (no restrictions)
• Example: "Do not use the tool $retrieve_documents$ for the domain dealing with company documents."

⚠️ Tool Ablation Behavior

If the agent requires a tool that is made unavailable (either through tool_availability or tool_usage), it is expected to:

Acknowledge its inability to gather the necessary information and terminate with a FINAL action indicating that it cannot answer the query due to tool limitations.

This ensures clarity in evaluation and enables benchmarking tool-dependency across models.

---

🤖 Supported Language Models

The following models are supported as agents:

Set the parameters for the agentic model to use via the key:

{
  "agent_model_name_or_path": "mistralai/Mistral-7B-Instruct-v0.3",
  "agent_template": "teacher_rits",
  "is_hf_agent": false,
  "temperature": 0.0,
  "max_new_tokens": 4096
}

Description of other parameters configuring model and its behaviour:

• is_hf_agent: Boolean whether the LLM to be used for the agent is provided by Huggingface or RITS
• agent_template: One of student_mistral, student_qwen3, student_granite3, teacher_rits
• max_new_tokens: Number of tokens to generate
• temperature: Sampling temperature
• path_to_hf_config: Path to json file specifying other parameters for running the HF model

---

📦 Data Generation

Training data is generated in two stages, both designed to provide supervision signals for learning i.e. supervised fine-tuning (SFT) and preference modeling (DPO). Each step of each trajectory is annotated with a reward category, for GRPO like training methods.

Here’s a clear and well-structured “📦 Data Generation” section for your README:

🛠️ Preliminary Setup

Before generating any data, begin by parsing the raw environment data into a pipeline-consumable format.

📄 Refer to the README in the ground_truth directory for details on how to generate this preprocessed format.

🧭 Stage 1: Generating Expert Trajectories (SFT Data)

Set the following configuration in configs/infer_agent.json:

{"expert_mode": "ground_truth"}

This will run the expert policy to solve each task fully, producing gold-standard trajectories used for supervised fine-tuning (SFT).

📁 Output: Trajectories in a standardized format, one per task instance

🔍 Stage 2: Generating Preference Data (Exploratory Agent with Expert Intervention)

Set the following configuration in configs/infer_agent.json:

{"expert_mode": "informed"}

In this mode, the agent interacts with the environment, and the expert intervenes when necessary (e.g., if the agent is stuck or making repeated errors). This generates step-level preference data, based on whether the agent or the expert took the action.

⚙️ You must ensure the underlying LLM agent is set up correctly:

• ✅ Either via RITS or a local Hugging Face endpoint
• 🔗 See the previous section on how to configure model-specific parameters like model_name_or_path and chat_template

[Future Work]: Work in progress to generate preference data at trajectory level by comparing trajectories of two separate agents (one better than other) here at create_pref_trajectories_for_dpo.py. For this we need trajectories from two separate agents.

🚀 Running the Pipeline

• If using the CCC cluster, set ENV_NAME, ROOT_DIR and HF_CACHE_DIR in the shell script, and then you can launch the data generation using the provided shell script:

  bsub < ./ccc_scripts/infer_agent.lsf.sh
  

  Example Usage:

  # Assign arguments to variables
  ENV_NAME="AgenticAI"
  ROOT_DIR="/u/aj05/project/Code"
  HF_CACHE_DIR="/u/aj05/project/hf_cache"
  

  bsub < ./ccc_scripts/infer_agent.lsf.sh
  

• Otherwise, run the entry point directly:

  python run.py
  

Ensure that your configs/infer_agent.json is correctly populated before running. The run results i.e. trajectories will be generated and saved in the logging/<curr_date_time>/trajectories with corresponding metadata.

Click to view example metadata

```json { "sample_id": 17, "truncated": false, "terminated": true, "success": true, "total_time_steps": 3, "expert_assistance": { "mode": "ground_truth", "init_limit": 2, "recent_limit": -2, "random_epsilon": 1.0, "tracker": [ 1, 1, 1 ] } } ```

Field Name	Explanation
sample_id	sample ID of the environment instance (not the same as actual data's sample IDs)
truncated	Whether the agent run was truncated by the environment because it exceeded max steps
terminated	Whether the agent run was terminated by the environment
success	Whether the agent successfully completed the task (determined by LLMasJudge)
total_time_steps	Number of time-steps the agent ran for (across multi-turns)
expert_assistance	Parameters with which expert assistance was configured
expert_assistance/tracker	For each step, whether the action was taken by expert (1) or agent (0)

---

🚀 Training: LoRA-Based Fine-Tuning of Agentic Language Models

To train the agentic framework, we use LoRA-based parameter-efficient fine-tuning of foundation models, optimizing for tool-use behavior grounded in expert demonstrations.

All training configurations are specified in 📄 config_files/train_lora.json

The training script is: 📄 tune.py

---

⚙️ Configure Training Parameters (config_files/train_lora.json)

Note the configuration parameters specified in config_files/train_lora.json overrides the default values specified for the same at hparams which is an extensive suite of hyper-params configuring the model training/evaluation.

Below we discuss the key ones:-

Model (overrides model_args.py)

To optimise memory footprint and training time, we have set the following additional params: flash_attn = fa2, disable_gradient_checkpointing = false, enable_liger_kernel = true

Training Data (overrides data_args.py)

DPO-specific parameters that are unique to our use-case are:

• preference_granularity: step (collected using single agent) or trajectory (collected using multiple agents). Default set to step
• mask_suboptimal_traces: Whether to mask suboptimal traces in preference samples. Used with trajectory-level preference granularity

Finetuning Method (overrides finetuning_args.py)

For DPO, the additional params that need to specified are pref_loss, pref_beta, ref_model, ref_model_adapters. We have implemented support for a wide range of preference losses.

Training Hyper-params (overrides training_args.py and built-in Transformer's)

---

🚀 Distributed Multi-GPU Training (with DeepSpeed)

We support multi-GPU distributed training via Hugging Face Accelerate and DeepSpeed, using:

• Stage 2: Optimized memory usage via gradient and optimizer state partitioning
• Stage 3: Full parameter, gradient, and optimizer sharding for large-scale setups

✅ Required Files:

⚠️ Important:

• The DeepSpeed config file path must be provided both in:
  • the deepspeed key inside train_lora.json
  • the --config_file (Accelerate YAML) when launching training
• num_processes field in Accelerate YAML file must be set to number of available GPUs

---

💻 Example Training Command

• If using the CCC cluster, set ENV_NAME, ROOT_DIR and HF_CACHE_DIR in the shell script, and then you can launch the training using the provided shell script:

  bsub < ./ccc_scripts/train_multi_gpu.lsf.sh
  

• Otherwise, run the entry point directly

  accelerate launch --config_file config_files/training/multi_gpu_ds_stage3.yml tune.py
  

📌 Notes

• For faster debugging or single-GPU setups, set deepspeed to null and use a basic accelerate config setup

---

✅ Evaluation

To evaluate an agent on held-out test data, you can reuse the same pipeline used for training data generation.

🚀 Run Evaluation

Use the same entry point:

python run.py

But with the following configuration changes in 📄 configs/infer_agent.json:

🔧 Required Configuration

🤖 Evaluating HF Base Models

For evaluating a base Hugging Face model:

• Set is_hf_agent = true
• Point path_to_hf_config to: 📄 ./config_files/infer_base.json

🧠 Evaluating Fine-Tuned (LoRA) Models

For evaluating trained models:

• Set is_hf_agent = true
• Point path_to_hf_config to: 📄 ./config_files/infer_lora.json
• In infer_lora.json, set:

{"adapter_name_or_path": "<path_to_trained_model>"}

This loads your fine-tuned weights via PEFT LoRA.

---

📊 Output

After evaluation, the following will be generated:

• ✅ Saved Trajectories (full step-by-step tool calls)
• 🧾 Metadata (intervention counts, turn-level info)
• 📈 Success Rate of the agent across the test set

The output format mirrors that of the training data generation pipeline, ensuring compatibility for post-hoc analysis or visualization.

---