Skip to content
 
 

Latest commit

 

History

8 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

MM-PRM logo

MM-PRM

📖Blog | 🤗MM-PRM-8B


MM-PRM: An open implementation of OmegaPRM and its corresponding training pipeline


🎯Overview

Multimodal Large Language Models (MLLMs) have shown promising performance on various reasoning tasks, yet their ability to reliably solve complex multi-step problems remains limited. Process Reward Models (PRMs) address this challenge by explicitly evaluating the correctness of each intermediate reasoning step, guiding models toward more robust solutions. However, effectively training PRMs typically requires substantial amounts of step-level supervision data, which are expensive and challenging to obtain.

We provide a complete implementation of OmegaPRM, an automated Monte Carlo Tree Search-based data pipeline, to generate scalable and high-quality multimodal step-level supervision data along with its corresponding training pipeline. Using this pipeline, we introduce a multimodal PRM based on the InternVL series.

We open-source our complete pipeline to foster further research in this area. We release all our codes, models, etc. at MM-PRM.

🗞️ News

🤖 Models

Case Study

Figure 1 | A case study showcasing the output of MM-PRM-8B. The PRM assigns confidence scores to each reasoning step and successfully distinguishes correct steps from incorrect ones.

Random (baseline) min last average max sum_log prob sum_logits mean_odd
K12 36.2 43.8 (+7.6) 41.8 (+5.6) 44.0 (+7.8) 42.2 (+6.0) 44.0 (+7.8) 43.8 (+7.6) 42.4 (+6.2)
OlympiadBench 15.0 19.3 (+4.3) 18.0 (+3.0) 19.3 (+4.3) 18.7 (+3.7) 18.7 (+3.7) 20.7 (+5.7) 18.7 (+3.7)

Table 1 | Evaluation results of MM-PRM-8B under Best-of-N (BoN, 16 rollout) metrics. K12 represents in-domain setting, while OlympiadBench serves as an out-of-domain benchmark. Our PRM consistently improves over the random baseline across multiple aggregation strategies.

🏁 Getting Started

📦 Installation

git clone https://github.com/ModalMinds/MM-PRM.git
cd MM-PRM
pip install -r requirements.txt

# install flash-attn==2.3.6:

pip install flash-attn==2.3.6 --no-build-isolation

# Alternatively you can compile from source:

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
git checkout v2.3.6
python setup.py install

📂 OmegaPRM Data Pipeline

  1. Seed dataset preparation

    To begin, prepare a seed dataset consisting of verifiable problems. Each example should be formatted as a JSON object containing the following fields:

    [
        {
            "id": "unique identifier for the problem",
            "question": "problem statement",
            "correct_answer": "ground-truth final answer for evaluation and verification",
            "image_path": "/path/to/image.png"
        },
        ...
    ]

    This dataset will be used as input to the OmegaPRM engine to generate annotated solution trees with step-wise correctness labels.

    To enable parallel data generation, you need to split the seed dataset into smaller chunks.

    cd data_pipeline
    python process_json.py
  2. API endpoint setup (Optional)

    The OmegaPRM data generation process requires an API endpoint to automatically verify whether the final answer in a rollout is correct. You can deploy a model (e.g., Qwen2.5) locally to act as the answer judge.

    We recommend using vLLM to deploy a local model.

  3. Run OmegaPRM

    Once you have all set, you can run the OmegaPRM pipeline to generate step-level supervision data.

    Before running, ensure that all necessary parameters are correctly set in the script or passed through the environment.

    sh run_omegaprm.sh
  4. Sampling Training Data from OmegaPRM annotation trees

    After generating annotated reasoning trees with OmegaPRM, you need to sample step-by-step solution paths from these trees to construct the training data for the Process Reward Model (PRM). This can be done using the script:

    python traverse.py

    The next step is to convert this data into the format required for PRM training. Use the following script to perform the formatting:

    python prm_data_format.py

🌐 Start PRM Training

Create a JSON file in internvl_chat/shell/data/

The format for the JSON file should be:

{
  "your-custom-prm_dataset": {
    "root": "/path/to/the/image/root",
    "annotation": "/path/to/the/jsonl/annotation",
    "data_augment": false,
    "repeat_time": 1,
    "length": "number of samples in the dataset"
  }
}

Once the dataset configuration is in place, you can start training the PRM model with:

GPUS=8 sh shell/internvl2.5/2nd_finetune/internvl2_5_38b_dynamic_res_2nd_finetune_full_prm.sh

📊 Evaluation

We provide our evaluation code in the eval/ directory.

⭐ Starchart

Star History Chart

🤝 Contribution

If you want to contribute, please feel free to make a pull request or create an issue.

Please refer to CONTRIBUTING.md before you dive in!

📬 Contact

If you have any questions or would like to engage with our community, feel free to scan the QR code below to join our WeChat group.

MM-PRM logo

🎓 Acknowledgements

We acknowledge the outstanding open-source contributions from OpenR and vLLM. We also extend our gratitude to InternVL for their open-source techniques and base models, which have enabled us to further our exploration.

📜 Citation

@misc{MM-PRM2025,
  title={MM-PRM: An open implementation of OmegaPRM and its corresponding training pipeline},
  author={ModalMinds Team},
  year={2025},
  howpublished={\url{https://github.com/ModalMinds/MM-PRM}},
}

About

MM-PRM: An open implementation of Multimodal OmegaPRM and its corresponding training pipeline

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages