How to run AgentFlow Benchmark with a locally trained Flow-GRPO checkpoint

[arghavan-kpm]

Hi,

I trained a model using the Flow-GRPO training pipeline following the README setup:

• Base model: Qwen/Qwen2.5-7B-Instruct (served via vLLM)
• MODEL_ENGINE in config.yaml: ["trainable", "gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini"]
  (i.e., I used gpt-4o-mini for the non-trainable model engines, while the trainable engine is the local vLLM-served model)

Now I’m trying to evaluate the trained checkpoint using the “AgentFlow Benchmark” section in the README. My understanding is:

1. Serve the model checkpoint with vLLM
2. Run the benchmark run.sh script

Question: Is that the correct evaluation flow?

If so, I’m confused about how to use my local training checkpoint with the provided benchmark scripts. The example scripts seem written for the published model on Hugging Face (AgentFlow-7B/agentflow-planner-7b). I’m not sure what to change in:

• serve_vllm.sh
• run.sh

…so they load my checkpoint saved under:
checkpoints/AgentFlow_pro/AgentFlow_pro/global_step_*

Could you please clarify what edits are needed and the recommended way to point the benchmark scripts to a locally trained checkpoint?

Thanks!

[IsaacGHX]

Hi @arghavan-kpm ,

Thanks for your interest in AgentFlow! Your understanding of the evaluation flow is correct:

1. Merge/Convert your checkpoint to a standard Hugging Face format.
2. Serve it using vLLM.
3. Run the benchmark using the provided scripts.

Regarding your local checkpoints in checkpoints/AgentFlow_pro/..., since we use FSDP (Fully Sharded Data Parallel) for training, the saved weights are typically sharded. To use them with vLLM, you first need to merge these shards into a single consolidated model file.

Here is a reference script you can use to merge the FSDP shards into a Hugging Face compatible format:

import os
import re
import torch
from torch.distributed._tensor import DTensor
from transformers.utils import WEIGHTS_NAME

def merge_fsdp_shards(ckpt_dir: str):
    # detect world size
    world_size = None
    for f in os.listdir(ckpt_dir):
        m = re.match(r"model_world_size_(\d+)_rank_0\.pt", f)
        if m:
            world_size = int(m.group(1))
            break
    assert world_size is not None, "No FSDP shards found"

    # load shards
    shards = []
    for rank in range(world_size):
        path = os.path.join(
            ckpt_dir, f"model_world_size_{world_size}_rank_{rank}.pt"
        )
        shards.append(torch.load(path, map_location="cpu", weights_only=True))

    merged_state_dict = {}

    for key in shards[0]:
        values = [sd[key] for sd in shards]
        if isinstance(values[0], DTensor):
            placement = values[0].placements[0]
            merged = torch.cat(
                [v._local_tensor for v in values],
                dim=placement.dim,
            )
            merged_state_dict[key] = merged
        else:
            merged_state_dict[key] = values[0]

    hf_dir = os.path.join(ckpt_dir, "huggingface")
    os.makedirs(hf_dir, exist_ok=True)
    torch.save(merged_state_dict, os.path.join(hf_dir, WEIGHTS_NAME))

    print(f"Merged HF checkpoint saved to: {hf_dir}")

# Usage:
# merge_fsdp_shards("checkpoints/AgentFlow_pro/AgentFlow_pro/global_step_XXXX")

Steps to update your scripts:

1. Run the merge script: Point it to your global_step_* directory. It will create a huggingface/ sub-folder with the pytorch_model.bin. Make sure to copy the config.json, tokenizer.json, etc., from the base Qwen2.5 model into this folder.
2. Update serve_vllm.sh:
   Change the --model path to your local hf_dir path:

# Inside serve_vllm.sh
vllm serve /path/to/your/checkpoints/AgentFlow_pro/.../huggingface --port 8000 ...

3. Update run.sh :
   Simply ensure the MODEL_PATH or the URL matches your local vLLM instance. If you follow the README's default config, as long as vLLM is serving your local model on the expected port, the benchmark script will talk to it correctly.

Hope this helps! Let us know if you have further questions.