VLM-R1 实战训练教程

VLM-R1 是deepseek R1 的GRPO训练方式在VLM中的实现。项目地址：https://github.com/om-ai-lab/VLM-R1

基础环境

构建docker环境：

bash
# docker pull pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel

docker pull pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel

# docker run -it --gpus all pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel bash 

docker run -it --gpus all pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel bash 

apt update

apt install vim git -y

git clone https://github.com/om-ai-lab/VLM-R1.git

cd VLM-R1/

bash setup.sh

pip install transformers==4.49.0

此docker环境已经被我上传为 kevinchina/deeplearning:2.5.1-cuda12.4-cudnn9-devel-vlmr1

使用此镜像：


docker run -it --gpus '"device=0,1,2,3,4,5,6,7"' --shm-size=64g -v /data/xiedong:/data/xiedong --net host kevinchina/deeplearning:2.5.1-cuda12.4-cudnn9-devel-vlmr1 bash

下载预训练模型


pip install modelscope

# 下载到这个目录
/data/xiedong# modelscope download --model Qwen/Qwen2.5-VL-3B-Instruct --local_dir Qwen2.5-VL-3B-Instruct

GRPO训练

准备数据：

wget https://huggingface.co/datasets/omlab/VLM-R1/resolve/main/train2014.zip
wget https://huggingface.co/datasets/omlab/VLM-R1/resolve/main/rec_jsons_processed.zip
在src/open-r1-multimodal/data_config/rec.yaml中写入数据路径

vim src/open-r1-multimodal/data_config/rec.yaml


datasets:
    - json_path: /data/xiedong/rec_jsons_processed/refcoco_train.json
    - json_path: /data/xiedong/rec_jsons_processed/refcocop_train.json
    - json_path: /data/xiedong/rec_jsons_processed/refcocog_train.json

运行训练bash src/open-r1-multimodal/run_scripts/run_grpo_rec.sh

这文件里面长这样： vim src/open-r1-multimodal/run_scripts/run_grpo_rec.sh


cd src/open-r1-multimodal

export DEBUG_MODE="true"
# export CUDA_VISIBLE_DEVICES=4,5,6,7

RUN_NAME="Qwen2.5-VL-3B-GRPO-REC"
export LOG_PATH="./debug_log_$RUN_NAME.txt"

torchrun --nproc_per_node="8" \
    --nnodes="1" \
    --node_rank="0" \
    --master_addr="127.0.0.1" \
    --master_port="12346" \
    src/open_r1/grpo_rec.py \
    --deepspeed local_scripts/zero3.json \
    --output_dir output/$RUN_NAME \
    --model_name_or_path /data/xiedong/Qwen2.5-VL-3B-Instruct \
    --dataset_name data_config/rec.yaml \
    --image_root /data/xiedong \
    --max_prompt_length 1024 \
    --num_generations 8 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --logging_steps 1 \
    --bf16 \
    --torch_dtype bfloat16 \
    --data_seed 42 \
    --report_to wandb \
    --gradient_checkpointing false \
    --attn_implementation flash_attention_2 \
    --num_train_epochs 2 \
    --run_name $RUN_NAME \
    --save_steps 100 \
    --save_only_model true

训练时候：

torchrun: 一个用于分布式训练的 PyTorch 脚本运行工具。
--nproc_per_node="8": 指定每个节点上的进程数为 8。
--nnodes="1": 指定训练使用的节点数为 1。
--node_rank="0": 指定当前节点的排名，这里是0号节点。
--master_addr="127.0.0.1": 指定 master 节点的地址，用于分布式训练的通信。
--master_port="12346": 指定 master 节点的端口号。
src/open_r1/grpo_rec.py: 运行的 python 脚本。
--deepspeed local_scripts/zero3.json: 使用 DeepSpeed 优化器并指定配置文件 local_scripts/zero3.json。
--output_dir output/$RUN_NAME: 指定输出目录，训练结果将保存到 output/$RUN_NAME。
--model_name_or_path /data/xiedong/Qwen2.5-VL-3B-Instruct: 指定预训练模型的路径。
--dataset_name data_config/rec.yaml: 指定数据集配置文件的路径。
--image_root /data/xiedong: 指定图像数据根目录。
--max_prompt_length 1024: 指定生成模型输入的最大 prompt 长度。
--num_generations 8: 指定每个输入生成的输出序列数。
--per_device_train_batch_size 1: 指定每个设备上的训练批次大小为 1。
--gradient_accumulation_steps 2: 指定梯度积累的步数，即在执行反向传播之前累积的梯度步数。
--logging_steps 1: 每隔多少步进行一次日志记录。
--bf16 和 --torch_dtype bfloat16: 指定使用 bfloat16 数据类型进行训练，以更好地利用 GPU 的计算能力。
--data_seed 42: 指定用于数据乱序的随机种子。
--report_to wandb: 使用 Weights & Biases 工具进行实验跟踪和报告。
--gradient_checkpointing false: 指定是否使用梯度检查点技术以减少显存使用。这里设置为 false。
--attn_implementation flash_attention_2: 指定注意力机制的实现方法。
--num_train_epochs 2: 指定训练的总轮数为 2。
--run_name $RUN_NAME: 指定运行名称，用于标识当前实验。
--save_steps 100: 指定每隔 100 步保存一次模型。
--save_only_model true: 指定仅保存模型权重，不保存优化器状态等信息。

多个数据文件和图像文件夹可以使用":"作为分隔符指定：


--data_file_paths /path/to/data1.jsonl:/path/to/data2.jsonl \
--image_folders /path/to/images1/:/path/to/images2/