> ## Documentation Index > Fetch the complete documentation index at: https://docs.flex.ai/llms.txt > Use this file to discover all available pages before exploring further. # Fine-Tune Vision-Language Models with GRPO on FlexAI > Train vision-language models using GRPO reinforcement learning with TRL, LoRA, vLLM, and DeepSpeed ZeRO-3 on FlexAI. Full tutorial with working code. This experiment demonstrates how to perform Group Relative Policy Optimization (GRPO) training on a Vision-Language Model (VLM) using the TRL library. GRPO is a reinforcement learning technique that optimizes the model based on relative rewards within groups of generated responses. We will use `Qwen2.5-VL-3B-Instruct` as our base VLM and apply LoRA (Low-Rank Adaptation) for efficient fine-tuning. The training leverages vLLM for fast inference during the policy optimization phase and DeepSpeed ZeRO-3 for distributed training. ## Prerequisites ### Connect to GitHub (if needed) If you haven't already connected FlexAI to GitHub, you'll need to set up a code registry connection: ```bash theme={null} flexai code-registry connect ``` This will allow FlexAI to pull repositories directly from GitHub using the `-u` flag in training commands. ### Create Secrets To access the Qwen2.5-VL-3B-Instruct model, you may need authentication with your HuggingFace account depending on the model's access requirements. Use the `flexai secret create` command to store your *HuggingFace Token* as a secret. Replace `` with your desired name for the secret: ```bash theme={null} flexai secret create ``` Then paste your *HuggingFace Token* value. ## Training To start the GRPO Training Job, run the following command: ```bash theme={null} flexai training run vlm-grpo-training --repository-url https://github.com/flexaihq/blueprints \ --requirements-path code/vlm-grpo/requirements.txt \ --secret HF_TOKEN= \ --runtime pytorch-28-vllm-0110-nvidia \ --nodes 1 --accels 2 \ -- accelerate launch \ --config_file=code/vlm-grpo/deepspeed_zero3.yaml \ code/vlm-grpo/grpo_vlm.py \ --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \ --output_dir /output-checkpoint/grpo-Qwen2.5-VL-3B-Instruct \ --learning_rate 1e-5 \ --gradient_checkpointing \ --dtype bfloat16 \ --max_prompt_length 2048 \ --max_completion_length 1024 \ --use_vllm \ --vllm_mode colocate \ --use_peft \ --lora_target_modules "q_proj" "v_proj" \ --log_completions \ --logging_steps 200 ``` ### Key Arguments Explained * `--model_name_or_path`: The base VLM model to fine-tune (Qwen2.5-VL-3B-Instruct) * `--output_dir`: Directory where checkpoints will be saved * `--learning_rate`: Learning rate for optimization (1e-5) * `--gradient_checkpointing`: Enables gradient checkpointing to reduce memory usage * `--dtype bfloat16`: Uses bfloat16 precision for training * `--max_prompt_length`: Maximum length of input prompts (2048 tokens) * `--max_completion_length`: Maximum length of generated completions (1024 tokens) * `--use_vllm`: Enables vLLM for fast inference during policy optimization * `--vllm_mode colocate`: Runs vLLM on the same GPUs as training * `--use_peft`: Enables LoRA for parameter-efficient fine-tuning * `--lora_target_modules`: Specifies which modules to apply LoRA (q\_proj, v\_proj) * `--log_completions`: Logs generated completions during training ### Configuration Files The experiment uses a DeepSpeed ZeRO-3 configuration file (`code/vlm-grpo/deepspeed_zero3.yaml`) that specifies: * ZeRO Stage 3 for memory-efficient distributed training * 2 processes (GPUs) for multi-GPU training * Mixed precision training with bfloat16 ## Monitoring the Training Job You can check the status and life cycle events of your Training Job by running: ```bash theme={null} flexai training inspect vlm-grpo-training ``` Additionally, you can view the logs of your Training Job by running: ```bash theme={null} flexai training logs vlm-grpo-training ``` ## Fetching the Trained Model Once the Training Job completes successfully, you can list all the produced checkpoints: ```bash theme={null} flexai training checkpoints vlm-grpo-training ``` Download a checkpoint with: ```bash theme={null} flexai checkpoint fetch "" ``` The checkpoint will contain the LoRA adapters that can be merged with the base model for inference. ## Additional Resources * [TRL Documentation](https://huggingface.co/docs/trl/index) * [GRPO Paper](https://arxiv.org/abs/2402.03300) * [Qwen2-VL Model Card](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) * [vLLM Documentation](https://docs.vllm.ai/) ## Code ### grpo\_vlm.py ```python theme={null} # Adapted from: https://github.com/huggingface/trl/blob/e622196097109080b73584d598d4162e64fc6bea/examples/scripts/grpo_vlm.py # Copyright 2020-2025 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # /// script # dependencies = [ # "trl", # "Pillow", # "peft", # "math-verify", # "latex2sympy2_extended", # "torchvision", # "trackio", # "kernels", # ] # /// """ pip install math_verify # For Qwen/Qwen2.5-VL-3B-Instruct accelerate launch \ --config_file examples/accelerate_configs/deepspeed_zero3.yaml \ examples/scripts/grpo_vlm.py \ --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \ --output_dir grpo-Qwen2.5-VL-3B-Instruct \ --learning_rate 1e-5 \ --gradient_checkpointing \ --dtype bfloat16 \ --max_prompt_length 2048 \ --max_completion_length 1024 \ --use_vllm \ --vllm_mode colocate \ --use_peft \ --lora_target_modules "q_proj", "v_proj" \ --log_completions # For HuggingFaceTB/SmolVLM2-2.2B-Instruct pip install num2words==0.5.14 accelerate launch \ --config_file examples/accelerate_configs/deepspeed_zero3.yaml \ examples/scripts/grpo_vlm.py \ --model_name_or_path HuggingFaceTB/SmolVLM2-2.2B-Instruct \ --output_dir grpo-SmolVLM2-2.2B-Instruct \ --learning_rate 1e-5 \ --dtype bfloat16 \ --max_prompt_length 2048 \ --max_completion_length 1024 \ --use_peft \ --lora_target_modules "q_proj", "v_proj" \ --log_completions \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 2 \ --num_generations 2 """ import os import torch from datasets import load_dataset from trl import ( GRPOConfig, GRPOTrainer, ModelConfig, ScriptArguments, TrlParser, get_kbit_device_map, get_peft_config, get_quantization_config, ) from trl.rewards import accuracy_reward, think_format_reward # Enable logging in a Hugging Face Space os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") if __name__ == "__main__": parser = TrlParser((ScriptArguments, GRPOConfig, ModelConfig)) script_args, training_args, model_args = parser.parse_args_and_config() ################ # Model ################ dtype = ( model_args.dtype if model_args.dtype in ["auto", None] else getattr(torch, model_args.dtype) ) training_args.model_init_kwargs = dict( revision=model_args.model_revision, attn_implementation=model_args.attn_implementation, dtype=dtype, ) quantization_config = get_quantization_config(model_args) if quantization_config is not None: # Passing None would not be treated the same as omitting the argument, so we include it only when valid. training_args.model_init_kwargs["device_map"] = get_kbit_device_map() training_args.model_init_kwargs["quantization_config"] = quantization_config ################ # Dataset ################ dataset = load_dataset("lmms-lab/multimodal-open-r1-8k-verified", split="train") dataset = dataset.train_test_split(test_size=100, seed=42) SYSTEM_PROMPT = ( "A conversation between user and assistant. The user asks a question, and the assistant solves it. The " "assistant first thinks about the reasoning process in the mind and then provides the user with the answer. " "The reasoning process and answer are enclosed within tags, i.e., \nThis is my " "reasoning.\n\nThis is my answer." ) def make_conversation(example): prompt = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": example["problem"]}, ] return {"prompt": prompt} dataset = dataset.map(make_conversation) # Filter have big images def filter_big_images(example): image = example["image"] return image.size[0] < 512 and image.size[1] < 512 dataset = dataset.filter(filter_big_images) def convert_to_rgb(example): image = example["image"] if image.mode != "RGB": image = image.convert("RGB") example["image"] = image return example dataset = dataset.map(convert_to_rgb) train_dataset = dataset["train"] eval_dataset = dataset["test"] if training_args.eval_strategy != "no" else None ################ # Training ################ trainer = GRPOTrainer( model=model_args.model_name_or_path, args=training_args, reward_funcs=[think_format_reward, accuracy_reward], train_dataset=train_dataset, eval_dataset=eval_dataset, peft_config=get_peft_config(model_args), ) trainer.train() # Save and push to hub trainer.save_model(training_args.output_dir) if training_args.push_to_hub: trainer.push_to_hub(dataset_name=script_args.dataset_name) ``` ### requirements.txt ```text theme={null} deepspeed==0.18.2 math_verify==0.8.0 peft==0.18.0 trl==0.25.1 ``` ### deepspeed\_zero3.yaml ```yaml theme={null} --- compute_environment: LOCAL_MACHINE debug: false deepspeed_config: deepspeed_multinode_launcher: standard offload_optimizer_device: none offload_param_device: none zero3_init_flag: true zero3_save_16bit_model: true zero_stage: 3 distributed_type: DEEPSPEED downcast_bf16: no machine_rank: 0 main_training_function: main mixed_precision: bf16 num_machines: 1 num_processes: 2 rdzv_backend: static same_network: true tpu_env: [] tpu_use_cluster: false tpu_use_sudo: false use_cpu: false ```

🚀 Run this on FlexAI

Managed checkpoints mean you never lose a run to preemption. Jobs launch in under 60 seconds — no infra setup, built-in observability.

Get started → Talk to us