Add files using upload-large-folder tool

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/baselines/qwen2_5_vl_3b_clevr.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+set -x
+
+export PYTHONUNBUFFERED=1
+
+MODEL_PATH=Qwen/Qwen2.5-VL-3B-Instruct  # replace it with your local file path
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=BUAADreamer/clevr_count_70k@train \
+    data.val_files=BUAADreamer/clevr_count_70k@test \
+    data.format_prompt=./examples/format_prompt/r1v_format.jinja \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    worker.rollout.tensor_parallel_size=1 \
+    worker.reward.reward_type=sequential \
+    worker.reward.reward_function=./examples/reward_function/r1v.py:compute_score \
+    trainer.experiment_name=qwen2_5_vl_3b_clevr \
+    trainer.n_gpus_per_node=2

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/baselines/qwen2_5_vl_3b_geoqa8k.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+set -x
+
+export PYTHONUNBUFFERED=1
+
+MODEL_PATH=Qwen/Qwen2.5-VL-3B-Instruct  # replace it with your local file path
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=leonardPKU/GEOQA_8K_R1V@train \
+    data.val_files=leonardPKU/GEOQA_8K_R1V@test \
+    data.format_prompt=./examples/format_prompt/r1v_format.jinja \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    worker.rollout.tensor_parallel_size=1 \
+    worker.reward.reward_type=sequential \
+    worker.reward.reward_function=./examples/reward_function/r1v.py:compute_score \
+    trainer.experiment_name=qwen2_5_vl_3b_geoqa8k \
+    trainer.n_gpus_per_node=8

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/config.yaml

@@ -0,0 +1,96 @@
+data:
+  train_files: hiyouga/math12k@train
+  val_files: hiyouga/math12k@test
+  prompt_key: problem
+  answer_key: answer
+  image_key: image
+  max_prompt_length: 4096
+  max_response_length: 2048
+  rollout_batch_size: 256
+  val_batch_size: 1024
+  format_prompt: ./examples/format_prompt/math_format.jinja
+  override_chat_template: null
+  shuffle: true
+  seed: 1
+  max_pixels: 4194304
+  min_pixels: 262144
+  filter_overlong_prompts: true
+
+algorithm:
+  adv_estimator: grpo
+  disable_kl: false
+  use_kl_loss: true
+  kl_penalty: low_var_kl
+  kl_coef: 1.0e-2
+
+worker:
+  actor:
+    global_batch_size: 128
+    micro_batch_size_per_device_for_update: 4
+    micro_batch_size_per_device_for_experience: 8
+    max_grad_norm: 1.0
+    padding_free: true
+    ulysses_sequence_parallel_size: 1
+    is_entropy_reg: False
+    entropy_coef: 0.001
+    model:
+      model_path: Qwen/Qwen2.5-7B-Instruct
+      enable_gradient_checkpointing: true
+      trust_remote_code: true
+      freeze_vision_tower: false
+    optim:
+      lr: 1.0e-6
+      weight_decay: 1.0e-2
+      strategy: adamw  # {adamw, adamw_bf16}
+      lr_warmup_ratio: 0.0
+    fsdp:
+      enable_full_shard: true
+      enable_cpu_offload: false
+      enable_rank0_init: true
+    offload:
+      offload_params: true  # true: more CPU memory; false: more GPU memory
+      offload_optimizer: true  # true: more CPU memory; false: more GPU memory
+
+  rollout:
+    n: 8
+    temperature: 1.0
+    top_p: 0.99
+    gpu_memory_utilization: 0.6
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    enforce_eager: false
+    enable_chunked_prefill: false
+    tensor_parallel_size: 2
+    limit_images: 10
+    val_override_config:
+      temperature: 0.5
+      n: 1
+
+  ref:
+    fsdp:
+      enable_full_shard: true
+      enable_cpu_offload: true  # true: more CPU memory; false: more GPU memory
+      enable_rank0_init: true
+    offload:
+      offload_params: false
+
+  reward:
+    reward_type: batch
+    reward_function: ./examples/reward_function/math.py:compute_score
+
+trainer:
+  total_epochs: 15
+  max_steps: null
+  project_name: easy_r1
+  experiment_name: qwen2_5_7b_math_grpo
+  logger: ["console", "wandb"]
+  nnodes: 1
+  n_gpus_per_node: 8
+  val_freq: -1  # -1 to disable
+  val_before_train: false
+  val_only: false
+  val_generations_to_log: 3
+  save_freq: 10  # -1 to disable
+  save_limit: -1  # -1 to disable
+  save_checkpoint_path: null
+  load_checkpoint_path: null

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/format_prompt/caption_format.jinja

@@ -0,0 +1 @@
+{{ content | trim }} You FIRST describe the given image in a detailed manner, THEN think about the reasoning process as an internal monologue based on the generated description and given question and then provide the final answer. The description MUST BE enclosed within <caption> </caption> tags. The reasoning process MUST BE enclosed within <think> </think> tags. The final answer MUST BE put in \boxed{}.

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/format_prompt/math_format.jinja

@@ -0,0 +1 @@
+{{ content | trim }} You FIRST think about the reasoning process as an internal monologue and then provide the final answer. The reasoning process MUST BE enclosed within <think> </think> tags. The final answer MUST BE put in \boxed{}.

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/format_prompt/r1v_format.jinja

@@ -0,0 +1 @@
+{{ content | trim }} 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 <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/llama3.2v_critic_grpo.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+set -x
+
+export PYTHONUNBUFFERED=1
+
+MODEL_PATH=meta-llama/Llama-3.2-11B-Vision-Instruct # replace it with your local file path
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.val_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.shuffle=True \
+    data.format_prompt=./examples/format_prompt/math_format.jinja \
+    data.rollout_batch_size=256 \
+    data.max_prompt_length=4096 \
+    data.prompt_key=question \
+    worker.actor.is_entropy_reg=False \
+    worker.actor.global_batch_size=128 \
+    worker.actor.optim.lr=5.0e-7 \
+    worker.actor.micro_batch_size_per_device_for_update=4 \
+    worker.actor.micro_batch_size_per_device_for_experience=8 \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    worker.rollout.tensor_parallel_size=8 \
+    worker.rollout.limit_images=1 \
+    worker.rollout.gpu_memory_utilization=0.6 \
+    worker.reward.reward_function=./examples/reward_function/math.py:compute_score \
+    trainer.experiment_name=llama3.2v_critic_grpo \
+    trainer.n_gpus_per_node=8 

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/mimo_vl_7b_llava_critic_grpo.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+MODEL_PATH=XiaomiMiMo/MiMo-VL-7B-SFT-2508
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.val_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.filter_overlong_prompts=True \
+    data.prompt_key=question \
+    data.max_prompt_length=6142 \
+    worker.actor.is_entropy_reg=False \
+    worker.actor.global_batch_size=128 \
+    worker.actor.micro_batch_size_per_device_for_update=4 \
+    worker.actor.micro_batch_size_per_device_for_experience=8 \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    trainer.experiment_name=mimo_vl_7b_llava_critic_grpo \
+    trainer.n_gpus_per_node=8 
+

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/mimo_vl_rl_7b_llava_critic_grpo.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+MODEL_PATH=XiaomiMiMo/MiMo-VL-7B-RL-2508
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.val_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.filter_overlong_prompts=True \
+    data.prompt_key=question \
+    data.max_prompt_length=6142 \
+    worker.actor.is_entropy_reg=False \
+    worker.actor.global_batch_size=128 \
+    worker.actor.micro_batch_size_per_device_for_update=4 \
+    worker.actor.micro_batch_size_per_device_for_experience=8 \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    trainer.experiment_name=mimo_vl_rl_7b_llava_critic_grpo \
+    trainer.n_gpus_per_node=8 
+

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/qwen2_5_vl_7b_llava_critic_grpo.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+MODEL_PATH=Qwen/Qwen2.5-VL-7B-Instruct
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.val_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.filter_overlong_prompts=True \
+    data.prompt_key=question \
+    data.max_prompt_length=6142 \
+    worker.actor.is_entropy_reg=False \
+    worker.actor.global_batch_size=128 \
+    worker.actor.micro_batch_size_per_device_for_update=4 \
+    worker.actor.micro_batch_size_per_device_for_experience=8 \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    trainer.experiment_name=qwen2_5_vl_7b_llava_critic_grpo \
+    trainer.n_gpus_per_node=8 
+

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/reward_function/caption.py

@@ -0,0 +1,46 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import re
+from typing import Dict, List
+
+from mathruler.grader import extract_boxed_content, grade_answer
+
+
+def format_reward(predict: str) -> float:
+    pattern = re.compile(r"<caption>.*</caption>.*<think>.*</think>.*\\boxed\{.*\}.*", re.DOTALL)
+    format_match = re.fullmatch(pattern, predict)
+    return 1.0 if format_match else 0.0
+
+
+def accuracy_reward(predict: str, ground_truth: str) -> float:
+    answer = extract_boxed_content(predict)
+    return 1.0 if grade_answer(answer, ground_truth) else 0.0
+
+
+def compute_score(predicts: List[str], ground_truths: List[str], format_weight: float = 0.1) -> List[Dict[str, float]]:
+    scores = []
+    for predict, ground_truth in zip(predicts, ground_truths):
+        predict = re.sub(r"\s*(<|>|/)\s*", r"\1", predict)  # handle qwen2.5vl-32b format
+        format_score = format_reward(predict)
+        accuracy_score = accuracy_reward(predict, ground_truth)
+        scores.append(
+            {
+                "overall": (1 - format_weight) * accuracy_score + format_weight * format_score,
+                "format": format_score,
+                "accuracy": accuracy_score,
+            }
+        )
+
+    return scores

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/reward_function/math.py

@@ -0,0 +1,46 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import re
+from typing import Dict, List
+
+from mathruler.grader import extract_boxed_content, grade_answer
+
+
+def format_reward(predict: str) -> float:
+    pattern = re.compile(r"<think>.*</think>.*\\boxed\{.*\}.*", re.DOTALL)
+    format_match = re.fullmatch(pattern, predict)
+    return 1.0 if format_match else 0.0
+
+
+def accuracy_reward(predict: str, ground_truth: str) -> float:
+    answer = extract_boxed_content(predict)
+    return 1.0 if grade_answer(answer, ground_truth) else 0.0
+
+
+def compute_score(predicts: List[str], ground_truths: List[str], format_weight: float = 0.1) -> List[Dict[str, float]]:
+    scores = []
+    for predict, ground_truth in zip(predicts, ground_truths):
+        predict = re.sub(r"\s*(<|>|/)\s*", r"\1", predict)  # handle qwen2.5vl-32b format
+        format_score = format_reward(predict)
+        accuracy_score = accuracy_reward(predict, ground_truth)
+        scores.append(
+            {
+                "overall": (1 - format_weight) * accuracy_score + format_weight * format_score,
+                "format": format_score,
+                "accuracy": accuracy_score,
+            }
+        )
+
+    return scores

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/reward_function/r1v.py

@@ -0,0 +1,47 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import re
+from typing import Dict
+
+from mathruler.grader import grade_answer
+
+
+def format_reward(predict: str) -> float:
+    pattern = re.compile(r"<think>.*?</think>\s*<answer>.*?</answer>", re.DOTALL)
+    format_match = re.fullmatch(pattern, predict)
+    return 1.0 if format_match else 0.0
+
+
+def accuracy_reward(predict: str, ground_truth: str) -> float:
+    try:
+        content_match = re.search(r"<answer>(.*?)</answer>", predict)
+        given_answer = content_match.group(1).strip() if content_match else predict.strip()
+        if grade_answer(given_answer, ground_truth.strip()):
+            return 1.0
+
+    except Exception:
+        pass
+
+    return 0.0
+
+
+def compute_score(predict: str, ground_truth: str, format_weight: float = 0.5) -> Dict[str, float]:
+    format_score = format_reward(predict)
+    accuracy_score = accuracy_reward(predict, ground_truth)
+    return {
+        "overall": (1 - format_weight) * accuracy_score + format_weight * format_score,
+        "format": format_score,
+        "accuracy": accuracy_score,
+    }

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/runtime_env.yaml

@@ -0,0 +1,9 @@
+working_dir: ./
+excludes: ["/.git/"]
+env_vars:
+  TOKENIZERS_PARALLELISM: "true"
+  NCCL_DEBUG: "WARN"
+  VLLM_LOGGING_LEVEL: "WARN"
+  TORCH_NCCL_AVOID_RECORD_STREAMS: "1"
+  PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:False"
+  PYTHONUNBUFFERED: "1"

LLaVA-NeXT/llava-critic-r1/EasyR1/examples/thinklite_vl_7b_llava_critic_grpo.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+MODEL_PATH=russwang/ThinkLite-VL-7B
+
+python3 -m verl.trainer.main \
+    config=examples/config.yaml \
+    data.train_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.val_files=lmms-lab/LLaVA-Critic-GRPO-dataset@train \
+    data.filter_overlong_prompts=True \
+    data.prompt_key=question \
+    data.max_prompt_length=6142 \
+    worker.actor.is_entropy_reg=False \
+    worker.actor.global_batch_size=128 \
+    worker.actor.optim.lr=5.0e-7 \
+    worker.actor.micro_batch_size_per_device_for_update=4 \
+    worker.actor.micro_batch_size_per_device_for_experience=8 \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    trainer.experiment_name=thinklite_vl_7b_llava_critic_grpo \
+    trainer.n_gpus_per_node=8 

LLaVA-NeXT/llava-critic-r1/EasyR1/verl.egg-info/PKG-INFO

@@ -0,0 +1,262 @@
+Metadata-Version: 2.4
+Name: verl
+Version: 0.3.1.dev0
+Summary: An Efficient, Scalable, Multi-Modality RL Training Framework based on veRL
+Home-page: https://github.com/volcengine/verl
+Author: verl
+Author-email: zhangchi.usc1992@bytedance.com, gmsheng@connect.hku.hk, hiyouga@buaa.edu.cn
+License: Apache 2.0 License
+Requires-Python: >=3.9.0
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: accelerate
+Requires-Dist: codetiming
+Requires-Dist: datasets
+Requires-Dist: liger-kernel
+Requires-Dist: mathruler
+Requires-Dist: numpy
+Requires-Dist: omegaconf
+Requires-Dist: pandas
+Requires-Dist: peft
+Requires-Dist: pillow
+Requires-Dist: pyarrow>=15.0.0
+Requires-Dist: pylatexenc
+Requires-Dist: qwen-vl-utils
+Requires-Dist: ray[default]
+Requires-Dist: tensordict
+Requires-Dist: torchdata
+Requires-Dist: transformers>=4.51.0
+Requires-Dist: vllm>=0.7.3
+Requires-Dist: wandb
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# EasyR1: An Efficient, Scalable, Multi-Modality RL Training Framework
+
+[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/EasyR1)](https://github.com/hiyouga/EasyR1/stargazers)
+[![Twitter](https://img.shields.io/twitter/follow/llamafactory_ai)](https://twitter.com/llamafactory_ai)
+
+This project is a clean fork of the original [veRL](https://github.com/volcengine/verl) project to support vision language models, we thank all the authors for providing such a high-performance RL training framework.
+
+EasyR1 is efficient and scalable due to the design of **[HybirdEngine](https://arxiv.org/abs/2409.19256)** and the latest release of **[vLLM](https://github.com/vllm-project/vllm)**'s SPMD mode.
+
+## Features
+
+- Supported models
+  - Llama3/Qwen2/Qwen2.5/Qwen3 language models
+  - Qwen2/Qwen2.5-VL vision language models
+  - DeepSeek-R1 distill models
+
+- Supported algorithms
+  - GRPO
+  - Reinforce++
+  - ReMax
+  - RLOO
+
+- Supported datasets
+  - Any text, vision-text dataset in a [specific format](#custom-dataset)
+
+- Supported tricks
+  - Padding-free training
+  - Resuming from checkpoint
+  - Wandb & SwanLab & Mlflow & Tensorboard tracking
+
+## Requirements
+
+### Software Requirements
+
+- Python 3.9+
+- transformers>=4.51.0
+- flash-attn>=2.4.3
+- vllm>=0.8.3
+
+We provide a [Dockerfile](./Dockerfile) to easily build environments.
+
+We recommend using the [pre-built docker image](https://hub.docker.com/r/hiyouga/verl) in EasyR1.
+
+```bash
+docker pull hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.4-flashinfer0.2.2-cxx11abi0
+```
+
+### Hardware Requirements
+
+\* *estimated*
+
+| Method                   | Bits |  1.5B  |   3B   |   7B   |   32B   |
+| ------------------------ | ---- | ------ | ------ | ------ | ------- |
+| GRPO Full Fine-Tuning    |  AMP | 2*24GB | 4*40GB | 8*40GB | 16*80GB |
+| GRPO Full Fine-Tuning    | BF16 | 1*24GB | 1*40GB | 4*40GB |  8*80GB |
+
+> [!NOTE]
+> Use `worker.actor.fsdp.torch_dtype=bf16` and `worker.actor.optim.strategy=adamw_bf16` to enable bf16 training.
+>
+> We are working hard to reduce the VRAM in RL training, LoRA support will be integrated in next updates.
+
+## Tutorial: Run Qwen2.5-VL GRPO on [Geometry3K](https://huggingface.co/datasets/hiyouga/geometry3k) Dataset in Just 3 Steps
+
+![image](assets/qwen2_5_vl_7b_geo.png)
+
+### Installation
+
+```bash
+git clone https://github.com/hiyouga/EasyR1.git
+cd EasyR1
+pip install -e .
+```
+
+### GRPO Training
+
+```bash
+bash examples/qwen2_5_vl_7b_geo3k_grpo.sh
+```
+
+### Merge Checkpoint in Hugging Face Format
+
+```bash
+python3 scripts/model_merger.py --local_dir checkpoints/easy_r1/exp_name/global_step_1/actor
+```
+
+> [!TIP]
+> If you encounter issues with connecting to Hugging Face, consider using `export HF_ENDPOINT=https://hf-mirror.com`.
+>
+> If you want to use SwanLab logger, consider using `bash examples/qwen2_5_vl_7b_geo3k_swanlab.sh`.
+
+## Custom Dataset
+
+Please refer to the example datasets to prepare your own dataset.
+
+- Text dataset: https://huggingface.co/datasets/hiyouga/math12k
+- Image-text dataset: https://huggingface.co/datasets/hiyouga/geometry3k
+- Multi-image-text dataset: https://huggingface.co/datasets/hiyouga/journeybench-multi-image-vqa
+
+## How to Understand GRPO in EasyR1
+
+![image](assets/easyr1_grpo.png)
+
+- To learn about the GRPO algorithm, you can refer to [Hugging Face's blog](https://huggingface.co/docs/trl/v0.16.1/en/grpo_trainer).
+
+## How to Run 70B+ Model in Multi-node Environment
+
+1. Start the Ray head node.
+
+```bash
+ray start --head --port=6379 --dashboard-host=0.0.0.0
+```
+
+2. Start the Ray worker node and connect to the head node.
+
+```bash
+ray start --address=<head_node_ip>:6379
+```
+
+3. Check the Ray resource pool.
+
+```bash
+ray status
+```
+
+4. Run training script on the Ray head node only.
+
+```bash
+bash examples/qwen2_5_vl_7b_geo3k_grpo.sh
+```
+
+See the **[veRL's official doc](https://verl.readthedocs.io/en/latest/start/multinode.html)** for more details about multi-node training and Ray debugger.
+
+## Other Baselines
+
+We also reproduced the following two baselines of the [R1-V](https://github.com/deep-agent/R1-V) project.
+- [CLEVR-70k-Counting](examples/baselines/qwen2_5_vl_3b_clevr.sh): Train the Qwen2.5-VL-3B-Instruct model on counting problem.
+- [GeoQA-8k](examples/baselines/qwen2_5_vl_3b_geoqa8k.sh): Train the Qwen2.5-VL-3B-Instruct model on GeoQA problem.
+
+## Performance Baselines
+
+See [baselines.md](assets/baselines.md).
+
+## Awesome Work using EasyR1
+
+- **MMR1**: Advancing the Frontiers of Multimodal Reasoning. [![[code]](https://img.shields.io/github/stars/LengSicong/MMR1)](https://github.com/LengSicong/MMR1)
+- **Vision-R1**: Incentivizing Reasoning Capability in Multimodal Large Language Models. [![[code]](https://img.shields.io/github/stars/Osilly/Vision-R1)](https://github.com/Osilly/Vision-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.06749-blue)](https://arxiv.org/abs/2503.06749)
+- **Seg-Zero**: Reasoning-Chain Guided Segmentation via Cognitive Reinforcement. [![[code]](https://img.shields.io/github/stars/dvlab-research/Seg-Zero)](https://github.com/dvlab-research/Seg-Zero) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.06520-blue)](https://arxiv.org/abs/2503.06520)
+- **MetaSpatial**: Reinforcing 3D Spatial Reasoning in VLMs for the Metaverse. [![[code]](https://img.shields.io/github/stars/PzySeere/MetaSpatial)](https://github.com/PzySeere/MetaSpatial) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.18470-blue)](https://arxiv.org/abs/2503.18470)
+- **Temporal-R1**: Envolving Temporal Reasoning Capability into LMMs via Temporal Consistent Reward. [![[code]](https://img.shields.io/github/stars/appletea233/Temporal-R1)](https://github.com/appletea233/Temporal-R1)
+- **NoisyRollout**: Reinforcing Visual Reasoning with Data Augmentation. [![[code]](https://img.shields.io/github/stars/John-AI-Lab/NoisyRollout)](https://github.com/John-AI-Lab/NoisyRollout) [![[arxiv]](https://img.shields.io/badge/arxiv-2504.13055-blue)](https://arxiv.org/pdf/2504.13055)
+- **GUI-R1**: A Generalist R1-Style Vision-Language Action Model For GUI Agents. [![[code]](https://img.shields.io/github/stars/ritzz-ai/GUI-R1)](https://github.com/ritzz-ai/GUI-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2504.10458-blue)](https://arxiv.org/abs/2504.10458)
+- **R1-Track**: Direct Application of MLLMs to Visual Object Tracking via Reinforcement Learning. [![[code]](https://img.shields.io/github/stars/Wangbiao2/R1-Track)](https://github.com/Wangbiao2/R1-Track)
+- **VisionReasoner**: Unified Visual Perception and Reasoning via Reinforcement Learning. [![[code]](https://img.shields.io/github/stars/dvlab-research/VisionReasoner)](https://github.com/dvlab-research/VisionReasoner) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.12081-blue)](https://arxiv.org/abs/2505.12081)
+- **MM-UPT**: Unsupervised Post-Training for Multi-Modal LLM Reasoning via GRPO. [![[code]](https://img.shields.io/github/stars/waltonfuture/MM-UPT)](https://github.com/waltonfuture/MM-UPT) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.22453-blue)](https://arxiv.org/pdf/2505.22453)
+- **RL-with-Cold-Start**: Advancing Multimodal Reasoning via Reinforcement Learning with Cold Start. [![[code]](https://img.shields.io/github/stars/waltonfuture/RL-with-Cold-Start)](https://github.com/waltonfuture/RL-with-Cold-Start) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.22334-blue)](https://arxiv.org/pdf/2505.22334)
+- **ViGoRL**: Grounded Reinforcement Learning for Visual Reasoning. [![[code]](https://img.shields.io/github/stars/Gabesarch/grounded-rl)](https://github.com/Gabesarch/grounded-rl) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.22334-blue)](https://arxiv.org/abs/2505.23678)
+- **Revisual-R1**: Advancing Multimodal Reasoning: From Optimized Cold Start to Staged Reinforcement Learning. [![[code]](https://img.shields.io/github/stars/CSfufu/Revisual-R1)](https://github.com/CSfufu/Revisual-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2506.04207-blue)](https://arxiv.org/abs/2506.04207)
+- **SophiaVL-R1**: Reinforcing MLLMs Reasoning with Thinking Reward. [![[code]](https://img.shields.io/github/stars/kxfan2002/SophiaVL-R1)](https://github.com/kxfan2002/SophiaVL-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.17018-blue)](https://arxiv.org/abs/2505.17018)
+  
+## TODO
+
+- Support LoRA (high priority).
+- Support ulysses parallelism for VLMs (middle priority).
+- Support more VLM architectures.
+
+> [!NOTE]
+> We will not provide scripts for supervised fine-tuning and inference in this project. If you have such requirements, we recommend using [LLaMA-Factory](https://github.com/hiyouga/LLaMA-Factory).
+
+### Known bugs
+
+These features are temporarily disabled for now, we plan to fix them one-by-one in the future updates.
+
+- Vision language models are not compatible with ulysses parallelism yet.
+
+## Discussion Group
+
+👋 Join our [WeChat group](assets/wechat.jpg).
+
+## FAQs
+
+> ValueError: Image features and image tokens do not match: tokens: 8192, features 9800
+
+Increase the `data.max_prompt_length` or reduce the `data.max_pixels`.
+
+> RuntimeError: CUDA Error: out of memory at /workspace/csrc/cumem_allocator.cpp:62
+
+Reduce the `worker.rollout.gpu_memory_utilization` and enable `worker.actor.offload.offload_params`.
+
+> RuntimeError: 0 active drivers ([]). There should only be one.
+
+Uninstall `deepspeed` from the current python environment.
+
+## Citation
+
+Core contributors: [Yaowei Zheng](https://github.com/hiyouga), [Junting Lu](https://github.com/AL-377), [Shenzhi Wang](https://github.com/Shenzhi-Wang), [Zhangchi Feng](https://github.com/BUAADreamer), [Dongdong Kuang](https://github.com/Kuangdd01) and Yuwen Xiong
+
+We also thank Guangming Sheng and Chi Zhang for helpful discussions.
+
+```bibtex
+@misc{zheng2025easyr1,
+  title        = {EasyR1: An Efficient, Scalable, Multi-Modality RL Training Framework},
+  author       = {Yaowei Zheng, Junting Lu, Shenzhi Wang, Zhangchi Feng, Dongdong Kuang, Yuwen Xiong},
+  howpublished = {\url{https://github.com/hiyouga/EasyR1}},
+  year         = {2025}
+}
+```
+
+We recommend to also cite the original work.
+
+```bibtex
+@article{sheng2024hybridflow,
+  title   = {HybridFlow: A Flexible and Efficient RLHF Framework},
+  author  = {Guangming Sheng and Chi Zhang and Zilingfeng Ye and Xibin Wu and Wang Zhang and Ru Zhang and Yanghua Peng and Haibin Lin and Chuan Wu},
+  year    = {2024},
+  journal = {arXiv preprint arXiv: 2409.19256}
+}
+```

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/Dockerfile

@@ -0,0 +1,76 @@
+# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10)
+# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html
+FROM nvcr.io/nvidia/pytorch:24.08-py3
+
+# Define environments
+ENV MAX_JOBS=32
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+ENV DEBIAN_FRONTEND=noninteractive
+ENV NODE_OPTIONS=""
+ENV PIP_ROOT_USER_ACTION=ignore
+ENV HF_HUB_ENABLE_HF_TRANSFER="1"
+
+# Define installation arguments
+ARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/
+ARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
+
+# Set apt source
+RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \
+    { \
+    echo "deb ${APT_SOURCE} jammy main restricted universe multiverse"; \
+    echo "deb ${APT_SOURCE} jammy-updates main restricted universe multiverse"; \
+    echo "deb ${APT_SOURCE} jammy-backports main restricted universe multiverse"; \
+    echo "deb ${APT_SOURCE} jammy-security main restricted universe multiverse"; \
+    } > /etc/apt/sources.list
+
+# Install systemctl
+RUN apt-get update && \
+    apt-get install -y -o Dpkg::Options::="--force-confdef" systemd && \
+    apt-get clean
+
+# Install tini
+RUN apt-get update && \
+    apt-get install -y tini && \
+    apt-get clean
+
+# Change pip source
+RUN pip config set global.index-url "${PIP_INDEX}" && \
+    pip config set global.extra-index-url "${PIP_INDEX}" && \
+    python -m pip install --upgrade pip
+
+# Uninstall nv-pytorch fork
+RUN pip uninstall -y torch torchvision torchaudio \
+    pytorch-quantization pytorch-triton torch-tensorrt \
+    transformer_engine flash_attn apex megatron-core \
+    xgboost opencv grpcio
+
+# Fix cv2
+RUN rm -rf /usr/local/lib/python3.10/dist-packages/cv2
+
+# Install torch-2.6.0+cu124 + vllm-0.8.4
+# torch-2.6.0+cu124: cxx11abi=False
+# torch-2.6.0+cu126: cxx11abi=True
+# see https://github.com/flashinfer-ai/flashinfer/issues/911
+RUN pip install --no-cache-dir "vllm==0.8.4" "torch==2.6.0" "torchvision==0.21.0" "torchaudio==2.6.0" tensordict torchdata \
+    "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \
+    "numpy<2.0.0" "pyarrow>=15.0.0" pandas \
+    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb liger-kernel mathruler \
+    pytest yapf py-spy pyext pre-commit ruff
+
+# Install flash-attn-2.7.4.post1 (cxx11abi=False)
+RUN wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl && \
+    pip install --no-cache-dir flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
+
+# Install flashinfer-0.2.2.post1+cu124 (cxx11abi=False)
+# vllm-0.8.3 does not support flashinfer>=0.2.3
+# see https://github.com/vllm-project/vllm/pull/15777
+RUN wget -nv https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \
+    pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl
+
+# Fix packages
+RUN pip uninstall -y pynvml nvidia-ml-py && \
+    pip install --no-cache-dir --upgrade "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1"
+
+# Reset pip config
+RUN pip config unset global.index-url && \
+    pip config unset global.extra-index-url

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/Makefile

@@ -0,0 +1,18 @@
+.PHONY: build commit quality style
+
+check_dirs := examples scripts verl setup.py
+
+build:
+	python3 setup.py sdist bdist_wheel
+
+commit:
+	pre-commit install
+	pre-commit run --all-files
+
+quality:
+	ruff check $(check_dirs)
+	ruff format --check $(check_dirs)
+
+style:
+	ruff check $(check_dirs) --fix
+	ruff format $(check_dirs)

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/README.md

@@ -0,0 +1,217 @@
+# EasyR1: An Efficient, Scalable, Multi-Modality RL Training Framework
+
+[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/EasyR1)](https://github.com/hiyouga/EasyR1/stargazers)
+[![Twitter](https://img.shields.io/twitter/follow/llamafactory_ai)](https://twitter.com/llamafactory_ai)
+
+This project is a clean fork of the original [veRL](https://github.com/volcengine/verl) project to support vision language models, we thank all the authors for providing such a high-performance RL training framework.
+
+EasyR1 is efficient and scalable due to the design of **[HybirdEngine](https://arxiv.org/abs/2409.19256)** and the latest release of **[vLLM](https://github.com/vllm-project/vllm)**'s SPMD mode.
+
+## Features
+
+- Supported models
+  - Llama3/Qwen2/Qwen2.5/Qwen3 language models
+  - Qwen2/Qwen2.5-VL vision language models
+  - DeepSeek-R1 distill models
+
+- Supported algorithms
+  - GRPO
+  - Reinforce++
+  - ReMax
+  - RLOO
+
+- Supported datasets
+  - Any text, vision-text dataset in a [specific format](#custom-dataset)
+
+- Supported tricks
+  - Padding-free training
+  - Resuming from checkpoint
+  - Wandb & SwanLab & Mlflow & Tensorboard tracking
+
+## Requirements
+
+### Software Requirements
+
+- Python 3.9+
+- transformers>=4.51.0
+- flash-attn>=2.4.3
+- vllm>=0.8.3
+
+We provide a [Dockerfile](./Dockerfile) to easily build environments.
+
+We recommend using the [pre-built docker image](https://hub.docker.com/r/hiyouga/verl) in EasyR1.
+
+```bash
+docker pull hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.4-flashinfer0.2.2-cxx11abi0
+```
+
+### Hardware Requirements
+
+\* *estimated*
+
+| Method                   | Bits |  1.5B  |   3B   |   7B   |   32B   |
+| ------------------------ | ---- | ------ | ------ | ------ | ------- |
+| GRPO Full Fine-Tuning    |  AMP | 2*24GB | 4*40GB | 8*40GB | 16*80GB |
+| GRPO Full Fine-Tuning    | BF16 | 1*24GB | 1*40GB | 4*40GB |  8*80GB |
+
+> [!NOTE]
+> Use `worker.actor.fsdp.torch_dtype=bf16` and `worker.actor.optim.strategy=adamw_bf16` to enable bf16 training.
+>
+> We are working hard to reduce the VRAM in RL training, LoRA support will be integrated in next updates.
+
+## Tutorial: Run Qwen2.5-VL GRPO on [Geometry3K](https://huggingface.co/datasets/hiyouga/geometry3k) Dataset in Just 3 Steps
+
+![image](assets/qwen2_5_vl_7b_geo.png)
+
+### Installation
+
+```bash
+git clone https://github.com/hiyouga/EasyR1.git
+cd EasyR1
+pip install -e .
+```
+
+### GRPO Training
+
+```bash
+bash examples/qwen2_5_vl_7b_geo3k_grpo.sh
+```
+
+### Merge Checkpoint in Hugging Face Format
+
+```bash
+python3 scripts/model_merger.py --local_dir checkpoints/easy_r1/exp_name/global_step_1/actor
+```
+
+> [!TIP]
+> If you encounter issues with connecting to Hugging Face, consider using `export HF_ENDPOINT=https://hf-mirror.com`.
+>
+> If you want to use SwanLab logger, consider using `bash examples/qwen2_5_vl_7b_geo3k_swanlab.sh`.
+
+## Custom Dataset
+
+Please refer to the example datasets to prepare your own dataset.
+
+- Text dataset: https://huggingface.co/datasets/hiyouga/math12k
+- Image-text dataset: https://huggingface.co/datasets/hiyouga/geometry3k
+- Multi-image-text dataset: https://huggingface.co/datasets/hiyouga/journeybench-multi-image-vqa
+
+## How to Understand GRPO in EasyR1
+
+![image](assets/easyr1_grpo.png)
+
+- To learn about the GRPO algorithm, you can refer to [Hugging Face's blog](https://huggingface.co/docs/trl/v0.16.1/en/grpo_trainer).
+
+## How to Run 70B+ Model in Multi-node Environment
+
+1. Start the Ray head node.
+
+```bash
+ray start --head --port=6379 --dashboard-host=0.0.0.0
+```
+
+2. Start the Ray worker node and connect to the head node.
+
+```bash
+ray start --address=<head_node_ip>:6379
+```
+
+3. Check the Ray resource pool.
+
+```bash
+ray status
+```
+
+4. Run training script on the Ray head node only.
+
+```bash
+bash examples/qwen2_5_vl_7b_geo3k_grpo.sh
+```
+
+See the **[veRL's official doc](https://verl.readthedocs.io/en/latest/start/multinode.html)** for more details about multi-node training and Ray debugger.
+
+## Other Baselines
+
+We also reproduced the following two baselines of the [R1-V](https://github.com/deep-agent/R1-V) project.
+- [CLEVR-70k-Counting](examples/baselines/qwen2_5_vl_3b_clevr.sh): Train the Qwen2.5-VL-3B-Instruct model on counting problem.
+- [GeoQA-8k](examples/baselines/qwen2_5_vl_3b_geoqa8k.sh): Train the Qwen2.5-VL-3B-Instruct model on GeoQA problem.
+
+## Performance Baselines
+
+See [baselines.md](assets/baselines.md).
+
+## Awesome Work using EasyR1
+
+- **MMR1**: Advancing the Frontiers of Multimodal Reasoning. [![[code]](https://img.shields.io/github/stars/LengSicong/MMR1)](https://github.com/LengSicong/MMR1)
+- **Vision-R1**: Incentivizing Reasoning Capability in Multimodal Large Language Models. [![[code]](https://img.shields.io/github/stars/Osilly/Vision-R1)](https://github.com/Osilly/Vision-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.06749-blue)](https://arxiv.org/abs/2503.06749)
+- **Seg-Zero**: Reasoning-Chain Guided Segmentation via Cognitive Reinforcement. [![[code]](https://img.shields.io/github/stars/dvlab-research/Seg-Zero)](https://github.com/dvlab-research/Seg-Zero) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.06520-blue)](https://arxiv.org/abs/2503.06520)
+- **MetaSpatial**: Reinforcing 3D Spatial Reasoning in VLMs for the Metaverse. [![[code]](https://img.shields.io/github/stars/PzySeere/MetaSpatial)](https://github.com/PzySeere/MetaSpatial) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.18470-blue)](https://arxiv.org/abs/2503.18470)
+- **Temporal-R1**: Envolving Temporal Reasoning Capability into LMMs via Temporal Consistent Reward. [![[code]](https://img.shields.io/github/stars/appletea233/Temporal-R1)](https://github.com/appletea233/Temporal-R1)
+- **NoisyRollout**: Reinforcing Visual Reasoning with Data Augmentation. [![[code]](https://img.shields.io/github/stars/John-AI-Lab/NoisyRollout)](https://github.com/John-AI-Lab/NoisyRollout) [![[arxiv]](https://img.shields.io/badge/arxiv-2504.13055-blue)](https://arxiv.org/pdf/2504.13055)
+- **GUI-R1**: A Generalist R1-Style Vision-Language Action Model For GUI Agents. [![[code]](https://img.shields.io/github/stars/ritzz-ai/GUI-R1)](https://github.com/ritzz-ai/GUI-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2504.10458-blue)](https://arxiv.org/abs/2504.10458)
+- **R1-Track**: Direct Application of MLLMs to Visual Object Tracking via Reinforcement Learning. [![[code]](https://img.shields.io/github/stars/Wangbiao2/R1-Track)](https://github.com/Wangbiao2/R1-Track)
+- **VisionReasoner**: Unified Visual Perception and Reasoning via Reinforcement Learning. [![[code]](https://img.shields.io/github/stars/dvlab-research/VisionReasoner)](https://github.com/dvlab-research/VisionReasoner) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.12081-blue)](https://arxiv.org/abs/2505.12081)
+- **MM-UPT**: Unsupervised Post-Training for Multi-Modal LLM Reasoning via GRPO. [![[code]](https://img.shields.io/github/stars/waltonfuture/MM-UPT)](https://github.com/waltonfuture/MM-UPT) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.22453-blue)](https://arxiv.org/pdf/2505.22453)
+- **RL-with-Cold-Start**: Advancing Multimodal Reasoning via Reinforcement Learning with Cold Start. [![[code]](https://img.shields.io/github/stars/waltonfuture/RL-with-Cold-Start)](https://github.com/waltonfuture/RL-with-Cold-Start) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.22334-blue)](https://arxiv.org/pdf/2505.22334)
+- **ViGoRL**: Grounded Reinforcement Learning for Visual Reasoning. [![[code]](https://img.shields.io/github/stars/Gabesarch/grounded-rl)](https://github.com/Gabesarch/grounded-rl) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.22334-blue)](https://arxiv.org/abs/2505.23678)
+- **Revisual-R1**: Advancing Multimodal Reasoning: From Optimized Cold Start to Staged Reinforcement Learning. [![[code]](https://img.shields.io/github/stars/CSfufu/Revisual-R1)](https://github.com/CSfufu/Revisual-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2506.04207-blue)](https://arxiv.org/abs/2506.04207)
+- **SophiaVL-R1**: Reinforcing MLLMs Reasoning with Thinking Reward. [![[code]](https://img.shields.io/github/stars/kxfan2002/SophiaVL-R1)](https://github.com/kxfan2002/SophiaVL-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2505.17018-blue)](https://arxiv.org/abs/2505.17018)
+  
+## TODO
+
+- Support LoRA (high priority).
+- Support ulysses parallelism for VLMs (middle priority).
+- Support more VLM architectures.
+
+> [!NOTE]
+> We will not provide scripts for supervised fine-tuning and inference in this project. If you have such requirements, we recommend using [LLaMA-Factory](https://github.com/hiyouga/LLaMA-Factory).
+
+### Known bugs
+
+These features are temporarily disabled for now, we plan to fix them one-by-one in the future updates.
+
+- Vision language models are not compatible with ulysses parallelism yet.
+
+## Discussion Group
+
+👋 Join our [WeChat group](assets/wechat.jpg).
+
+## FAQs
+
+> ValueError: Image features and image tokens do not match: tokens: 8192, features 9800
+
+Increase the `data.max_prompt_length` or reduce the `data.max_pixels`.
+
+> RuntimeError: CUDA Error: out of memory at /workspace/csrc/cumem_allocator.cpp:62
+
+Reduce the `worker.rollout.gpu_memory_utilization` and enable `worker.actor.offload.offload_params`.
+
+> RuntimeError: 0 active drivers ([]). There should only be one.
+
+Uninstall `deepspeed` from the current python environment.
+
+## Citation
+
+Core contributors: [Yaowei Zheng](https://github.com/hiyouga), [Junting Lu](https://github.com/AL-377), [Shenzhi Wang](https://github.com/Shenzhi-Wang), [Zhangchi Feng](https://github.com/BUAADreamer), [Dongdong Kuang](https://github.com/Kuangdd01) and Yuwen Xiong
+
+We also thank Guangming Sheng and Chi Zhang for helpful discussions.
+
+```bibtex
+@misc{zheng2025easyr1,
+  title        = {EasyR1: An Efficient, Scalable, Multi-Modality RL Training Framework},
+  author       = {Yaowei Zheng, Junting Lu, Shenzhi Wang, Zhangchi Feng, Dongdong Kuang, Yuwen Xiong},
+  howpublished = {\url{https://github.com/hiyouga/EasyR1}},
+  year         = {2025}
+}
+```
+
+We recommend to also cite the original work.
+
+```bibtex
+@article{sheng2024hybridflow,
+  title   = {HybridFlow: A Flexible and Efficient RLHF Framework},
+  author  = {Guangming Sheng and Chi Zhang and Zilingfeng Ye and Xibin Wu and Wang Zhang and Ru Zhang and Yanghua Peng and Haibin Lin and Chuan Wu},
+  year    = {2024},
+  journal = {arXiv preprint arXiv: 2409.19256}
+}
+```

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import os
+
+from .utils.py_functional import is_package_available
+
+
+if is_package_available("modelscope"):
+    from modelscope.utils.hf_util import patch_hub  # type: ignore
+
+
+__version__ = "0.3.1.dev0"
+
+
+if os.getenv("USE_MODELSCOPE_HUB", "0").lower() in ["true", "y", "1"]:
+    # Patch hub to download models from modelscope to speed up.
+    if not is_package_available("modelscope"):
+        raise ImportError("You are using the modelscope hub, please install modelscope by `pip install modelscope`.")
+
+    patch_hub()

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/init

@@ -0,0 +1 @@
+

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/models/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/models/monkey_patch.py

@@ -0,0 +1,32 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from .transformers.flash_attention_utils import flash_attention_forward
+from .transformers.qwen2_vl import qwen2_vl_attn_forward
+
+
+def apply_ulysses_patch(model_type: str) -> None:
+    if model_type in ("llama", "mllama", "gemma", "gemma2", "mistral", "qwen2", "qwen3", "qwen3_moe"):
+        ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = flash_attention_forward
+    elif model_type in ("qwen2_vl", "qwen2_5_vl"):
+        from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLFlashAttention2
+        from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLFlashAttention2
+
+        Qwen2VLFlashAttention2.forward = qwen2_vl_attn_forward
+        Qwen2_5_VLFlashAttention2.forward = qwen2_vl_attn_forward
+    else:
+        raise NotImplementedError(f"Model architecture {model_type} is not supported yet.")

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/models/transformers/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/models/transformers/flash_attention_utils.py

@@ -0,0 +1,191 @@
+# Copyright 2024 The Fairseq Authors and the HuggingFace Inc. team
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Based on https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/modeling_flash_attention_utils.py
+#
+# 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.
+
+import inspect
+import os
+from typing import Optional, Tuple
+
+import torch
+import torch.distributed as dist
+from transformers.modeling_flash_attention_utils import _flash_attention_forward, fa_peft_integration_check
+from transformers.utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10
+
+from ...utils.ulysses import (
+    gather_heads_scatter_seq,
+    gather_seq_scatter_heads,
+    get_ulysses_sequence_parallel_group,
+    get_ulysses_sequence_parallel_world_size,
+)
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+
+    _flash_supports_window_size = "window_size" in inspect.signature(flash_attn_func).parameters
+    _flash_supports_deterministic = "deterministic" in inspect.signature(flash_attn_func).parameters
+    _flash_deterministic_enabled = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
+    _flash_use_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+
+def prepare_fa2_from_position_ids(
+    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, position_ids: torch.Tensor
+):
+    query = query.view(-1, query.size(-2), query.size(-1))
+    key = key.contiguous().view(-1, key.size(-2), key.size(-1))
+    value = value.contiguous().view(-1, value.size(-2), value.size(-1))
+    position_ids = position_ids.flatten()
+    indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32)
+    cu_seqlens = torch.cat(
+        (
+            indices_q[position_ids == 0],
+            torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),
+        )
+    )
+    max_length = cu_seqlens.diff().max()  # use cu_seqlens to infer max_length for qwen2vl mrope
+    return (query, key, value, indices_q, (cu_seqlens, cu_seqlens), (max_length, max_length))
+
+
+def _custom_flash_attention_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    query_length: int,
+    is_causal: bool = True,
+    position_ids: Optional[torch.Tensor] = None,
+    sliding_window: Optional[int] = None,
+    use_top_left_mask: bool = False,
+    deterministic: Optional[bool] = None,
+    **kwargs,
+):
+    """
+    Patches flash attention forward to handle 3D position ids in mrope. (3, batch_size, seq_length)
+    """
+    if not use_top_left_mask:
+        causal = is_causal
+    else:
+        causal = is_causal and query_length != 1
+
+    # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).
+    use_sliding_windows = (
+        _flash_supports_window_size and sliding_window is not None and key_states.shape[1] > sliding_window
+    )
+    flash_kwargs = {"window_size": (sliding_window, sliding_window)} if use_sliding_windows else {}
+
+    if _flash_supports_deterministic:
+        flash_kwargs["deterministic"] = deterministic if deterministic is not None else _flash_deterministic_enabled
+
+    if kwargs.get("softcap") is not None:
+        flash_kwargs["softcap"] = kwargs.pop("softcap")
+
+    query_states, key_states, value_states = fa_peft_integration_check(
+        query_states, key_states, value_states, target_dtype=torch.bfloat16
+    )
+
+    sp_size = get_ulysses_sequence_parallel_world_size()
+    if sp_size > 1:
+        # (batch_size, seq_length, num_head, head_size)
+        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)
+        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)
+        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)
+        position_ids_lst = [torch.empty_like(position_ids) for _ in range(sp_size)]
+        position_ids = dist.all_gather(position_ids_lst, position_ids, group=get_ulysses_sequence_parallel_group())
+        position_ids = torch.cat(position_ids_lst, dim=-1)  # (..., batch_size, seq_length)
+
+    if position_ids is not None and position_ids.dim() == 3:  # qwen2vl mrope
+        position_ids = position_ids[0]
+
+    if position_ids is not None and query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all():
+        batch_size = query_states.size(0)
+        query_states, key_states, value_states, _, cu_seq_lens, max_seq_lens = prepare_fa2_from_position_ids(
+            query_states, key_states, value_states, position_ids
+        )
+        cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+        max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+        attn_output = flash_attn_varlen_func(
+            query_states,
+            key_states,
+            value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_in_batch_q,
+            max_seqlen_k=max_seqlen_in_batch_k,
+            dropout_p=kwargs.pop("dropout", 0.0),
+            softmax_scale=kwargs.pop("softmax_scale", None),
+            causal=causal,
+            **flash_kwargs,
+        )
+        attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1))
+    else:
+        attn_output = _flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            query_length,
+            is_causal=is_causal,
+            sliding_window=sliding_window,
+            use_top_left_mask=use_top_left_mask,
+            deterministic=deterministic,
+            **kwargs,
+        )  # do not pass position_ids to old flash_attention_forward
+
+    if sp_size > 1:
+        # (batch_size, seq_length, num_head, head_size)
+        attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1)
+
+    return attn_output
+
+
+def flash_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    dropout: float = 0.0,
+    scaling: Optional[float] = None,
+    sliding_window: Optional[int] = None,
+    softcap: Optional[float] = None,
+    **kwargs,
+) -> Tuple[torch.Tensor, None]:
+    # This is before the transpose
+    q_len = query.shape[2]
+
+    # FA2 uses non-transposed inputs
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
+    value = value.transpose(1, 2)
+
+    # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice
+    kwargs.pop("is_causal", None)
+
+    attn_output = _custom_flash_attention_forward(
+        query,
+        key,
+        value,
+        attention_mask,
+        query_length=q_len,
+        is_causal=True,
+        dropout=dropout,
+        softmax_scale=scaling,
+        sliding_window=sliding_window,
+        softcap=softcap,
+        use_top_left_mask=_flash_use_top_left_mask,
+        **kwargs,
+    )
+
+    return attn_output, None

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/models/transformers/qwen2_vl.py

@@ -0,0 +1,189 @@
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Based on:
+# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py
+#
+# 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.
+
+from typing import Optional, Tuple
+
+import torch
+
+from .flash_attention_utils import flash_attention_forward
+
+
+try:
+    from transformers.models.qwen2_vl.modeling_qwen2_vl import (
+        Qwen2VLAttention,
+        apply_multimodal_rotary_pos_emb,
+        repeat_kv,
+    )
+    from transformers.models.qwen2_vl.processing_qwen2_vl import Qwen2VLProcessor
+except ImportError:
+    pass
+
+
+def get_rope_index(
+    processor: "Qwen2VLProcessor",
+    input_ids: torch.Tensor,
+    image_grid_thw: Optional[torch.Tensor] = None,
+    video_grid_thw: Optional[torch.Tensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Gets the position ids for Qwen2-VL, it should be generated before sharding the sequence.
+    The batch dim has been removed and the input_ids should be a 1D tensor representing a single example.
+    https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L1546
+    """
+    spatial_merge_size = processor.image_processor.merge_size
+    tokens_per_second = 2
+    image_token_id = processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")
+    video_token_id = processor.tokenizer.convert_tokens_to_ids("<|video_pad|>")
+    vision_start_token_id = processor.tokenizer.convert_tokens_to_ids("<|vision_start|>")
+    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        position_ids = torch.ones(3, input_ids.size(0), dtype=input_ids.dtype, device=input_ids.device)  # (3, seqlen)
+        image_index, video_index = 0, 0
+        input_ids = input_ids[attention_mask == 1]
+        image_nums, video_nums = 0, 0
+        vision_start_indices = torch.argwhere(input_ids == vision_start_token_id)
+        vision_tokens = input_ids[vision_start_indices + 1]
+        image_nums = (vision_tokens == image_token_id).sum()
+        video_nums = (vision_tokens == video_token_id).sum()
+        input_tokens = input_ids.tolist()
+        llm_pos_ids_list: list = []
+        st = 0
+        remain_images, remain_videos = image_nums, video_nums
+        for _ in range(image_nums + video_nums):
+            if image_token_id in input_tokens and remain_images > 0:
+                ed_image = input_tokens.index(image_token_id, st)
+            else:
+                ed_image = len(input_tokens) + 1
+            if video_token_id in input_tokens and remain_videos > 0:
+                ed_video = input_tokens.index(video_token_id, st)
+            else:
+                ed_video = len(input_tokens) + 1
+            if ed_image < ed_video:
+                t, h, w = (
+                    image_grid_thw[image_index][0],
+                    image_grid_thw[image_index][1],
+                    image_grid_thw[image_index][2],
+                )
+                second_per_grid_t = 0
+                image_index += 1
+                remain_images -= 1
+                ed = ed_image
+            else:
+                t, h, w = (
+                    video_grid_thw[video_index][0],
+                    video_grid_thw[video_index][1],
+                    video_grid_thw[video_index][2],
+                )
+                if second_per_grid_ts is not None:
+                    second_per_grid_t = second_per_grid_ts[video_index]
+                else:
+                    second_per_grid_t = 1.0
+
+                video_index += 1
+                remain_videos -= 1
+                ed = ed_video
+
+            llm_grid_t, llm_grid_h, llm_grid_w = (
+                t.item(),
+                h.item() // spatial_merge_size,
+                w.item() // spatial_merge_size,
+            )
+            text_len = ed - st
+
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)
+            t_index = (t_index * second_per_grid_t * tokens_per_second).long().flatten()
+            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+            llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+            st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+        if st < len(input_tokens):
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            text_len = len(input_tokens) - st
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+        position_ids[..., attention_mask == 1] = llm_positions.to(position_ids.device)
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1).to(input_ids.device)
+        else:
+            position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).view(1, -1).expand(3, -1)
+
+    return position_ids
+
+
+def qwen2_vl_attn_forward(
+    self: "Qwen2VLAttention",
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+    **kwargs,
+) -> Tuple[torch.Tensor, None, None]:
+    bsz, q_len, _ = hidden_states.size()  # q_len = seq_length / sp_size
+    query_states = self.q_proj(hidden_states)  # (batch_size, seq_length / sp_size, num_heads * head_size)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    # Because the input can be padded, the absolute sequence length depends on the max position id.
+    if position_embeddings is None:
+        cos, sin = self.rotary_emb(value_states, position_ids)
+    else:
+        cos, sin = position_embeddings
+
+    query_states, key_states = apply_multimodal_rotary_pos_emb(
+        query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
+    )
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+    dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+    sliding_window = None
+    if (
+        self.config.use_sliding_window
+        and getattr(self.config, "sliding_window", None) is not None
+        and self.layer_idx >= self.config.max_window_layers
+    ):
+        sliding_window = self.config.sliding_window
+
+    attn_output, _ = flash_attention_forward(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        dropout=dropout_rate,
+        sliding_window=sliding_window,
+        position_ids=position_ids,  # important: pass position ids
+    )  # (batch_size, seq_length, num_head / sp_size, head_size)
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+    attn_output = self.o_proj(attn_output)
+    return attn_output, None, None

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/protocol.py

@@ -0,0 +1,689 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+Implement base data transfer protocol between any two functions, modules.
+We can subclass Protocol to define more detailed batch info with specific keys
+"""
+
+import copy
+import io
+import pickle
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import ray
+import torch
+from numpy.typing import NDArray
+from tensordict import TensorDict
+from torch.distributed import ProcessGroup
+from torch.utils.data import DataLoader
+
+from .utils.py_functional import union_two_dict
+
+
+try:
+    import tensordict
+
+    tensordict.set_lazy_legacy(False).set()
+except Exception:
+    pass
+
+
+__all__ = ["DataProto", "union_tensor_dict"]
+
+
+def pad_dataproto_to_divisor(data: "DataProto", size_divisor: int) -> Tuple["DataProto", int]:
+    """Pad a DataProto to size divisible by size_divisor
+
+    Args:
+        data (DataProto): the unpadded DataProto
+        size_divisor (int): size divisor
+
+    Returns:
+        data (DataProto): the padded DataProto
+        pad_size (int)
+    """
+    assert isinstance(data, DataProto), "data must be a DataProto"
+    if len(data) % size_divisor != 0:
+        pad_size = size_divisor - len(data) % size_divisor
+        padding_protos = []
+        remaining_pad = pad_size
+        while remaining_pad > 0:
+            take_size = min(remaining_pad, len(data))
+            padding_protos.append(data[:take_size])
+            remaining_pad -= take_size
+
+        data_padded = DataProto.concat([data] + padding_protos)
+    else:
+        pad_size = 0
+        data_padded = data
+
+    return data_padded, pad_size
+
+
+def unpad_dataproto(data: "DataProto", pad_size: int) -> "DataProto":
+    if pad_size != 0:
+        data = data[:-pad_size]
+
+    return data
+
+
+def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict:
+    """Union two tensordicts."""
+    if tensor_dict1.batch_size != tensor_dict2.batch_size:
+        raise ValueError(
+            f"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}"
+        )
+
+    for key in tensor_dict2.keys():
+        if key in tensor_dict1 and not torch.equal(tensor_dict1[key], tensor_dict2[key]):
+            raise ValueError(f"Key already exists: {key}.")
+
+        tensor_dict1[key] = tensor_dict2[key]
+
+    return tensor_dict1
+
+
+def union_numpy_dict(tensor_dict1: Dict[str, NDArray], tensor_dict2: Dict[str, NDArray]) -> Dict[str, NDArray]:
+    for key in tensor_dict2.keys():
+        if key in tensor_dict1:
+            assert isinstance(tensor_dict2[key], np.ndarray)
+            assert isinstance(tensor_dict1[key], np.ndarray)
+            if not np.all(tensor_dict1[key] == tensor_dict2[key]):
+                raise ValueError(f"Key already exists: {key}.")
+
+        tensor_dict1[key] = tensor_dict2[key]
+
+    return tensor_dict1
+
+
+def batch_collate(features: List[Dict[str, Any]]) -> Dict[str, List[Any]]:
+    if len(features) == 0:
+        return {}
+
+    batch_features = defaultdict(list)
+    for feature in features:
+        for key, value in feature.items():
+            batch_features[key].append(value)
+
+    return batch_features
+
+
+def fold_batch_dim(data: "DataProto", new_batch_size: int):
+    """
+    Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx]
+    """
+    batch_size = data.batch.batch_size[0]
+
+    assert batch_size % new_batch_size == 0
+
+    tensor: TensorDict = data.batch
+    non_tensor = data.non_tensor_batch
+
+    tensor = tensor.view(new_batch_size, -1)
+    tensor.auto_batch_size_(batch_dims=1)
+
+    for key, value in non_tensor.items():
+        non_tensor[key] = np.reshape(value, newshape=(new_batch_size, -1, *value.shape[1:]))
+
+    return DataProto(batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info)
+
+
+def collate_fn(data_items: list["DataProtoItem"]):
+    batch = []
+    non_tensor_batch = []
+    for data in data_items:
+        batch.append(data.batch)
+        non_tensor_batch.append(data.non_tensor_batch)
+
+    batch = torch.stack(batch).contiguous()
+    non_tensor_batch = batch_collate(non_tensor_batch)
+    non_tensor_batch = {key: np.array(value, dtype=object) for key, value in non_tensor_batch.items()}
+    return DataProto(batch=batch, non_tensor_batch=non_tensor_batch)
+
+
+@dataclass
+class DataProtoItem:
+    batch: Optional[TensorDict] = None
+    non_tensor_batch: Dict[str, NDArray] = field(default_factory=dict)
+    meta_info: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class DataProto:
+    """
+    A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions.
+    It contains a batch (TensorDict) and a meta_info (Dict). The batch is a TensorDict https://pytorch.org/tensordict/.
+    TensorDict allows you to manipulate a dictionary of Tensors like a single Tensor. Ideally, the tensors with the
+    same batch size should be put inside batch.
+    """
+
+    batch: Optional[TensorDict] = None
+    non_tensor_batch: Dict[str, NDArray] = field(default_factory=dict)
+    meta_info: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        self.check_consistency()  # perform necessary checking
+
+    def __len__(self) -> int:
+        if self.batch is not None:
+            return self.batch.batch_size[0]
+        elif self.non_tensor_batch is not None and len(self.non_tensor_batch) > 0:
+            pivot_key = list(self.non_tensor_batch.keys())[0]
+            return self.non_tensor_batch[pivot_key].shape[0]
+        else:
+            return 0
+
+    def __getitem__(self, item: Union[int, slice]) -> Union["DataProto", "DataProtoItem"]:
+        tensor_data = self.batch[item]
+        non_tensor_data = {key: value[item] for key, value in self.non_tensor_batch.items()}
+        return_type = DataProto if isinstance(item, slice) else DataProtoItem
+        return return_type(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)
+
+    def __getstate__(self) -> Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]:
+        buffer = io.BytesIO()
+        if self.batch is not None:
+            self.batch: TensorDict = self.batch.contiguous()
+            self.batch: TensorDict = self.batch.consolidate()
+
+        torch.save(self.batch, buffer)
+        buffer_bytes = buffer.getvalue()
+        return buffer_bytes, self.non_tensor_batch, self.meta_info
+
+    def __setstate__(self, data: Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]) -> None:
+        batch_deserialized_bytes, non_tensor_batch, meta_info = data
+        batch_deserialized = io.BytesIO(batch_deserialized_bytes)
+        batch = torch.load(batch_deserialized, weights_only=False, map_location="cpu")
+        self.batch = batch
+        self.non_tensor_batch = non_tensor_batch
+        self.meta_info = meta_info
+
+    def save_to_disk(self, filepath: str) -> None:
+        with open(filepath, "wb") as f:
+            pickle.dump(self, f)
+
+    @staticmethod
+    def load_from_disk(filepath: str) -> "DataProto":
+        with open(filepath, "rb") as f:
+            data = pickle.load(f)
+            return data
+
+    def print_size(self, prefix: str = "") -> None:
+        size_of_tensordict = 0
+        if self.batch is not None:
+            for tensor in self.batch.values():
+                if isinstance(tensor, torch.Tensor):
+                    size_of_tensordict += tensor.element_size() * tensor.numel()
+
+        size_of_numpy_array = 0
+        for value in self.non_tensor_batch.values():
+            size_of_numpy_array += value.nbytes
+
+        size_of_numpy_array /= 1024**3
+        size_of_tensordict /= 1024**3
+
+        message = f"Size of tensordict: {size_of_tensordict} GB, size of non_tensor_batch: {size_of_numpy_array} GB."
+        print({prefix}, {message})
+
+    def check_consistency(self):
+        """Check the consistency of the DataProto. Mainly for batch and non_tensor_batch
+        We expose this function as a public one so that user can call themselves directly
+        """
+        if self.batch is not None:
+            assert len(self.batch.batch_size) == 1, "only support num_batch_dims=1"
+
+        if self.batch is not None and len(self.non_tensor_batch) != 0:
+            # TODO: we can actually lift this restriction if needed
+            assert len(self.batch.batch_size) == 1, "only support num_batch_dims=1 when non_tensor_batch is not empty."
+
+            batch_size = self.batch.batch_size[0]
+            for key, value in self.non_tensor_batch.items():
+                assert len(value) == batch_size, f"key {key} length {len(value)} is not equal to bsz {batch_size}."
+
+    @classmethod
+    def from_single_dict(
+        cls,
+        data: Dict[str, Union[torch.Tensor, NDArray]],
+        meta_info: Optional[Dict[str, Any]] = None,
+    ) -> "DataProto":
+        tensors, non_tensors = {}, {}
+        for key, value in data.items():
+            if isinstance(value, torch.Tensor):
+                tensors[key] = value
+            elif isinstance(value, np.ndarray):
+                non_tensors[key] = value
+            else:
+                raise ValueError(f"Unsupported type in data {type(value)}")
+
+        return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)
+
+    @classmethod
+    def from_dict(
+        cls,
+        tensors: Dict[str, torch.Tensor],
+        non_tensors: Dict[str, NDArray] = None,
+        meta_info: Optional[Dict[str, Any]] = None,
+        num_batch_dims: int = 1,
+    ) -> "DataProto":
+        """Create a DataProto from a dict of tensors. This assumes that
+        1. All the tensor in tensors have the same dim0
+        2. Only dim0 is the batch dim
+        """
+        assert len(tensors) > 0, "tensors must not be empty"
+        assert num_batch_dims > 0, "num_batch_dims must be greater than zero"
+        if non_tensors is not None:
+            assert num_batch_dims == 1, "only support num_batch_dims=1 when non_tensors is not None."
+
+        meta_info = meta_info or {}
+        non_tensors = non_tensors or {}
+        assert isinstance(non_tensors, dict), "non_tensors should be a dictionary."
+
+        # get and check batch size
+        batch_size = None
+        pivot_key = None
+        for key, tensor in tensors.items():
+            if batch_size is None:
+                batch_size = tensor.shape[:num_batch_dims]
+                pivot_key = key
+            else:
+                current_batch = tensor.shape[:num_batch_dims]
+                assert batch_size == current_batch, (
+                    f"Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. "
+                    f"Got {pivot_key} has {batch_size}, {key} has {current_batch}"
+                )
+
+        tensor_dict = TensorDict(source=tensors, batch_size=batch_size)
+        return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info)
+
+    def to(self, device: torch.device) -> "DataProto":
+        """move the batch to device
+
+        Args:
+            device (torch.device, str): torch device
+
+        Returns:
+            DataProto: the current DataProto
+
+        """
+        if self.batch is not None:
+            self.batch = self.batch.to(device)
+
+        return self
+
+    def select(
+        self,
+        batch_keys: Optional[List[str]] = None,
+        non_tensor_batch_keys: Optional[List[str]] = None,
+        meta_info_keys: Optional[List[str]] = None,
+        deepcopy: bool = False,
+    ) -> "DataProto":
+        """Select a subset of the DataProto via batch_keys and meta_info_keys
+
+        Args:
+            batch_keys (list, optional): a list of strings indicating the keys in batch to select
+            meta_info_keys (list, optional): a list of keys indicating the meta info to select
+
+        Returns:
+            DataProto: the DataProto with the selected batch_keys and meta_info_keys
+        """
+        # TODO (zhangchi.usc1992) whether to copy
+        if batch_keys is not None:
+            batch_keys = tuple(batch_keys)
+            sub_batch = self.batch.select(*batch_keys)
+        else:
+            sub_batch = self.batch
+
+        if non_tensor_batch_keys is not None:
+            non_tensor_batch = {k: v for k, v in self.non_tensor_batch.items() if k in non_tensor_batch_keys}
+        else:
+            non_tensor_batch = self.non_tensor_batch
+
+        if deepcopy:
+            non_tensor_batch = copy.deepcopy(non_tensor_batch)
+
+        if meta_info_keys is not None:
+            sub_meta_info = {k: v for k, v in self.meta_info.items() if k in meta_info_keys}
+        else:
+            sub_meta_info = self.meta_info
+
+        if deepcopy:
+            sub_meta_info = copy.deepcopy(sub_meta_info)
+
+        return DataProto(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info)
+
+    def pop(
+        self,
+        batch_keys: Optional[List[str]] = None,
+        non_tensor_batch_keys: Optional[List[str]] = None,
+        meta_info_keys: Optional[List[str]] = None,
+    ) -> "DataProto":
+        """Pop a subset of the DataProto via `batch_keys` and `meta_info_keys`
+
+        Args:
+            batch_keys (list, optional): a list of strings indicating the keys in batch to pop
+            meta_info_keys (list, optional): a list of keys indicating the meta info to pop
+
+        Returns:
+            DataProto: the DataProto with the poped batch_keys and meta_info_keys
+        """
+        assert batch_keys is not None
+        non_tensor_batch_keys = non_tensor_batch_keys or []
+        meta_info_keys = meta_info_keys or []
+
+        tensors = {}
+        for key in batch_keys:
+            tensors[key] = self.batch.pop(key)
+
+        non_tensors = {}
+        for key in non_tensor_batch_keys:
+            non_tensors[key] = self.non_tensor_batch.pop(key)
+
+        meta_info = {}
+        for key in meta_info_keys:
+            meta_info[key] = self.meta_info.pop(key)
+
+        return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)
+
+    def rename(
+        self, old_keys: Optional[Union[str, List[str]]] = None, new_keys: Optional[Union[str, List[str]]] = None
+    ) -> "DataProto":
+        """
+        Note that this function only rename the key in the batch
+        """
+
+        def validate_input(keys):
+            if keys is not None:
+                if isinstance(keys, str):
+                    keys = [keys]
+                elif isinstance(keys, list):
+                    pass
+                else:
+                    raise TypeError(f"keys must be a list or a string, but got {type(keys)}")
+            return keys
+
+        old_keys = validate_input(old_keys)
+        new_keys = validate_input(new_keys)
+
+        if len(new_keys) != len(old_keys):
+            raise ValueError(
+                f"new_keys and old_keys must have the same length, but got {len(new_keys)} and {len(old_keys)}"
+            )
+
+        self.batch.rename_key_(tuple(old_keys), tuple(new_keys))
+
+        return self
+
+    def union(self, other: "DataProto") -> "DataProto":
+        """Union with another DataProto. Union batch and meta_info separately.
+        Throw an error if
+        - there are conflict keys in batch and they are not equal
+        - the batch size of two data batch is not the same
+        - there are conflict keys in meta_info and they are not the same.
+
+        Args:
+            other (DataProto): another DataProto to union
+
+        Returns:
+            DataProto: the DataProto after union
+        """
+        self.batch = union_tensor_dict(self.batch, other.batch)
+        self.non_tensor_batch = union_numpy_dict(self.non_tensor_batch, other.non_tensor_batch)
+        self.meta_info = union_two_dict(self.meta_info, other.meta_info)
+        return self
+
+    def make_iterator(
+        self, mini_batch_size: int, epochs: int, seed: int = None, dataloader_kwargs: Dict[str, Any] = None
+    ):
+        """Make an iterator from the DataProto. This is built upon that TensorDict can be used as a normal Pytorch
+        dataset. See https://pytorch.org/tensordict/tutorials/data_fashion for more details.
+
+        Args:
+            mini_batch_size (int): mini-batch size when iterating the dataset. We require that
+                ``batch.batch_size[0] % mini_batch_size == 0``
+            epochs (int): number of epochs when iterating the dataset.
+            dataloader_kwargs: internally, it returns a DataLoader over the batch.
+                The dataloader_kwargs is the kwargs passed to the DataLoader
+
+        Returns:
+            Iterator: an iterator that yields a mini-batch data at a time. The total number of iteration steps is
+            ``self.batch.batch_size * epochs // mini_batch_size``
+        """
+        assert self.batch.batch_size[0] % mini_batch_size == 0, f"{self.batch.batch_size[0]} % {mini_batch_size} != 0"
+        # we can directly create a dataloader from TensorDict
+        if dataloader_kwargs is None:
+            dataloader_kwargs = {}
+
+        if seed is not None:
+            generator = torch.Generator()
+            generator.manual_seed(seed)
+        else:
+            generator = None
+
+        assert isinstance(dataloader_kwargs, Dict)
+        train_dataloader = DataLoader(
+            dataset=self, batch_size=mini_batch_size, collate_fn=collate_fn, generator=generator, **dataloader_kwargs
+        )
+
+        def get_data():
+            for _ in range(epochs):
+                for d in train_dataloader:
+                    d.meta_info = self.meta_info
+                    yield d
+
+        return iter(get_data())
+
+    def chunk(self, chunks: int) -> List["DataProto"]:
+        """Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split.
+
+        Args:
+            chunks (int): the number of chunks to split on dim=0
+
+        Returns:
+            List[DataProto]: a list of DataProto after splitting
+        """
+        assert len(self) % chunks == 0, (
+            f"only support equal chunk. Got size of DataProto {len(self)} and chunk {chunks}."
+        )
+        if self.batch is not None:
+            batch_lst = self.batch.chunk(chunks=chunks, dim=0)
+        else:
+            batch_lst = [None for _ in range(chunks)]
+
+        non_tensor_batch_lst = [{} for _ in range(chunks)]
+        for key, value in self.non_tensor_batch.items():
+            assert isinstance(value, np.ndarray)
+            non_tensor_lst = np.array_split(value, chunks)
+            assert len(non_tensor_lst) == chunks
+            for i in range(chunks):
+                non_tensor_batch_lst[i][key] = non_tensor_lst[i]
+
+        output = []
+        for i in range(chunks):
+            output.append(
+                DataProto(batch=batch_lst[i], non_tensor_batch=non_tensor_batch_lst[i], meta_info=self.meta_info)
+            )
+
+        return output
+
+    def split(self, split_size: int) -> List["DataProto"]:
+        chunks = len(self) // split_size
+        return self.chunk(chunks)
+
+    @staticmethod
+    def concat(data: List["DataProto"]) -> "DataProto":
+        """Concat a list of DataProto. The batch is concatenated among dim=0.
+        The meta_info is assumed to be identical and will use the first one.
+
+        Args:
+            data (List[DataProto]): list of DataProto
+
+        Returns:
+            DataProto: concatenated DataProto
+        """
+        batch_lst = [batch.batch for batch in data]
+        if batch_lst[0] is not None:
+            new_batch = torch.cat(batch_lst, dim=0)
+        else:
+            new_batch = None
+
+        non_tensor_batch = batch_collate([d.non_tensor_batch for d in data])
+        for key, value in non_tensor_batch.items():
+            non_tensor_batch[key] = np.concatenate(value, axis=0)
+
+        return DataProto(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info)
+
+    def reorder(self, indices: torch.Tensor) -> None:
+        """
+        Note that this operation is in-place
+        """
+        indices_np = indices.detach().numpy()
+        self.batch = self.batch[indices]
+        self.non_tensor_batch = {key: value[indices_np] for key, value in self.non_tensor_batch.items()}
+
+    def repeat(self, repeat_times: int = 2, interleave: bool = True) -> "DataProto":
+        """
+        Repeat the batch data a specified number of times.
+
+        Args:
+            repeat_times (int): Number of times to repeat the data.
+            interleave (bool): Whether to interleave the repeated data.
+
+        Returns:
+            DataProto: A new DataProto with repeated data.
+        """
+        if self.batch is not None:
+            if interleave:
+                # Interleave the data
+                repeated_tensors = {
+                    key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items()
+                }
+            else:
+                # Stack the data
+                repeated_tensors = {
+                    key: tensor.unsqueeze(0).expand(repeat_times, *tensor.shape).reshape(-1, *tensor.shape[1:])
+                    for key, tensor in self.batch.items()
+                }
+
+            repeated_batch = TensorDict(
+                source=repeated_tensors,
+                batch_size=(self.batch.batch_size[0] * repeat_times,),
+            )
+        else:
+            repeated_batch = None
+
+        repeated_non_tensor_batch = {}
+        for key, value in self.non_tensor_batch.items():
+            if interleave:
+                repeated_non_tensor_batch[key] = np.repeat(value, repeat_times, axis=0)
+            else:
+                repeated_non_tensor_batch[key] = np.tile(value, (repeat_times,) + (1,) * (value.ndim - 1))
+
+        return DataProto(
+            batch=repeated_batch,
+            non_tensor_batch=repeated_non_tensor_batch,
+            meta_info=self.meta_info,
+        )
+
+
+@dataclass
+class DataProtoFuture:
+    """
+    DataProtoFuture aims to eliminate actual data fetching on driver. By doing so, the driver doesn't have to wait
+    for data so that asynchronous execution becomes possible.
+    DataProtoFuture contains a list of futures from another WorkerGroup of size world_size.
+    - collect_fn is a Callable that reduces the list of futures to a DataProto
+    - dispatch_fn is a Callable that partitions the DataProto into a list of DataProto of size world_size and then select
+
+    Potential issue: we can optimize dispatch_fn(collect_fn) such that only needed data is fetched on destination
+    - DataProtoFuture only supports directly passing from the output of a method to another input. You can't perform any
+    operation on the DataProtoFuture in driver.
+    """
+
+    collect_fn: Callable
+    futures: List[ray.ObjectRef]
+    dispatch_fn: Callable = None
+
+    @staticmethod
+    def concat(data: List[ray.ObjectRef]) -> "DataProtoFuture":
+        output = DataProtoFuture(collect_fn=DataProto.concat, futures=data)
+        return output
+
+    def chunk(self, chunks: int) -> List["DataProtoFuture"]:
+        from functools import partial
+
+        arg_future_lst = []
+        for i in range(chunks):
+            # note that we can't directly pass i and chunks
+            def dispatch_fn(x, i, chunks):
+                return x.chunk(chunks=chunks)[i]
+
+            arg_future = DataProtoFuture(
+                collect_fn=self.collect_fn, dispatch_fn=partial(dispatch_fn, i=i, chunks=chunks), futures=self.futures
+            )
+            arg_future_lst.append(arg_future)
+        return arg_future_lst
+
+    def get(self):
+        outputs = ray.get(self.futures)  # dp_size.
+        for output in outputs:
+            assert isinstance(output, DataProto)
+
+        outputs = self.collect_fn(outputs)  # select dp, concat
+        if self.dispatch_fn is not None:
+            outputs = self.dispatch_fn(outputs)  # split in batch dim, select using dp
+
+        return outputs
+
+
+def allgather_dict_tensors(
+    tensors: Union[Dict[str, torch.Tensor], TensorDict], size: int, group: ProcessGroup, dim: int = 0
+) -> Union[Dict[str, torch.Tensor], TensorDict]:
+    """
+    TODO: optimize this.
+    - We can use async ops
+    - We can use only one allgather
+    """
+    if isinstance(tensors, TensorDict):
+        is_tensor_dict = True
+        tensors_as_dict = tensors.to_dict()
+    else:
+        tensors_as_dict = tensors
+        is_tensor_dict = False
+
+    output = {}
+    sorted_keys = sorted(tensors_as_dict.keys())
+    for key in sorted_keys:
+        value = tensors_as_dict[key]
+        output[key] = [torch.empty_like(value) for _ in range(size)]
+        torch.distributed.all_gather(output[key], value, group=group, async_op=False)
+        output[key] = torch.cat(output[key], dim=dim)
+
+    if is_tensor_dict:
+        output = TensorDict(source=output, batch_size=tensors.batch_size[0] * size)
+
+    return output
+
+
+def all_gather_data_proto(data: DataProto, size: int, group: ProcessGroup) -> None:
+    # Note that this is an inplace operator just like torch.distributed.all_gather
+    prev_device = data.batch.device
+    data.batch = data.batch.cuda(device=torch.cuda.current_device())
+    data.batch = allgather_dict_tensors(data.batch.contiguous(), size=size, group=group, dim=0)
+    data.batch = data.batch.to(prev_device)
+    # all gather non_tensor_batch
+    all_non_tensor_batch = [None for _ in range(size)]
+    torch.distributed.all_gather_object(all_non_tensor_batch, data.non_tensor_batch, group=group)
+    data.non_tensor_batch = {k: np.concatenate([d[k] for d in all_non_tensor_batch]) for k in data.non_tensor_batch}

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "verl"
+dynamic = [
+    "version",
+    "dependencies",
+    "optional-dependencies",
+    "requires-python",
+    "authors",
+    "description",
+    "readme",
+    "license"
+]
+
+[tool.ruff]
+target-version = "py39"
+line-length = 119
+indent-width = 4
+
+[tool.ruff.lint]
+ignore = ["C901", "E501", "E741", "W605", "C408"]
+select = ["C", "E", "F", "I", "W", "RUF022"]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403", "F811"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
+known-first-party = ["verl"]
+known-third-party = ["torch", "transformers", "wandb"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/requirements.txt

@@ -0,0 +1,19 @@
+accelerate
+codetiming
+datasets
+liger-kernel
+mathruler
+numpy
+omegaconf
+pandas
+peft
+pillow
+pyarrow>=15.0.0
+pylatexenc
+qwen-vl-utils
+ray[default]
+tensordict
+torchdata
+transformers>=4.51.0
+vllm>=0.7.3
+wandb

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/base/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from .worker import Worker
+from .worker_group import ClassWithInitArgs, ResourcePool, WorkerGroup
+
+
+__all__ = ["ClassWithInitArgs", "ResourcePool", "Worker", "WorkerGroup"]

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/base/decorator.py

@@ -0,0 +1,213 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from enum import Enum, auto
+from functools import wraps
+from types import FunctionType
+from typing import TYPE_CHECKING, Dict, List, Literal, Union
+
+import ray
+
+from ...protocol import DataProto, DataProtoFuture
+
+
+if TYPE_CHECKING:
+    from .worker_group import WorkerGroup
+
+
+# here we add a magic number of avoid user-defined function already have this attribute
+MAGIC_ATTR = "attrs_3141562937"
+
+
+class Dispatch(Enum):
+    RANK_ZERO = auto()
+    ONE_TO_ALL = auto()
+    ALL_TO_ALL = auto()
+    DP_COMPUTE = auto()
+    DP_COMPUTE_PROTO = auto()
+    DP_COMPUTE_PROTO_WITH_FUNC = auto()
+    DP_COMPUTE_METRIC = auto()
+
+
+class Execute(Enum):
+    ALL = 0
+    RANK_ZERO = 1
+
+
+def _split_args_kwargs_data_proto(chunks: int, *args, **kwargs):
+    splitted_args = []
+    for arg in args:
+        assert isinstance(arg, (DataProto, DataProtoFuture))
+        splitted_args.append(arg.chunk(chunks=chunks))
+
+    splitted_kwargs = {}
+    for key, value in kwargs.items():
+        assert isinstance(value, (DataProto, DataProtoFuture))
+        splitted_kwargs[key] = value.chunk(chunks=chunks)
+
+    return splitted_args, splitted_kwargs
+
+
+def dispatch_one_to_all(worker_group: "WorkerGroup", *args, **kwargs):
+    args = tuple([arg] * worker_group.world_size for arg in args)
+    kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()}
+    return args, kwargs
+
+
+def dispatch_all_to_all(worker_group: "WorkerGroup", *args, **kwargs):
+    return args, kwargs
+
+
+def collect_all_to_all(worker_group: "WorkerGroup", output):
+    return output
+
+
+def _concat_data_proto_or_future(outputs: List[DataProto]) -> DataProto:
+    # make sure all the elements in output has the same type
+    for output in outputs:
+        assert type(output) is type(outputs[0])
+
+    output = outputs[0]
+
+    if isinstance(output, DataProto):
+        return DataProto.concat(outputs)
+    elif isinstance(output, ray.ObjectRef):
+        return DataProtoFuture.concat(outputs)
+    else:
+        raise NotImplementedError
+
+
+def dispatch_dp_compute(worker_group: "WorkerGroup", *args, **kwargs):
+    for arg in args:
+        assert isinstance(arg, (tuple, list)) and len(arg) == worker_group.world_size
+
+    for value in kwargs.values():
+        assert isinstance(value, (tuple, list)) and len(value) == worker_group.world_size
+
+    return args, kwargs
+
+
+def collect_dp_compute(worker_group: "WorkerGroup", outputs: List[DataProto]) -> List[DataProto]:
+    assert len(outputs) == worker_group.world_size
+    return outputs
+
+
+def dispatch_dp_compute_data_proto(worker_group: "WorkerGroup", *args, **kwargs):
+    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args, **kwargs)
+    return splitted_args, splitted_kwargs
+
+
+def dispatch_dp_compute_data_proto_with_func(worker_group: "WorkerGroup", *args, **kwargs):
+    assert type(args[0]) is FunctionType  # NOTE: The first one args is a function!
+    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args[1:], **kwargs)
+    splitted_args_with_func = [[args[0]] * worker_group.world_size] + splitted_args
+    return splitted_args_with_func, splitted_kwargs
+
+
+def collect_dp_compute_data_proto(worker_group: "WorkerGroup", outputs: List[DataProto]) -> DataProto:
+    for output in outputs:
+        assert isinstance(output, (DataProto, ray.ObjectRef)), f"Expect a DataProto, but got {type(output)}"
+
+    outputs = collect_dp_compute(worker_group, outputs)
+    return _concat_data_proto_or_future(outputs)
+
+
+def get_predefined_dispatch_fn(dispatch_mode: Dispatch):
+    predefined_dispatch_mode_fn = {
+        Dispatch.ONE_TO_ALL: {
+            "dispatch_fn": dispatch_one_to_all,
+            "collect_fn": collect_all_to_all,
+        },
+        Dispatch.ALL_TO_ALL: {
+            "dispatch_fn": dispatch_all_to_all,
+            "collect_fn": collect_all_to_all,
+        },
+        Dispatch.DP_COMPUTE: {
+            "dispatch_fn": dispatch_dp_compute,
+            "collect_fn": collect_dp_compute,
+        },
+        Dispatch.DP_COMPUTE_PROTO: {
+            "dispatch_fn": dispatch_dp_compute_data_proto,
+            "collect_fn": collect_dp_compute_data_proto,
+        },
+        Dispatch.DP_COMPUTE_PROTO_WITH_FUNC: {
+            "dispatch_fn": dispatch_dp_compute_data_proto_with_func,
+            "collect_fn": collect_dp_compute_data_proto,
+        },
+        Dispatch.DP_COMPUTE_METRIC: {
+            "dispatch_fn": dispatch_dp_compute_data_proto,
+            "collect_fn": collect_dp_compute,
+        },
+    }
+    return predefined_dispatch_mode_fn[dispatch_mode]
+
+
+def get_predefined_execute_fn(execute_mode: Execute):
+    """
+    Note that here we only asks execute_all and execute_rank_zero to be implemented
+    Leave the choice of how these two functions handle argument 'blocking' to users
+    """
+    predefined_execute_mode_fn = {
+        Execute.ALL: {"execute_fn_name": "execute_all"},
+        Execute.RANK_ZERO: {"execute_fn_name": "execute_rank_zero"},
+    }
+    return predefined_execute_mode_fn[execute_mode]
+
+
+def _check_dispatch_mode(dispatch_mode: Union[Dispatch, Dict[Literal["dispatch_fn", "collect_fn"], FunctionType]]):
+    assert isinstance(dispatch_mode, (Dispatch, dict)), (
+        f"dispatch_mode must be a Dispatch or a Dict. Got {dispatch_mode}"
+    )
+    if isinstance(dispatch_mode, dict):
+        necessary_keys = ["dispatch_fn", "collect_fn"]
+        for key in necessary_keys:
+            assert key in dispatch_mode, f"key {key} should be in dispatch_mode if it is a dictionary"
+
+
+def _check_execute_mode(execute_mode: Execute):
+    assert isinstance(execute_mode, Execute), f"execute_mode must be a Execute. Got {execute_mode}"
+
+
+def _materialize_futures(*args, **kwargs):
+    new_args = []
+    for arg in args:
+        if isinstance(arg, DataProtoFuture):
+            arg = arg.get()
+        # add more type to materialize
+        new_args.append(arg)
+
+    for key, value in kwargs.items():
+        if isinstance(value, DataProtoFuture):
+            kwargs[key] = value.get()
+
+    new_args = tuple(new_args)
+    return new_args, kwargs
+
+
+def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True):
+    _check_dispatch_mode(dispatch_mode=dispatch_mode)
+    _check_execute_mode(execute_mode=execute_mode)
+
+    def decorator(func):
+        @wraps(func)
+        def inner(*args, **kwargs):
+            if materialize_futures:
+                args, kwargs = _materialize_futures(*args, **kwargs)
+            return func(*args, **kwargs)
+
+        attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, "blocking": blocking}
+        setattr(inner, MAGIC_ATTR, attrs)
+        return inner
+
+    return decorator

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/base/register_center/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/base/register_center/ray.py

@@ -0,0 +1,28 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import ray
+
+
+@ray.remote
+class WorkerGroupRegisterCenter:
+    def __init__(self, rank_zero_info):
+        self.rank_zero_info = rank_zero_info
+
+    def get_rank_zero_info(self):
+        return self.rank_zero_info
+
+
+def create_worker_group_register_center(name, info):
+    return WorkerGroupRegisterCenter.options(name=name).remote(info)

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/base/worker.py

@@ -0,0 +1,202 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+the class for Worker
+"""
+
+import os
+import socket
+from dataclasses import dataclass
+from typing import Tuple
+
+import ray
+import torch
+
+from .decorator import Dispatch, Execute, register
+from .register_center.ray import create_worker_group_register_center
+
+
+@dataclass
+class DistRankInfo:
+    tp_rank: int
+    dp_rank: int
+    pp_rank: int
+
+
+@dataclass
+class DistGlobalInfo:
+    tp_size: int
+    dp_size: int
+    pp_size: int
+
+
+class WorkerHelper:
+    def _get_node_ip(self) -> str:
+        host_ipv4 = os.getenv("MY_HOST_IP", None)
+        host_ipv6 = os.getenv("MY_HOST_IPV6", None)
+        host_ip_by_env = host_ipv4 or host_ipv6
+        host_ip_by_sdk = ray._private.services.get_node_ip_address()
+
+        host_ip = host_ip_by_env or host_ip_by_sdk
+        return host_ip
+
+    def _get_free_port(self) -> int:
+        with socket.socket() as sock:
+            sock.bind(("", 0))
+            return sock.getsockname()[1]
+
+    def get_availale_master_addr_port(self) -> Tuple[str, str]:
+        return self._get_node_ip(), str(self._get_free_port())
+
+    def _get_pid(self):
+        return
+
+
+class WorkerMeta:
+    keys = [
+        "WORLD_SIZE",
+        "RANK",
+        "LOCAL_WORLD_SIZE",
+        "LOCAL_RANK",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+        "CUDA_VISIBLE_DEVICES",
+    ]
+
+    def __init__(self, store) -> None:
+        self._store = store
+
+    def to_dict(self):
+        return {f"_{key.lower()}": self._store.get(f"_{key.lower()}", None) for key in WorkerMeta.keys}
+
+
+# we assume that in each WorkerGroup, there is a Master Worker
+class Worker(WorkerHelper):
+    """A (distributed) worker."""
+
+    _world_size: int
+    _rank: int
+    _local_world_size: int
+    _local_rank: int
+    _master_addr: str
+    _master_port: str
+    _cuda_visible_devices: str
+
+    def __new__(cls, *args, **kwargs):
+        instance = super().__new__(cls)
+
+        # note that here we use int to distinguish
+        disable_worker_init = int(os.getenv("DISABLE_WORKER_INIT", 0))
+        if disable_worker_init:
+            return instance
+
+        rank = os.getenv("RANK", None)
+        worker_group_prefix = os.getenv("WG_PREFIX", None)
+
+        # when decorator @ray.remote applies, __new__ will be called while we don't want to apply _configure_before_init
+        if None not in [rank, worker_group_prefix] and "ActorClass(" not in cls.__name__:
+            instance._configure_before_init(f"{worker_group_prefix}_register_center", int(rank))
+
+        return instance
+
+    def _configure_before_init(self, register_center_name: str, rank: int):
+        assert isinstance(rank, int), f"rank must be int, instead of {type(rank)}"
+
+        if rank == 0:
+            master_addr, master_port = self.get_availale_master_addr_port()
+            rank_zero_info = {
+                "MASTER_ADDR": master_addr,
+                "MASTER_PORT": master_port,
+            }
+            self.register_center = create_worker_group_register_center(name=register_center_name, info=rank_zero_info)
+            os.environ.update(rank_zero_info)
+
+    def __init__(self, cuda_visible_devices=None) -> None:
+        # construct a meta from envrionment variable. Note that the import must be inside the class because it is executed remotely
+        world_size = int(os.getenv("WORLD_SIZE"))
+        rank = int(os.getenv("RANK"))
+        self._rank = rank
+        self._world_size = world_size
+
+        if "AMD" in torch.cuda.get_device_name():
+            os.environ["CUDA_VISIBLE_DEVICES"] = os.getenv("ROCR_VISIBLE_DEVICES")
+            os.environ["LOCAL_RANK"] = os.getenv("RAY_LOCAL_RANK")
+            cuda_visible_devices = os.getenv("LOCAL_RANK", "0")
+            torch.cuda.set_device(int(cuda_visible_devices))
+
+        master_addr = os.getenv("MASTER_ADDR")
+        master_port = os.getenv("MASTER_PORT")
+
+        local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", "1"))
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        store = {
+            "_world_size": world_size,
+            "_rank": rank,
+            "_local_world_size": local_world_size,
+            "_local_rank": local_rank,
+            "_master_addr": master_addr,
+            "_master_port": master_port,
+        }
+        if cuda_visible_devices is not None:
+            store["_cuda_visible_devices"] = cuda_visible_devices
+
+        meta = WorkerMeta(store=store)
+        self._configure_with_meta(meta=meta)
+
+    def _configure_with_meta(self, meta: WorkerMeta):
+        """
+        This function should only be called inside by WorkerGroup
+        """
+        assert isinstance(meta, WorkerMeta)
+        self.__dict__.update(meta.to_dict())  # this is hacky
+        # print(f"__dict__: {self.__dict__}")
+        for key in WorkerMeta.keys:
+            val = self.__dict__.get(f"_{key.lower()}", None)
+            if val is not None:
+                # print(f"set {key} to {val}")
+                os.environ[key] = str(val)
+
+        os.environ["REDIS_STORE_SERVER_HOST"] = (
+            str(self._master_addr).replace("[", "").replace("]", "") if self._master_addr else ""
+        )
+
+    def get_master_addr_port(self):
+        return self._master_addr, self._master_port
+
+    def get_cuda_visible_devices(self):
+        cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", "not set")
+        return cuda_visible_devices
+
+    def print_rank0(self, *args, **kwargs):
+        if self.rank == 0:
+            print(*args, **kwargs)
+
+    @property
+    def world_size(self):
+        return self._world_size
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO_WITH_FUNC)
+    def execute_with_func_generator(self, func, *args, **kwargs):
+        ret_proto = func(self, *args, **kwargs)
+        return ret_proto
+
+    @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)
+    def execute_func_rank_zero(self, func, *args, **kwargs):
+        result = func(*args, **kwargs)
+        return result

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/base/worker_group.py

@@ -0,0 +1,194 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+the class of WorkerGroup
+"""
+
+import logging
+import signal
+import threading
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from .decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn
+
+
+class ResourcePool:
+    """The resource pool with meta info such as world size."""
+
+    def __init__(
+        self, process_on_nodes: Optional[Any] = None, max_colocate_count: int = 10, n_gpus_per_node: int = 8
+    ) -> None:
+        if process_on_nodes is None:
+            process_on_nodes = []
+
+        self._store = process_on_nodes
+        self.max_colocate_count = max_colocate_count
+        self.n_gpus_per_node = n_gpus_per_node  # this is left for future huawei GPU that contains 16 GPUs per node
+
+    def add_node(self, process_count):
+        self._store.append(process_count)
+
+    @property
+    def world_size(self):
+        return sum(self._store)
+
+    def __call__(self) -> Any:
+        return self._store
+
+    @property
+    def store(self):
+        return self._store
+
+    def local_world_size_list(self) -> List[int]:
+        nested_local_world_size_list = [
+            [local_world_size for _ in range(local_world_size)] for local_world_size in self._store
+        ]
+        return [item for row in nested_local_world_size_list for item in row]
+
+    def local_rank_list(self) -> List[int]:
+        nested_local_rank_list = [[i for i in range(local_world_size)] for local_world_size in self._store]  # noqa: C416
+        return [item for row in nested_local_rank_list for item in row]
+
+
+class ClassWithInitArgs:
+    """
+    This class stores a class constructor and the args/kwargs to construct the class.
+    It is used to instantiate the remote class.
+    """
+
+    def __init__(self, cls, *args, **kwargs) -> None:
+        self.cls = cls
+        self.args = args
+        self.kwargs = kwargs
+
+    def __call__(self) -> Any:
+        return self.cls(*self.args, **self.kwargs)
+
+
+def check_workers_alive(workers: List, is_alive: Callable, gap_time: float = 1) -> None:
+    while True:
+        for worker in workers:
+            if not is_alive(worker):
+                logging.warning(f"Worker {worker} is not alive, sending signal to main thread")
+                signal.raise_signal(signal.SIGABRT)
+
+        time.sleep(gap_time)
+
+
+class WorkerGroup:
+    """A group of workers"""
+
+    def __init__(self, resource_pool: ResourcePool, **kwargs) -> None:
+        self._is_init_with_detached_workers = True if resource_pool is None else False
+
+        if resource_pool is not None:
+            # handle the case when WorkGroup is attached to an existing one
+            self._procecss_dispatch_config = resource_pool()
+        else:
+            self._procecss_dispatch_config = None
+
+        self._workers = []
+        self._worker_names = []
+
+        self._master_addr = None
+        self._master_port = None
+
+        self._checker_thread: threading.Thread = None
+
+    def _is_worker_alive(self, worker):
+        raise NotImplementedError("WorkerGroup._is_worker_alive called, should be implemented in derived class.")
+
+    def _block_until_all_workers_alive(self) -> None:
+        while True:
+            all_state = [self._is_worker_alive(worker) for worker in self._workers]
+            if False in all_state:
+                time.sleep(1)
+            else:
+                break
+
+    def start_worker_aliveness_check(self, every_n_seconds=1) -> None:
+        # before starting checking worker aliveness, make sure all workers are already alive
+        self._block_until_all_workers_alive()
+
+        self._checker_thread = threading.Thread(
+            target=check_workers_alive, args=(self._workers, self._is_worker_alive, every_n_seconds)
+        )
+        self._checker_thread.start()
+
+    @property
+    def world_size(self):
+        return len(self._workers)
+
+    def _bind_worker_method(self, user_defined_cls, func_generator):
+        """
+        Bind the worker method to the WorkerGroup
+        """
+        for method_name in dir(user_defined_cls):
+            try:
+                method = getattr(user_defined_cls, method_name)
+                assert callable(method), f"{method_name} in {user_defined_cls} is not callable"
+            except Exception:
+                # if it is a property, it will fail because Class doesn't have instance property
+                continue
+
+            if hasattr(method, MAGIC_ATTR):
+                # this method is decorated by register
+                attribute = getattr(method, MAGIC_ATTR)
+                assert isinstance(attribute, Dict), f"attribute must be a dictionary. Got {type(attribute)}"
+                assert "dispatch_mode" in attribute, "attribute must contain dispatch_mode in its key"
+
+                dispatch_mode = attribute["dispatch_mode"]
+                execute_mode = attribute["execute_mode"]
+                blocking = attribute["blocking"]
+
+                # get dispatch fn
+                if isinstance(dispatch_mode, Dispatch):
+                    # get default dispatch fn
+                    fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode)
+                    dispatch_fn = fn["dispatch_fn"]
+                    collect_fn = fn["collect_fn"]
+                else:
+                    assert isinstance(dispatch_mode, dict)
+                    assert "dispatch_fn" in dispatch_mode
+                    assert "collect_fn" in dispatch_mode
+                    dispatch_fn = dispatch_mode["dispatch_fn"]
+                    collect_fn = dispatch_mode["collect_fn"]
+
+                # get execute_fn_name
+                execute_mode = get_predefined_execute_fn(execute_mode=execute_mode)
+                wg_execute_fn_name = execute_mode["execute_fn_name"]
+
+                # get execute_fn from string
+                try:
+                    execute_fn = getattr(self, wg_execute_fn_name)
+                    assert callable(execute_fn), "execute_fn must be callable"
+                except Exception:
+                    print(f"execute_fn {wg_execute_fn_name} is invalid")
+                    raise
+
+                # bind a new method to the RayWorkerGroup
+                func = func_generator(
+                    self,
+                    method_name,
+                    dispatch_fn=dispatch_fn,
+                    collect_fn=collect_fn,
+                    execute_fn=execute_fn,
+                    blocking=blocking,
+                )
+
+                try:
+                    setattr(self, method_name, func)
+                except Exception:
+                    raise ValueError(f"Fail to set method_name {method_name}")

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/ray/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from .base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, create_colocated_worker_cls
+
+
+__all__ = ["RayClassWithInitArgs", "RayResourcePool", "RayWorkerGroup", "create_colocated_worker_cls"]

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/single_controller/ray/base.py

@@ -0,0 +1,493 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import os
+import random
+import re
+import string
+import time
+from typing import Any, Dict, List, Optional, Tuple
+from unittest.mock import patch
+
+import ray
+from ray.actor import ActorHandle
+from ray.experimental.state.api import get_actor
+from ray.util import list_named_actors
+from ray.util.placement_group import PlacementGroup, placement_group
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy
+
+from ..base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup
+from ..base.decorator import MAGIC_ATTR
+
+
+__all__ = ["Worker"]
+
+
+def get_random_string(length: int) -> str:
+    letters_digits = string.ascii_letters + string.digits
+    return "".join(random.choice(letters_digits) for _ in range(length))
+
+
+def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):
+    def func(*args, **kwargs):
+        args, kwargs = dispatch_fn(self, *args, **kwargs)
+        output = execute_fn(method_name, *args, **kwargs)
+        if blocking:
+            output = ray.get(output)
+        output = collect_fn(self, output)
+        return output
+
+    return func
+
+
+def sort_placement_group_by_node_ip(pgs: List[PlacementGroup]) -> List[PlacementGroup]:
+    """
+    Sort the placement groups by node ip, all bundles in a single placement group should be on the same node.
+
+    FSDPCheckpointManager saves sharded model states and optimizer states in local storage, which requires RANK
+    to be consistent across nodes when resume from checkpoint.
+
+    With this function, if there's only one resource pool and there's no node change, RANK should be consistent
+    across nodes in multiple ray jobs, even if the whole ray cluster is restarted.
+    """
+    node_ip = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes()}
+    pg_ip = {}
+    for pg in pgs:
+        specs = ray._private.state.state.placement_group_table(pg.id)
+        # all bunles should be on the same node
+        node_id = specs["bundles_to_node_id"][0]
+        pg_ip[pg.id] = node_ip[node_id]
+
+    return sorted(pgs, key=lambda pg: pg_ip[pg.id])
+
+
+class RayResourcePool(ResourcePool):
+    def __init__(
+        self,
+        process_on_nodes: List[int] = None,
+        use_gpu: bool = True,
+        name_prefix: str = "",
+        max_colocate_count: int = 5,
+        detached: bool = False,
+    ) -> None:
+        super().__init__(process_on_nodes, max_colocate_count)
+        self.use_gpu = use_gpu
+        # print(f"in RayProcessDispatchConfiguration: name_prefix = {name_prefix}")
+        self.name_prefix = name_prefix
+        self.pgs = None
+        self.detached = detached
+
+    def get_placement_groups(self, strategy: str = "STRICT_PACK", name: Optional[str] = None) -> List[PlacementGroup]:
+        if self.pgs is not None:
+            return self.pgs
+
+        pg_name_prefix = (
+            name if name else f"{self.name_prefix}verl_group_{'_'.join([str(count) for count in self._store])}:"
+        )
+        # print(f"pg_name_prefix = {pg_name_prefix}")
+        pg_scheme = [
+            [
+                {"CPU": self.max_colocate_count, "GPU": 1} if self.use_gpu else {"CPU": self.max_colocate_count}
+                for _ in range(process_count)
+            ]
+            for process_count in self._store
+        ]
+
+        lifetime = "detached" if self.detached else None
+
+        pgs = [
+            placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime)
+            for idx, bundles in enumerate(pg_scheme)
+        ]
+
+        ray.get([pg.ready() for pg in pgs])
+
+        self.pgs = pgs
+        return pgs
+
+
+def extract_pg_from_exist(
+    resource_pools: Dict[str, RayResourcePool], src_role_names: List[str], resource_pool: RayResourcePool
+) -> List[PlacementGroup]:
+    src_pgs = [
+        pg
+        for role_name, resource_pool in resource_pools.items()
+        for pg in resource_pool.get_placement_groups()
+        if role_name in src_role_names
+    ]
+
+    sorted_src_pgs = sorted(src_pgs, key=lambda pg: pg.bundle_count, reverse=True)
+    sorted_process_on_nodes = sorted([(val, idx) for idx, val in enumerate(resource_pool.store)], reverse=True)
+
+    unsorted_pgs: List[Tuple[int, PlacementGroup]] = []
+    searching_idx = 0
+    for request_process, original_idx in sorted_process_on_nodes:
+        assert searching_idx < len(sorted_src_pgs), f"no enough nodes for request: searching {searching_idx} th node"
+        assert request_process <= sorted_src_pgs[searching_idx].bundle_count, (
+            f"requesting {request_process} processes, bundle count cannot satisfy"
+        )
+        unsorted_pgs.append((original_idx, sorted_src_pgs[searching_idx]))
+        searching_idx += 1
+
+    return [pg for _, pg in sorted(unsorted_pgs)]
+
+
+def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool:
+    assert rp1.use_gpu == rp2.use_gpu, "Both RayResourcePool must either use_gpu or not"
+    assert rp1.max_colocate_count == rp2.max_colocate_count, (
+        "Both RayResourcePool must has the same max_colocate_count"
+    )
+    assert rp1.n_gpus_per_node == rp2.n_gpus_per_node, "Both RayResourcePool must has the same n_gpus_per_node"
+    assert rp1.detached == rp2.detached, "Detached ResourcePool cannot be merged with non-detached ResourcePool"
+
+    new_store = rp1.store + rp2.store
+
+    merged = RayResourcePool(new_store, rp1.use_gpu, f"{rp1.name_prefix}_{rp2.name_prefix}")
+    merged.pgs = rp1.get_placement_groups() + rp2.get_placement_groups()
+
+    return merged
+
+
+class RayClassWithInitArgs(ClassWithInitArgs):
+    def __init__(self, cls, *args, **kwargs) -> None:
+        # self._options = kwargs.pop('options', dict())
+        super().__init__(cls, *args, **kwargs)
+        self._options = {}
+        self._additional_resource = {}
+
+    def set_additional_resource(self, additional_resource):
+        self._additional_resource = additional_resource
+
+    def update_options(self, options: Dict):
+        self._options.update(options)
+
+    def __call__(
+        self,
+        placement_group: PlacementGroup,
+        placement_group_bundle_idx: int,
+        use_gpu: bool = True,
+        num_gpus: int = 1,
+        sharing_with: Worker = None,
+    ) -> Any:
+        if sharing_with is not None:
+            target_node_id = ray.get(sharing_with.get_node_id.remote())
+            cuda_visible_devices = ray.get(sharing_with.get_cuda_visible_devices.remote())
+            options = {"scheduling_strategy": NodeAffinitySchedulingStrategy(node_id=target_node_id, soft=False)}
+            return self.cls.options(**options).remote(
+                *self.args, cuda_visible_devices=cuda_visible_devices, **self.kwargs
+            )
+
+        options = {
+            "scheduling_strategy": PlacementGroupSchedulingStrategy(
+                placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_idx
+            )
+        }
+        options.update(self._options)
+
+        if use_gpu:
+            options["num_gpus"] = num_gpus
+
+        if len(self._additional_resource) > 1:
+            for k, v in self._additional_resource.items():
+                options[k] = v
+
+        # print("cls:", self.cls)
+        # print("args: ", self.args)
+        # print("kwargs: ", self.kwargs)
+        return self.cls.options(**options).remote(*self.args, **self.kwargs)
+
+
+class RayWorkerGroup(WorkerGroup):
+    def __init__(
+        self,
+        resource_pool: RayResourcePool = None,
+        ray_cls_with_init: RayClassWithInitArgs = None,
+        bin_pack: bool = True,
+        name_prefix: str = None,
+        detached: bool = False,
+        worker_names: List[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(resource_pool=resource_pool, **kwargs)
+        self.ray_cls_with_init = ray_cls_with_init
+        self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix
+
+        if worker_names is not None:
+            assert self._is_init_with_detached_workers
+            self._worker_names = worker_names
+
+        if self._is_init_with_detached_workers:
+            self._init_with_detached_workers(worker_names=worker_names)
+        else:
+            self._init_with_resource_pool(
+                resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, bin_pack=bin_pack, detached=detached
+            )
+
+        if ray_cls_with_init is not None:
+            self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)
+
+    def _is_worker_alive(self, worker: ActorHandle) -> bool:
+        worker_state_dict = get_actor(worker._actor_id.hex())
+        return worker_state_dict.get("state", "undefined") == "ALIVE" if worker_state_dict is not None else False
+
+    def _init_with_detached_workers(self, worker_names: List[str]) -> None:
+        workers = [ray.get_actor(name=name) for name in worker_names]
+        self._workers = workers
+        self._world_size = len(worker_names)
+
+    def _init_with_resource_pool(
+        self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, bin_pack: bool, detached: bool
+    ):
+        use_gpu = resource_pool.use_gpu
+
+        strategy = "PACK"
+        if bin_pack:
+            strategy = "STRICT_PACK"
+
+        pgs = resource_pool.get_placement_groups(strategy=strategy)
+        world_size = resource_pool.world_size
+        self._world_size = world_size
+        # cia.add_kwarg("_world_size", world_size)
+        num_gpus = 1 / resource_pool.max_colocate_count
+
+        rank = -1
+        local_world_size = resource_pool.store[0]
+        for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)):
+            assert local_world_size <= pg.bundle_count, f"when generating for {self.name_prefix}, for the "
+            for local_rank in range(local_world_size):
+                rank += 1
+
+                # we pass in environment variable at option so that Worker can use environment variable to set
+                env_vars = {
+                    "WORLD_SIZE": str(world_size),
+                    "RANK": str(rank),
+                    "WG_PREFIX": self.name_prefix,
+                    "WG_BACKEND": "ray",
+                    "RAY_LOCAL_WORLD_SIZE": str(local_world_size),
+                    "RAY_LOCAL_RANK": str(local_rank),
+                }
+                if rank != 0:
+                    env_vars["MASTER_ADDR"] = self._master_addr
+                    env_vars["MASTER_PORT"] = self._master_port
+
+                cia_name = type(ray_cls_with_init.cls).__name__
+                match = re.search(r"ActorClass\(([^)]+)\)", cia_name)  # ray.remote(Obj) -> "ActorClass(Obj)"
+                cia_name = match.group(1) if match else cia_name  # "ActorClass(Obj)" -> "Obj"
+                name = f"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}"  # e.g. Worker_2:5
+
+                ray_cls_with_init.update_options({"runtime_env": {"env_vars": env_vars}, "name": name})
+
+                if detached:
+                    ray_cls_with_init.update_options({"lifetime": "detached"})
+
+                # create a worker
+                worker = ray_cls_with_init(
+                    placement_group=pg, placement_group_bundle_idx=local_rank, use_gpu=use_gpu, num_gpus=num_gpus
+                )
+                self._workers.append(worker)
+                self._worker_names.append(name)
+
+                if rank == 0:
+                    register_center_actor = None
+                    for _ in range(120):
+                        if f"{self.name_prefix}_register_center" not in list_named_actors():
+                            time.sleep(1)
+                        else:
+                            register_center_actor = ray.get_actor(f"{self.name_prefix}_register_center")
+                            break
+                    assert register_center_actor is not None, (
+                        f"failed to get register_center_actor: {self.name_prefix}_register_center in {list_named_actors(all_namespaces=True)}"
+                    )
+                    rank_zero_info = ray.get(register_center_actor.get_rank_zero_info.remote())
+                    self._master_addr, self._master_port = rank_zero_info["MASTER_ADDR"], rank_zero_info["MASTER_PORT"]
+                    # print(f"rank_zero_info: {rank_zero_info}")
+                    # print(f"master_addr: {self._master_addr}, master_port: {self._master_port}")
+
+    @property
+    def worker_names(self):
+        return self._worker_names
+
+    @classmethod
+    def from_detached(cls, worker_names=None, ray_cls_with_init=None):
+        worker_group = cls(
+            resource_pool=None, ray_cls_with_init=ray_cls_with_init, name_prefix=None, worker_names=worker_names
+        )
+        return worker_group
+
+    def spawn(self, prefix_set):
+        """
+        spawn to a dictionary of worker groups, each with a subset of method with prefix.
+
+        """
+
+        def _rebind_actor_methods(worker_group, actor_name):
+            """
+            bind the method with actor_prefix to its original name
+            """
+            prefix: str = actor_name + "_"
+            for method_name in dir(worker_group):
+                if method_name.startswith(prefix):
+                    # only valid when Python >= 3.9
+                    original_method_name = method_name.removeprefix(prefix)
+                    method = getattr(worker_group, method_name)
+                    setattr(worker_group, original_method_name, method)
+
+        new_worker_group_dict = {}
+        for prefix in prefix_set:
+            new_worker_group = self.from_detached(
+                worker_names=self._worker_names, ray_cls_with_init=self.ray_cls_with_init
+            )
+
+            _rebind_actor_methods(new_worker_group, prefix)
+            new_worker_group_dict[prefix] = new_worker_group
+        return new_worker_group_dict
+
+    def execute_rank_zero_sync(self, method_name: str, *args, **kwargs):
+        return ray.get(self.execute_rank_zero_async(method_name, *args, **kwargs))
+
+    def execute_rank_zero_async(self, method_name: str, *args, **kwargs):
+        remote_call = getattr(self._workers[0], method_name)
+        return remote_call.remote(*args, **kwargs)
+
+    def execute_rank_zero(self, method_name: str, *args, **kwargs):
+        return self.execute_rank_zero_async(method_name, *args, **kwargs)
+
+    def execute_all(self, method_name: str, *args, **kwargs):
+        return self.execute_all_async(method_name, *args, **kwargs)
+
+    def execute_all_sync(self, method_name: str, *args, **kwargs):
+        return ray.get(self.execute_all_async(method_name, *args, **kwargs))
+
+    def execute_all_async(self, method_name: str, *args, **kwargs):
+        # Here we assume that if all the parameters in args and kwargs are lists,
+        # and the lengths of all these lists are the same as len(self._workers),
+        # then we will send each element in the list to the corresponding worker.
+        # print(f"execute_all_async: method {method_name}({args}, {kwargs})")
+        length = len(self._workers)
+        if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()):
+            if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()):
+                # print(f"splitting args and kwargs into {length} shards")
+                result = []
+                for i in range(length):
+                    sliced_args = tuple(arg[i] for arg in args)
+                    sliced_kwargs = {k: v[i] for k, v in kwargs.items()}
+                    remote_call = getattr(self._workers[i], method_name)
+                    result.append(remote_call.remote(*sliced_args, **sliced_kwargs))
+                return result
+
+        return [getattr(worker, method_name).remote(*args, **kwargs) for worker in self._workers]
+
+    @property
+    def master_address(self):
+        return self._master_addr
+
+    @property
+    def master_port(self):
+        return self._master_port
+
+    @property
+    def workers(self):
+        return self._workers
+
+    @property
+    def world_size(self):
+        return self._world_size
+
+
+"""
+Utilities that enables creating workers inside the same ray.Actor,
+with code written in separate ray.Actors.
+"""
+
+
+def _bind_workers_method_to_parent(cls, key, user_defined_cls):
+    """
+    Binds the methods of each worker to the WorkerDict.
+    Note that we only bind public methods that are decorated by register
+    """
+    for method_name in dir(user_defined_cls):
+        try:
+            method = getattr(user_defined_cls, method_name)
+            assert callable(method), f"{method_name} in {user_defined_cls} is not callable"
+        except Exception:
+            # if it is a property, it will fail because Class doesn't have instance property
+            continue
+
+        if hasattr(method, MAGIC_ATTR):
+
+            def generate_function(name):
+                def func(self, *args, **kwargs):
+                    # dispatch to the actual worker
+                    return getattr(self.worker_dict[key], name)(*args, **kwargs)
+
+                return func
+
+            func = generate_function(method_name)
+            # pass MAGIC_ATTR for outer worker group
+            setattr(func, MAGIC_ATTR, getattr(method, MAGIC_ATTR))
+            try:
+                method_name_with_prefix = key + "_" + method_name
+                setattr(cls, method_name_with_prefix, func)
+                # print(f'Binding {method_name_with_prefix}')
+            except Exception:
+                raise ValueError(f"Fail to set method_name {method_name}")
+
+
+def _unwrap_ray_remote(cls):
+    if hasattr(cls, "__ray_actor_class__"):
+        cls = cls.__ray_actor_class__
+    return cls
+
+
+def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]):
+    """
+    This function should return a class instance that delegates the calls to every
+    cls in cls_dict
+    """
+    cls_dict = {}
+    init_args_dict = {}
+    worker_cls = None
+    for key, cls in class_dict.items():
+        if worker_cls is None:
+            worker_cls = cls.cls.__ray_actor_class__.__base__
+        else:
+            assert worker_cls == cls.cls.__ray_actor_class__.__base__, (
+                "the worker class should be the same when share the same process"
+            )
+        cls_dict[key] = cls.cls
+        init_args_dict[key] = {"args": cls.args, "kwargs": cls.kwargs}
+
+    assert cls_dict.keys() == init_args_dict.keys()
+
+    # TODO: create a class with customizable name
+    class WorkerDict(worker_cls):
+        def __init__(self):
+            super().__init__()
+            self.worker_dict = {}
+            for key, user_defined_cls in cls_dict.items():
+                user_defined_cls = _unwrap_ray_remote(user_defined_cls)
+                # directly instantiate the class without remote
+                with patch.dict(os.environ, {"DISABLE_WORKER_INIT": "1"}):
+                    self.worker_dict[key] = user_defined_cls(
+                        *init_args_dict[key].get("args", ()), **init_args_dict[key].get("kwargs", {})
+                    )
+
+    # now monkey-patch the methods from inner class to WorkerDict
+    for key, user_defined_cls in cls_dict.items():
+        user_defined_cls = _unwrap_ray_remote(user_defined_cls)
+        _bind_workers_method_to_parent(WorkerDict, key, user_defined_cls)
+
+    remote_cls = ray.remote(WorkerDict)
+    remote_cls = RayClassWithInitArgs(cls=remote_cls)
+    return remote_cls

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/config.py

@@ -0,0 +1,123 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+PPO config
+"""
+
+import os
+from dataclasses import asdict, dataclass, field, fields, is_dataclass
+from typing import Optional, Tuple
+
+from ..workers.config import WorkerConfig
+
+
+def recursive_post_init(dataclass_obj):
+    if hasattr(dataclass_obj, "post_init"):
+        dataclass_obj.post_init()
+
+    for attr in fields(dataclass_obj):
+        if is_dataclass(getattr(dataclass_obj, attr.name)):
+            recursive_post_init(getattr(dataclass_obj, attr.name))
+
+
+@dataclass
+class DataConfig:
+    train_files: str = ""
+    val_files: str = ""
+    prompt_key: str = "prompt"
+    answer_key: str = "answer"
+    image_key: str = "images"
+    max_prompt_length: int = 512
+    max_response_length: int = 512
+    rollout_batch_size: int = 512
+    val_batch_size: int = -1
+    format_prompt: Optional[str] = None
+    override_chat_template: Optional[str] = None
+    shuffle: bool = True
+    seed: int = 1
+    max_pixels: int = 4194304
+    min_pixels: int = 262144
+    filter_overlong_prompts: bool = True
+
+    def post_init(self):
+        if self.format_prompt is not None:
+            if os.path.exists(self.format_prompt):  # ray job uses absolute path
+                self.format_prompt = os.path.abspath(self.format_prompt)
+            else:
+                self.format_prompt = None
+
+
+@dataclass
+class AlgorithmConfig:
+    gamma: float = 1.0
+    lam: float = 1.0
+    adv_estimator: str = "grpo"
+    disable_kl: bool = False
+    use_kl_loss: bool = False
+    kl_penalty: str = "kl"
+    kl_coef: float = 1e-3
+    kl_type: str = "fixed"
+    kl_horizon: float = 0.0
+    kl_target: float = 0.0
+
+
+@dataclass
+class TrainerConfig:
+    total_epochs: int = 10
+    max_steps: Optional[int] = None
+    project_name: str = "easy_r1"
+    experiment_name: str = "demo"
+    logger: Tuple[str] = ("console", "wandb")
+    nnodes: int = 1
+    n_gpus_per_node: int = 8
+    critic_warmup: int = 0
+    val_freq: int = -1
+    val_before_train: bool = True
+    val_only: bool = False
+    val_generations_to_log: int = 0
+    save_freq: int = -1
+    save_limit: int = -1
+    save_checkpoint_path: Optional[str] = None
+    load_checkpoint_path: Optional[str] = None
+
+    def post_init(self):
+        if self.save_checkpoint_path is None:
+            self.save_checkpoint_path = os.path.join("checkpoints", self.project_name, self.experiment_name)
+
+        self.save_checkpoint_path = os.path.abspath(self.save_checkpoint_path)  # ray job uses absolute path
+        if self.load_checkpoint_path is not None:
+            self.load_checkpoint_path = os.path.abspath(self.load_checkpoint_path)
+
+
+@dataclass
+class PPOConfig:
+    data: DataConfig = field(default_factory=DataConfig)
+    worker: WorkerConfig = field(default_factory=WorkerConfig)
+    algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig)
+    trainer: TrainerConfig = field(default_factory=TrainerConfig)
+
+    def post_init(self):
+        self.worker.rollout.prompt_length = self.data.max_prompt_length
+        self.worker.rollout.response_length = self.data.max_response_length
+        self.worker.rollout.trust_remote_code = self.worker.actor.model.trust_remote_code
+        self.worker.actor.disable_kl = self.algorithm.disable_kl
+        self.worker.actor.use_kl_loss = self.algorithm.use_kl_loss
+        self.worker.actor.kl_penalty = self.algorithm.kl_penalty
+        self.worker.actor.kl_coef = self.algorithm.kl_coef
+
+    def deep_post_init(self):
+        recursive_post_init(self)
+
+    def to_dict(self):
+        return asdict(self)

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/core_algos.py

@@ -0,0 +1,428 @@
+# Copyright 2022 The HuggingFace Team
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+Core functions to implement PPO algorithms.
+The function implemented in this file should be used by trainer with different distributed strategies to
+implement PPO
+"""
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import TYPE_CHECKING, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from ..utils import torch_functional as VF
+
+
+if TYPE_CHECKING:
+    from .config import AlgorithmConfig
+
+
+class KLController(ABC):
+    kl_coef: float
+    """KL coefficient."""
+
+    @abstractmethod
+    def update(self, current_kl: float, n_steps: int) -> None:
+        """Update kl_coef according to current KL."""
+        ...
+
+
+class AdaptiveKLController(KLController):
+    """Adaptive KL controller described in: https://arxiv.org/pdf/1909.08593.pdf
+
+    Copied from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L54"""
+
+    def __init__(self, init_kl_coef: float, target_kl: float, horizon: float):
+        self.kl_coef = init_kl_coef
+        self.target = target_kl
+        self.horizon = horizon
+
+    def update(self, current_kl: float, n_steps: int) -> None:
+        target = self.target
+        proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2)
+        mult = 1 + proportional_error * n_steps / self.horizon
+        self.kl_coef *= mult
+
+
+class FixedKLController(KLController):
+    """Fixed KL controller.
+
+    Copeid from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L72"""
+
+    def __init__(self, init_kl_coef: float):
+        self.kl_coef = init_kl_coef
+
+    def update(self, current_kl: float, n_steps: int) -> None:
+        pass
+
+
+def get_kl_controller(algorithm_config: "AlgorithmConfig") -> KLController:
+    """Adapted from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L319"""
+    if algorithm_config.kl_type == "fixed":
+        kl_ctrl = FixedKLController(init_kl_coef=algorithm_config.kl_coef)
+    elif algorithm_config.kl_type == "adaptive":
+        assert algorithm_config.kl_horizon > 0, f"horizon must be larger than 0. Got {algorithm_config.kl_horizon}."
+        kl_ctrl = AdaptiveKLController(
+            init_kl_coef=algorithm_config.kl_coef,
+            target_kl=algorithm_config.kl_target,
+            horizon=algorithm_config.kl_horizon,
+        )
+    else:
+        raise ValueError(f"Unknown kl type: {algorithm_config.kl_type}.")
+
+    return kl_ctrl
+
+
+@torch.no_grad()
+def compute_gae_advantage_return(
+    token_level_rewards: torch.Tensor,
+    values: torch.Tensor,
+    response_mask: torch.Tensor,
+    gamma: torch.Tensor,
+    lam: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Adapted from https://github.com/huggingface/trl/blob/v0.16.0/trl/trainer/ppo_trainer.py#L513
+
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        values: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length). The token after eos tokens have mask zero.
+        gamma: `(float)`
+            discounted factor used in RL
+        lam: `(float)`
+            lambda value when computing Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438)
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    """
+    lastgaelam = 0
+    advantages_reversed = []
+    gen_len = token_level_rewards.shape[-1]
+    for t in reversed(range(gen_len)):
+        nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
+        delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
+        lastgaelam = delta + gamma * lam * lastgaelam
+        advantages_reversed.append(lastgaelam)
+
+    advantages = torch.stack(advantages_reversed[::-1], dim=1)
+    returns = advantages + values
+    advantages = VF.masked_whiten(advantages, response_mask)
+    return advantages, returns
+
+
+# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar.
+@torch.no_grad()
+def compute_grpo_outcome_advantage(
+    token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index: torch.Tensor, eps: float = 1e-6
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute advantage for GRPO, operating only on Outcome reward
+    (with only one scalar reward for each response).
+
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    """
+    scores = token_level_rewards.sum(dim=-1)
+    id2score = defaultdict(list)
+    id2mean, id2std = {}, {}
+
+    bsz = scores.shape[0]
+    for i in range(bsz):
+        id2score[index[i]].append(scores[i])
+
+    for idx in id2score:
+        assert len(id2score[idx]) > 1, "GRPO needs rollout.n > 1."
+        id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
+        id2std[idx] = torch.std(torch.tensor(id2score[idx]))
+
+    for i in range(bsz):
+        scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + eps)
+
+    returns = scores.unsqueeze(-1) * response_mask
+    return returns, returns
+
+
+@torch.no_grad()
+def compute_rloo_outcome_advantage(
+    token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute advantage for RLOO based on https://arxiv.org/abs/2402.14740
+
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    """
+    scores = token_level_rewards.sum(dim=-1)
+
+    id2score = defaultdict(list)
+    id2sum = {}
+    bsz = scores.shape[0]
+    for i in range(bsz):
+        id2score[index[i]].append(scores[i])
+
+    for idx in id2score:
+        id2sum[idx] = torch.sum(torch.tensor(id2score[idx]))
+
+    for i in range(bsz):
+        sample_num = len(id2score[index[i]])
+        assert sample_num > 1, "RLOO needs rollout.n > 1."
+        baseline = (id2sum[index[i]] - scores[i]) / (sample_num - 1)
+        scores[i] = scores[i] - baseline
+
+    returns = scores.unsqueeze(-1) * response_mask
+    return returns, returns
+
+
+@torch.no_grad()
+def compute_reinforce_plus_plus_outcome_advantage(
+    token_level_rewards: torch.Tensor, response_mask: torch.Tensor, gamma: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute advantage for REINFORCE++.
+    This implementation is based on the paper: https://arxiv.org/abs/2501.03262
+
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    """
+    returns = torch.zeros_like(token_level_rewards)
+    running_return = 0
+    for t in reversed(range(token_level_rewards.shape[1])):
+        running_return = token_level_rewards[:, t] + gamma * running_return
+        returns[:, t] = running_return
+        # Reset after EOS
+        running_return = running_return * response_mask[:, t]
+
+    advantages = VF.masked_whiten(returns, response_mask)
+    return advantages, returns
+
+
+@torch.no_grad()
+def compute_remax_outcome_advantage(
+    token_level_rewards: torch.Tensor, reward_baselines: torch.Tensor, response_mask: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute advantage for ReMax, operating only on Outcome reward
+    This implementation is based on the paper: https://arxiv.org/abs/2310.10505
+
+    (with only one scalar reward for each response).
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        reward_baselines: `(torch.Tensor)`
+            shape: (bs,)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+
+    """
+    scores = token_level_rewards.sum(dim=-1) - reward_baselines
+    returns = scores.unsqueeze(-1) * response_mask
+    return returns, returns
+
+
+def compute_rewards(
+    token_level_scores: torch.Tensor,
+    log_probs: torch.Tensor,
+    ref_log_probs: torch.Tensor,
+    kl_ratio: float,
+) -> torch.Tensor:
+    kl = log_probs - ref_log_probs
+    return token_level_scores - kl * kl_ratio
+
+
+def compute_policy_loss(
+    old_log_probs: torch.Tensor,
+    log_probs: torch.Tensor,
+    advantages: torch.Tensor,
+    response_mask: torch.Tensor,
+    clip_ratio_low: float,
+    clip_ratio_high: float,
+    clip_ratio_dual: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compute the policy loss.
+
+    Adapted from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L568
+
+    Args:
+        old_log_prob: `(torch.Tensor)`
+            shape: (bs, response_length)
+        log_prob: `(torch.Tensor)`
+            shape: (bs, response_length)
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        response_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+        clip_ratio_low: (float)
+            The lower clip range used in PPO. See https://arxiv.org/abs/1707.06347
+        clip_ratio_high: (float)
+            The higher clip range used in DAPO. See https://arxiv.org/pdf/2503.14476
+        clip_ratio_dual: (float)
+            The dual clip range used in Dual-clip PPO. See https://arxiv.org/pdf/1912.09729
+
+    Returns:
+        pg_loss: `a scalar torch.Tensor`
+            policy gradient loss computed via PPO
+        pg_clipfrac_higher: (float)
+            a float number indicating the fraction of policy gradient loss being clipped to a higher value
+        pg_clipfrac_lower: (float)
+            a float number indicating the fraction of policy gradient loss being clipped to a lower value
+        ppo_kl: (float)
+            a float number indicating the mean KL divergence between the old policy and the new policy
+
+    """
+    negative_approx_kl = log_probs - old_log_probs
+    # clamp the ratio before exp to avoid nan
+    # see: https://github.com/pytorch/pytorch/issues/10729
+    ratio = torch.exp(negative_approx_kl)
+    clipped_ratio = torch.exp(
+        torch.clamp(negative_approx_kl, np.log(1.0 - clip_ratio_low), np.log(1.0 + clip_ratio_high))
+    )
+
+    pg_loss = -advantages * ratio
+    pg_loss2 = -advantages * clipped_ratio
+    pg_loss3 = -advantages * clip_ratio_dual
+
+    clipped_pg_loss_higher = torch.max(pg_loss, pg_loss2)  # clip if pg_loss < pg_loss2
+    pg_clipfrac_higher = (pg_loss < pg_loss2).float()
+    clipped_pg_loss_lower = torch.min(clipped_pg_loss_higher, pg_loss3)  # clip if pg_loss > pg_loss3 and adv < 0
+    final_pg_loss = torch.where(advantages < 0, clipped_pg_loss_lower, clipped_pg_loss_higher)
+    pg_clipfrac_lower = (clipped_pg_loss_higher > pg_loss3).float() * (advantages < 0).float()
+
+    final_pg_loss = VF.masked_mean(final_pg_loss, response_mask)
+    pg_clipfrac_higher = VF.masked_mean(pg_clipfrac_higher, response_mask)
+    pg_clipfrac_lower = VF.masked_mean(pg_clipfrac_lower, response_mask)
+    ppo_kl = VF.masked_mean(-negative_approx_kl, response_mask)
+    return final_pg_loss, pg_clipfrac_higher, pg_clipfrac_lower, ppo_kl
+
+
+def compute_value_loss(
+    vpreds: torch.Tensor,
+    returns: torch.Tensor,
+    values: torch.Tensor,
+    action_mask: torch.Tensor,
+    cliprange_value: float,
+) -> Tuple[torch.Tensor, float]:
+    """Compute the value loss.
+
+    Adapted from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L556
+
+    Args:
+        vpreds (`torch.FloatTensor`):
+            Predicted values of the value head, shape (`batch_size`, `response_length`)
+        returns: (`torch.FloatTensor`):
+            Ground truth returns, shape (`batch_size`, `response_length`)
+        values (`torch.FloatTensor`):
+            Old values of value head, shape (`batch_size`, `response_length`)
+        action_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+        cliprange_value: (float)
+            The clip range for value net used in PPO. See https://arxiv.org/abs/1707.06347
+
+    Returns:
+        vf_loss: a scalar (`torch.FloatTensor`):
+            value function loss
+        vf_clipfrac: a float
+            The ratio of vf being clipped
+
+    """
+    vpredclipped = torch.clamp(vpreds, values - cliprange_value, values + cliprange_value)
+    vf_loss1 = torch.square(vpreds - returns)
+    vf_loss2 = torch.square(vpredclipped - returns)
+    vf_loss = 0.5 * VF.masked_mean(torch.max(vf_loss1, vf_loss2), action_mask)  # clip if vf_loss1 < vf_loss2
+    vf_clipfrac = VF.masked_mean((vf_loss1 < vf_loss2).float(), action_mask)
+    return vf_loss, vf_clipfrac
+
+
+def compute_kl(log_probs: torch.FloatTensor, ref_log_probs: torch.FloatTensor, kl_penalty: str) -> torch.Tensor:
+    """Compute KL divergence given log_probs and ref_log_probs.
+
+    Adapted from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L1150
+
+    Args:
+        log_probs: torch.Tensor
+        ref_log_probs: torch.Tensor
+        kl_penalty: str
+
+    Returns:
+        kl_div: torch.Tensor
+
+    """
+    log_probs, ref_log_probs = log_probs.float(), ref_log_probs.float()
+    if kl_penalty == "kl":
+        return log_probs - ref_log_probs
+
+    if kl_penalty == "abs":
+        return (log_probs - ref_log_probs).abs()
+
+    if kl_penalty == "mse":
+        return 0.5 * (log_probs - ref_log_probs).square()
+
+    # J. Schulman. Approximating kl divergence, 2020.
+    # URL http://joschu.net/blog/kl-approx.html
+    if kl_penalty == "low_var_kl":
+        kl = ref_log_probs - log_probs
+        kld = (kl.exp() - kl - 1).contiguous()
+        return torch.clamp(kld, min=-10, max=10)
+
+    if kl_penalty == "full":
+        return F.kl_div(ref_log_probs, log_probs, log_target=True, reduction="none").sum(-1)
+
+    raise NotImplementedError(f"Unknown KL penalty: {kl_penalty}.")

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/data_loader.py

@@ -0,0 +1,88 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from typing import Optional
+
+import torch
+from torch.utils.data import RandomSampler, SequentialSampler
+from torchdata.stateful_dataloader import StatefulDataLoader
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+from ..utils.dataset import RLHFDataset, collate_fn
+from .config import DataConfig
+
+
+def create_dataloader(config: DataConfig, tokenizer: PreTrainedTokenizer, processor: Optional[ProcessorMixin]) -> None:
+    train_dataset = RLHFDataset(
+        data_path=config.train_files,
+        tokenizer=tokenizer,
+        processor=processor,
+        prompt_key=config.prompt_key,
+        answer_key=config.answer_key,
+        image_key=config.image_key,
+        max_prompt_length=config.max_prompt_length,
+        truncation="right",
+        format_prompt=config.format_prompt,
+        min_pixels=config.min_pixels,
+        max_pixels=config.max_pixels,
+        filter_overlong_prompts=config.filter_overlong_prompts,
+    )
+    # use sampler for better ckpt resume
+    if config.shuffle:
+        train_dataloader_generator = torch.Generator()
+        train_dataloader_generator.manual_seed(config.seed)
+        print("train_dataset size:", len(train_dataset))
+        sampler = RandomSampler(data_source=train_dataset, generator=train_dataloader_generator)
+    else:
+        sampler = SequentialSampler(data_source=train_dataset)
+
+    train_dataloader = StatefulDataLoader(
+        dataset=train_dataset,
+        batch_size=config.rollout_batch_size,
+        sampler=sampler,
+        num_workers=8,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        drop_last=True,
+    )
+
+    val_dataset = RLHFDataset(
+        data_path=config.val_files,
+        tokenizer=tokenizer,
+        processor=processor,
+        prompt_key=config.prompt_key,
+        answer_key=config.answer_key,
+        image_key=config.image_key,
+        max_prompt_length=config.max_prompt_length,
+        truncation="right",
+        format_prompt=config.format_prompt,
+        min_pixels=config.min_pixels,
+        max_pixels=config.max_pixels,
+        filter_overlong_prompts=config.filter_overlong_prompts,
+    )
+    val_dataloader = StatefulDataLoader(
+        dataset=val_dataset,
+        batch_size=len(val_dataset) if config.val_batch_size == -1 else config.val_batch_size,
+        shuffle=False,
+        num_workers=8,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        drop_last=False,
+    )
+
+    assert len(train_dataloader) >= 1
+    assert len(val_dataloader) >= 1
+    print(f"Size of train dataloader: {len(train_dataloader)}")
+    print(f"Size of val dataloader: {len(val_dataloader)}")
+    return train_dataloader, val_dataloader

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/main.py

@@ -0,0 +1,130 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import json
+
+import ray
+from omegaconf import OmegaConf
+
+from ..single_controller.ray import RayWorkerGroup
+from ..utils.tokenizer import get_processor, get_tokenizer
+from ..workers.fsdp_workers import FSDPWorker
+from ..workers.reward import BatchFunctionRewardManager, SequentialFunctionRewardManager
+from .config import PPOConfig
+from .data_loader import create_dataloader
+from .ray_trainer import RayPPOTrainer, ResourcePoolManager, Role
+
+
+# please make sure main_task is not scheduled on head
+@ray.remote(num_cpus=1)
+class Runner:
+    """A runner for RL training."""
+
+    def run(self, config: PPOConfig):
+        # print config
+        print(json.dumps(config.to_dict(), indent=2))
+
+        # instantiate tokenizer
+        tokenizer = get_tokenizer(
+            config.worker.actor.model.model_path,
+            override_chat_template=config.data.override_chat_template,
+            trust_remote_code=config.worker.actor.model.trust_remote_code,
+            use_fast=True,
+        )
+        processor = get_processor(
+            config.worker.actor.model.model_path,
+            override_chat_template=config.data.override_chat_template,
+            trust_remote_code=config.worker.actor.model.trust_remote_code,
+            use_fast=True,
+        )
+
+        # define worker classes
+        ray_worker_group_cls = RayWorkerGroup
+        role_worker_mapping = {
+            Role.ActorRollout: ray.remote(FSDPWorker),
+            Role.Critic: ray.remote(FSDPWorker),
+            Role.RefPolicy: ray.remote(FSDPWorker),
+        }
+        global_pool_id = "global_pool"
+        resource_pool_spec = {
+            global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
+        }
+        mapping = {
+            Role.ActorRollout: global_pool_id,
+            Role.Critic: global_pool_id,
+            Role.RefPolicy: global_pool_id,
+        }
+        resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
+
+        if config.worker.reward.reward_type == "sequential":
+            RewardManager = SequentialFunctionRewardManager
+        elif config.worker.reward.reward_type == "batch":
+            RewardManager = BatchFunctionRewardManager
+        else:
+            raise NotImplementedError(f"Unknown reward type {config.worker.reward.reward_type}.")
+
+        RemoteRewardManager = ray.remote(RewardManager).options(num_cpus=config.worker.reward.num_cpus)
+        reward_fn = RemoteRewardManager.remote(config.worker.reward, tokenizer)
+        val_reward_fn = RemoteRewardManager.remote(config.worker.reward, tokenizer)
+
+        train_dataloader, val_dataloader = create_dataloader(config.data, tokenizer, processor)
+
+        trainer = RayPPOTrainer(
+            config=config,
+            tokenizer=tokenizer,
+            processor=processor,
+            train_dataloader=train_dataloader,
+            val_dataloader=val_dataloader,
+            role_worker_mapping=role_worker_mapping,
+            resource_pool_manager=resource_pool_manager,
+            ray_worker_group_cls=ray_worker_group_cls,
+            reward_fn=reward_fn,
+            val_reward_fn=val_reward_fn,
+        )
+        trainer.init_workers()
+        trainer.fit()
+
+
+def main():
+    cli_args = OmegaConf.from_cli()
+    default_config = OmegaConf.structured(PPOConfig())
+
+    if hasattr(cli_args, "config"):
+        config_path = cli_args.pop("config", None)
+        file_config = OmegaConf.load(config_path)
+        default_config = OmegaConf.merge(default_config, file_config)
+
+    ppo_config = OmegaConf.merge(default_config, cli_args)
+    ppo_config: PPOConfig = OmegaConf.to_object(ppo_config)
+    ppo_config.deep_post_init()
+
+    if not ray.is_initialized():
+        runtime_env = {
+            "env_vars": {
+                "TOKENIZERS_PARALLELISM": "true",
+                "NCCL_DEBUG": "WARN",
+                "VLLM_LOGGING_LEVEL": "WARN",
+                "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
+                "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False",
+                "PYTHONUNBUFFERED": "1",
+            }
+        }
+        ray.init(runtime_env=runtime_env)
+
+    runner = Runner.remote()
+    ray.get(runner.run.remote(ppo_config))
+
+
+if __name__ == "__main__":
+    main()

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/metrics.py

@@ -0,0 +1,120 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+
+from ..protocol import DataProto
+
+
+def reduce_metrics(metrics: Dict[str, List[Any]]) -> Dict[str, Any]:
+    return {key: np.mean(value) for key, value in metrics.items()}
+
+
+def compute_data_metrics(batch: DataProto, use_critic: bool = False) -> Dict[str, Any]:
+    sequence_score = batch.batch["token_level_scores"].sum(-1)
+    sequence_reward = batch.batch["token_level_rewards"].sum(-1)
+
+    advantages = batch.batch["advantages"]
+    returns = batch.batch["returns"]
+
+    max_response_length = batch.batch["responses"].size(-1)
+
+    prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool()
+    response_mask = batch.batch["attention_mask"][:, -max_response_length:].bool()
+
+    max_prompt_length = prompt_mask.size(-1)
+    prompt_length = prompt_mask.sum(-1).float()
+    response_length = response_mask.sum(-1).float()
+
+    valid_adv = torch.masked_select(advantages, response_mask)
+    valid_returns = torch.masked_select(returns, response_mask)
+
+    if use_critic:
+        values = batch.batch["values"]
+        valid_values = torch.masked_select(values, response_mask)
+        return_diff_var = torch.var(valid_returns - valid_values)
+        return_var = torch.var(valid_returns)
+
+    metrics = {
+        # score
+        "critic/score/mean": torch.mean(sequence_score).detach().item(),
+        "critic/score/max": torch.max(sequence_score).detach().item(),
+        "critic/score/min": torch.min(sequence_score).detach().item(),
+        # reward
+        "critic/rewards/mean": torch.mean(sequence_reward).detach().item(),
+        "critic/rewards/max": torch.max(sequence_reward).detach().item(),
+        "critic/rewards/min": torch.min(sequence_reward).detach().item(),
+        # adv
+        "critic/advantages/mean": torch.mean(valid_adv).detach().item(),
+        "critic/advantages/max": torch.max(valid_adv).detach().item(),
+        "critic/advantages/min": torch.min(valid_adv).detach().item(),
+        # returns
+        "critic/returns/mean": torch.mean(valid_returns).detach().item(),
+        "critic/returns/max": torch.max(valid_returns).detach().item(),
+        "critic/returns/min": torch.min(valid_returns).detach().item(),
+        **(
+            {
+                # values
+                "critic/values/mean": torch.mean(valid_values).detach().item(),
+                "critic/values/max": torch.max(valid_values).detach().item(),
+                "critic/values/min": torch.min(valid_values).detach().item(),
+                # vf explained var
+                "critic/vf_explained_var": (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
+            }
+            if use_critic
+            else {}
+        ),
+        # response length
+        "response_length/mean": torch.mean(response_length).detach().item(),
+        "response_length/max": torch.max(response_length).detach().item(),
+        "response_length/min": torch.min(response_length).detach().item(),
+        "response_length/clip_ratio": torch.mean(torch.eq(response_length, max_response_length).float())
+        .detach()
+        .item(),
+        # prompt length
+        "prompt_length/mean": torch.mean(prompt_length).detach().item(),
+        "prompt_length/max": torch.max(prompt_length).detach().item(),
+        "prompt_length/min": torch.min(prompt_length).detach().item(),
+        "prompt_length/clip_ratio": torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),
+    }
+    return metrics
+
+
+def compute_timing_metrics(batch: DataProto, timing_raw: Dict[str, float]) -> Dict[str, Any]:
+    num_response_tokens = torch.sum(batch.batch["response_mask"]).item()
+    num_overall_tokens = sum(batch.meta_info["global_token_num"])
+    num_tokens_of_section = {
+        **dict.fromkeys(["gen", "reward"], num_response_tokens),
+        **dict.fromkeys(["ref", "old", "values", "adv", "update_critic", "update_actor"], num_overall_tokens),
+    }
+    return {
+        **{f"timing_s/{name}": value for name, value in timing_raw.items()},
+        **{
+            f"timing_per_token_ms/{name}": timing_raw[name] * 1000 / num_tokens_of_section[name]
+            for name in set(num_tokens_of_section.keys()) & set(timing_raw.keys())
+        },
+    }
+
+
+def compute_throughout_metrics(batch: DataProto, timing_raw: Dict[str, float], num_gpus: int) -> Dict[str, Any]:
+    total_num_tokens = sum(batch.meta_info["global_token_num"])
+    time = timing_raw["step"]
+    return {
+        "perf/total_num_tokens": total_num_tokens,
+        "perf/time_per_step": time,
+        "perf/throughput": total_num_tokens / (time * num_gpus),
+    }

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/trainer/ray_trainer.py

@@ -0,0 +1,637 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""
+FSDP PPO Trainer with Ray-based single controller.
+This trainer supports model-agonistic model initialization with huggingface
+"""
+
+import os
+import uuid
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass, field
+from enum import Enum, IntEnum, auto
+from typing import Any, Dict, List, Optional, Type
+
+import numpy as np
+import ray
+import torch
+from ray.experimental.tqdm_ray import tqdm
+from torchdata.stateful_dataloader import StatefulDataLoader
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+from ..protocol import DataProto, pad_dataproto_to_divisor, unpad_dataproto
+from ..single_controller.base import Worker
+from ..single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
+from ..single_controller.ray.base import create_colocated_worker_cls
+from ..utils import torch_functional as VF
+from ..utils.checkpoint import CHECKPOINT_TRACKER, remove_obsolete_ckpt
+from ..utils.logger import Tracker
+from ..utils.py_functional import convert_dict_to_str, timer
+from ..utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance
+from ..workers.fsdp_workers import FSDPWorker
+from ..workers.reward import FunctionRewardManager
+from . import core_algos
+from .config import PPOConfig
+from .metrics import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics, reduce_metrics
+
+
+class Role(IntEnum):
+    """
+    To create more roles dynamically, you can subclass Role and add new members
+    """
+
+    Actor = auto()
+    Rollout = auto()
+    ActorRollout = auto()
+    Critic = auto()
+    RefPolicy = auto()
+    RewardModel = auto()
+    ActorRolloutRef = auto()
+
+
+class AdvantageEstimator(str, Enum):
+    """
+    Using an enumeration class to avoid spelling errors in adv_estimator
+    """
+
+    GAE = "gae"
+    GRPO = "grpo"
+    REINFORCE_PLUS_PLUS = "reinforce_plus_plus"
+    REMAX = "remax"
+    RLOO = "rloo"
+
+
+@dataclass
+class ResourcePoolManager:
+    """
+    Define a resource pool specification. Resource pool will be initialized first.
+    """
+
+    resource_pool_spec: dict[str, list[int]]
+    mapping: dict[Role, str]
+    resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict)
+
+    def create_resource_pool(self):
+        for resource_pool_name, process_on_nodes in self.resource_pool_spec.items():
+            # max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool
+            # For FSDP backend, we recommend using max_colocate_count=1 that merge all WorkerGroups into one.
+            # For Megatron backend, we recommend using max_colocate_count>1 that can utilize different WorkerGroup for differnt models
+            resource_pool = RayResourcePool(
+                process_on_nodes=process_on_nodes, use_gpu=True, max_colocate_count=1, name_prefix=resource_pool_name
+            )
+            self.resource_pool_dict[resource_pool_name] = resource_pool
+
+        self._check_resource_available()
+
+    def get_resource_pool(self, role: Role) -> RayResourcePool:
+        """Get the resource pool of the worker."""
+        return self.resource_pool_dict[self.mapping[role]]
+
+    def get_num_gpus(self) -> int:
+        """Get the number of gpus in this cluster."""
+        return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes])
+
+    def _check_resource_available(self):
+        """Check if the resource pool can be satisfied in this ray cluster."""
+        gpus_available = ray.available_resources().get("GPU", 0)
+        gpus_required = self.get_num_gpus()
+        if gpus_available < gpus_required:
+            raise ValueError(f"Total available GPUs {gpus_available} is less than total desired GPUs {gpus_required}.")
+
+
+def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.KLController, kl_penalty="kl"):
+    token_level_scores = data.batch["token_level_scores"]
+    batch_size = data.batch.batch_size[0]
+    response_mask = data.batch["response_mask"]
+
+    # compute kl between ref_policy and current policy
+    kld = core_algos.compute_kl(data.batch["old_log_probs"], data.batch["ref_log_probs"], kl_penalty=kl_penalty)
+    kld = kld * response_mask  # (batch_size, response_length)
+
+    data.batch["token_level_rewards"] = token_level_scores - kl_ctrl.kl_coef * kld
+
+    current_kl = VF.masked_mean(kld, mask=response_mask, dim=-1)  # average over sequence
+    current_kl = torch.mean(current_kl, dim=0).item()
+    metrics = {"critic/kl": current_kl, "critic/kl_coef": kl_ctrl.kl_coef}
+
+    # According to https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L880
+    kl_ctrl.update(current_kl=current_kl, n_steps=batch_size)
+    return data, metrics
+
+
+def compute_advantage(data: DataProto, adv_estimator: AdvantageEstimator, gamma: float = 1.0, lam: float = 1.0):
+    token_level_rewards = data.batch["token_level_rewards"]
+    response_mask = data.batch["response_mask"]
+    index = data.non_tensor_batch["uid"]
+    if adv_estimator == AdvantageEstimator.GAE:
+        values = data.batch["values"]
+        advantages, returns = core_algos.compute_gae_advantage_return(
+            token_level_rewards, values, response_mask, gamma, lam
+        )
+    elif adv_estimator == AdvantageEstimator.GRPO:
+        advantages, returns = core_algos.compute_grpo_outcome_advantage(token_level_rewards, response_mask, index)
+    elif adv_estimator == AdvantageEstimator.REINFORCE_PLUS_PLUS:
+        advantages, returns = core_algos.compute_reinforce_plus_plus_outcome_advantage(
+            token_level_rewards, response_mask, gamma
+        )
+    elif adv_estimator == AdvantageEstimator.REMAX:
+        reward_baselines = data.batch["reward_baselines"]
+        advantages, returns = core_algos.compute_remax_outcome_advantage(
+            token_level_rewards, reward_baselines, response_mask
+        )
+    elif adv_estimator == AdvantageEstimator.RLOO:
+        advantages, returns = core_algos.compute_rloo_outcome_advantage(token_level_rewards, response_mask, index)
+    else:
+        raise NotImplementedError
+
+    data.batch["advantages"] = advantages
+    data.batch["returns"] = returns
+    return data
+
+
+class RayPPOTrainer:
+    """
+    Note that this trainer runs on the driver process on a single CPU/GPU node.
+    """
+
+    def __init__(
+        self,
+        config: PPOConfig,
+        tokenizer: PreTrainedTokenizer,
+        processor: Optional[ProcessorMixin],
+        train_dataloader: StatefulDataLoader,
+        val_dataloader: StatefulDataLoader,
+        role_worker_mapping: dict[Role, Type[Worker]],
+        resource_pool_manager: ResourcePoolManager,
+        ray_worker_group_cls: Type[RayWorkerGroup] = RayWorkerGroup,
+        reward_fn: Optional[FunctionRewardManager] = None,
+        val_reward_fn: Optional[FunctionRewardManager] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.processor = processor
+        self.train_dataloader = train_dataloader
+        self.val_dataloader = val_dataloader
+        self.config = config
+        self.reward_fn = reward_fn
+        self.val_reward_fn = val_reward_fn
+
+        self.hybrid_engine = config.worker.hybrid_engine
+        if self.hybrid_engine:
+            assert Role.ActorRollout in role_worker_mapping, (
+                f"ActorRollout should be included in {role_worker_mapping.keys()}."
+            )
+        else:
+            raise NotImplementedError
+
+        self.role_worker_mapping = role_worker_mapping
+        self.resource_pool_manager = resource_pool_manager
+        self.use_reward_model = Role.RewardModel in role_worker_mapping
+        self.ray_worker_group_cls = ray_worker_group_cls
+
+        # define KL control
+        if Role.RefPolicy in role_worker_mapping and not config.algorithm.disable_kl:
+            self.use_reference_policy = True
+            self.kl_ctrl = core_algos.get_kl_controller(config.algorithm)
+        else:
+            self.use_reference_policy = False
+            self.kl_ctrl = core_algos.FixedKLController(init_kl_coef=0.0)
+            print("KL is disabled, no KL metrics will be logged. Please set `kl_coef=0` to log KL metrics.")
+
+        if config.algorithm.adv_estimator == AdvantageEstimator.GAE:
+            self.use_critic = True
+        else:
+            self.use_critic = False
+
+        if config.algorithm.adv_estimator not in list(AdvantageEstimator):
+            raise NotImplementedError(f"Unknown advantage estimator: {config.algorithm.adv_estimator}.")
+
+        if config.data.rollout_batch_size % config.worker.actor.global_batch_size != 0:
+            raise ValueError("Rollout batch size must be divisible by actor global batch size.")
+
+        if (
+            config.data.rollout_batch_size * config.worker.rollout.n
+        ) % config.worker.actor.micro_batch_size_per_device_for_experience != 0:
+            raise ValueError(
+                "Rollout batch size * rollout.n must be divisible by actor micro batch size for experience."
+            )
+
+        if self.use_critic:
+            if config.data.rollout_batch_size % config.worker.critic.global_batch_size != 0:
+                raise ValueError("Rollout batch size must be divisible by critic global batch size.")
+
+            if (
+                config.data.rollout_batch_size * config.worker.rollout.n
+            ) % config.worker.critic.micro_batch_size_per_device_for_experience != 0:
+                raise ValueError(
+                    "Rollout batch size * rollout.n must be divisible by critic micro batch size for experience."
+                )
+
+        if (
+            config.algorithm.adv_estimator in (AdvantageEstimator.GRPO, AdvantageEstimator.RLOO)
+            and config.worker.rollout.n == 1
+        ):
+            raise ValueError("GRPO and RLOO algorithm need `config.worker.rollout.n > 1`.")
+
+        if config.trainer.max_steps is not None:
+            self.training_steps = config.trainer.max_steps
+        else:
+            self.training_steps = len(train_dataloader) * config.trainer.total_epochs
+
+        config.worker.actor.optim.training_steps = self.training_steps
+        config.worker.critic.optim.training_steps = self.training_steps
+        print(f"Total training steps: {self.training_steps}")
+
+    def _maybe_log_val_generations(
+        self, inputs: List[str], outputs: List[str], labels: List[str], scores: List[float]
+    ) -> None:
+        """Log a table of validation samples"""
+        if self.config.trainer.val_generations_to_log <= 0:
+            return
+
+        # Create tuples of (input, output, score) and sort by input text
+        samples = list(zip(inputs, outputs, labels, scores))
+        samples.sort(key=lambda x: x[0])  # Sort by input text
+
+        # Use fixed random seed for deterministic shuffling
+        rng = np.random.RandomState(42)
+        rng.shuffle(samples)
+
+        samples = samples[: self.config.trainer.val_generations_to_log]
+        self.logger.log_generation(samples, self.global_step)
+
+    def _validate(self) -> Dict[str, Any]:
+        reward_tensor_lst = []
+        # Lists to collect samples for the table
+        sample_inputs, sample_outputs, sample_labels, sample_scores = [], [], [], []
+        reward_metrics_lst = defaultdict(list)
+        for batch_dict in self.val_dataloader:
+            test_batch = DataProto.from_single_dict(batch_dict)
+            # Store original inputs
+            input_ids = test_batch.batch["input_ids"]
+            input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
+            sample_inputs.extend(input_texts)
+
+            if "multi_modal_data" in test_batch.non_tensor_batch.keys():
+                test_gen_batch = test_batch.pop(
+                    batch_keys=["input_ids", "attention_mask", "position_ids"],
+                    non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data"],
+                )
+            else:
+                test_gen_batch = test_batch.pop(
+                    batch_keys=["input_ids", "attention_mask", "position_ids"],
+                    non_tensor_batch_keys=["raw_prompt_ids"],
+                )
+
+            test_gen_batch.meta_info = self.config.worker.rollout.val_override_config
+            test_gen_batch.meta_info.update({
+                "min_pixels": self.config.data.min_pixels,
+                "max_pixels": self.config.data.max_pixels,
+            })
+            test_gen_batch, pad_size = pad_dataproto_to_divisor(test_gen_batch, self.actor_rollout_wg.world_size)
+            test_output_gen_batch = self.actor_rollout_wg.generate_sequences(test_gen_batch)
+            test_output_gen_batch = unpad_dataproto(test_output_gen_batch, pad_size=pad_size)
+
+            # Store generated outputs
+            output_ids = test_output_gen_batch.batch["responses"]
+            output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]
+            sample_outputs.extend(output_texts)
+            sample_labels.extend(test_batch.non_tensor_batch["ground_truth"].tolist())
+            test_batch = test_batch.union(test_output_gen_batch)
+
+            # evaluate using reward_function
+            reward_tensor, reward_metrics = ray.get(self.val_reward_fn.compute_reward.remote(test_batch))
+
+            # Store scores
+            scores = reward_tensor.sum(-1).cpu().tolist()
+            sample_scores.extend(scores)
+
+            reward_tensor_lst.append(reward_tensor)
+            for key, value in reward_metrics.items():
+                reward_metrics_lst[key].extend(value)
+
+        self._maybe_log_val_generations(sample_inputs, sample_outputs, sample_labels, sample_scores)
+        reward_score = torch.cat(reward_tensor_lst, dim=0).sum(-1).mean().item()
+        val_reward_metrics = {f"val/{key}_reward": value for key, value in reduce_metrics(reward_metrics_lst).items()}
+        return {"val/reward_score": reward_score, **val_reward_metrics}
+
+    def init_workers(self) -> None:
+        """Init resource pool and worker group"""
+        self.resource_pool_manager.create_resource_pool()
+        self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
+
+        # create actor and rollout
+        if self.hybrid_engine:
+            resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout)
+            actor_rollout_cls = RayClassWithInitArgs(
+                cls=self.role_worker_mapping[Role.ActorRollout], config=self.config.worker, role="actor_rollout"
+            )
+            self.resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls
+        else:
+            raise NotImplementedError
+
+        # create critic
+        if self.use_critic:
+            resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)
+            critic_cls = RayClassWithInitArgs(
+                cls=self.role_worker_mapping[Role.Critic], config=self.config.worker, role="critic"
+            )
+            self.resource_pool_to_cls[resource_pool]["critic"] = critic_cls
+
+        # create reference policy if needed
+        if self.use_reference_policy:
+            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)
+            ref_policy_cls = RayClassWithInitArgs(
+                self.role_worker_mapping[Role.RefPolicy], config=self.config.worker, role="ref"
+            )
+            self.resource_pool_to_cls[resource_pool]["ref"] = ref_policy_cls
+
+        # create a reward model if reward_fn is None
+        if self.use_reward_model:
+            # we create a RM here
+            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
+            rm_cls = RayClassWithInitArgs(
+                cls=self.role_worker_mapping[Role.RewardModel], config=self.config.worker, role="reward"
+            )
+            self.resource_pool_to_cls[resource_pool]["rm"] = rm_cls
+
+        # initialize WorkerGroup
+        # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,
+        # you should not use `create_colocated_worker_cls`. Instead, directly pass different resource pool to different worker groups.
+        # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
+        all_wg: Dict[str, FSDPWorker] = {}
+        self.wg_dicts = []
+        for resource_pool, class_dict in self.resource_pool_to_cls.items():
+            worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
+            wg_dict = self.ray_worker_group_cls(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)
+            spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
+            all_wg.update(spawn_wg)
+            # keep the referece of WorkerDict to support ray >= 2.31. Ref: https://github.com/ray-project/ray/pull/45699
+            self.wg_dicts.append(wg_dict)
+
+        if self.use_critic:
+            self.critic_wg = all_wg["critic"]
+            self.critic_wg.init_model()
+
+        if self.use_reference_policy:
+            self.ref_policy_wg = all_wg["ref"]
+            self.ref_policy_wg.init_model()
+
+        if self.use_reward_model:
+            self.rm_wg = all_wg["rm"]
+            self.rm_wg.init_model()
+
+        # we should create rollout at the end so that vllm can have a better estimation of kv cache memory
+        self.actor_rollout_wg = all_wg["actor_rollout"]
+        self.actor_rollout_wg.init_model()
+
+    def _save_checkpoint(self) -> None:
+        # path: {save_checkpoint_path}/global_step_{global_step}/{actor,critic}
+        remove_obsolete_ckpt(
+            self.config.trainer.save_checkpoint_path, self.global_step, self.config.trainer.save_limit
+        )
+        folder_path = os.path.join(self.config.trainer.save_checkpoint_path, f"global_step_{self.global_step}")
+        actor_path = os.path.join(folder_path, "actor")
+        self.actor_rollout_wg.save_checkpoint(actor_path)
+
+        if self.use_critic:
+            critic_path = os.path.join(folder_path, "critic")
+            self.critic_wg.save_checkpoint(critic_path)
+
+        dataloader_path = os.path.join(folder_path, "dataloader.pt")
+        dataloader_state_dict = self.train_dataloader.state_dict()
+        torch.save(dataloader_state_dict, dataloader_path)
+
+        last_global_step_path = os.path.join(self.config.trainer.save_checkpoint_path, CHECKPOINT_TRACKER)
+        with open(last_global_step_path, "w") as f:
+            f.write(str(self.global_step))
+
+    def _load_checkpoint(self) -> None:
+        if self.config.trainer.load_checkpoint_path is None:
+            return
+
+        if "global_step_" not in self.config.trainer.load_checkpoint_path.strip(os.path.sep).split(os.path.sep)[-1]:
+            raise ValueError("`load_checkpoint_path` should end with `global_step_*`.")
+
+        print(f"Load from checkpoint: {self.config.trainer.load_checkpoint_path}.")
+        self.global_step = int(self.config.trainer.load_checkpoint_path.strip(os.path.sep).split("global_step_")[-1])
+        actor_path = os.path.join(self.config.trainer.load_checkpoint_path, "actor")
+        self.actor_rollout_wg.load_checkpoint(actor_path)
+        if self.use_critic:
+            critic_path = os.path.join(self.config.trainer.load_checkpoint_path, "critic")
+            self.critic_wg.load_checkpoint(critic_path)
+
+        dataloader_path = os.path.join(self.config.trainer.load_checkpoint_path, "dataloader.pt")
+        if os.path.exists(dataloader_path):
+            dataloader_state_dict = torch.load(dataloader_path, weights_only=False)
+            self.train_dataloader.load_state_dict(dataloader_state_dict)
+        else:
+            print(f"No dataloader state found at {dataloader_path}, will start from scratch.")
+
+    def _balance_batch(self, batch: DataProto, metrics: Dict[str, Any], logging_prefix: str = "global_seqlen") -> None:
+        """Reorder the data on single controller such that each dp rank gets similar total tokens"""
+        attention_mask = batch.batch["attention_mask"]
+        batch_size = attention_mask.shape[0]
+        global_seqlen_lst = batch.batch["attention_mask"].view(batch_size, -1).sum(-1).tolist()  # (train_batch_size,)
+        world_size = self.actor_rollout_wg.world_size
+        global_partition_lst = get_seqlen_balanced_partitions(
+            global_seqlen_lst, k_partitions=world_size, equal_size=True
+        )
+        # reorder based on index. The data will be automatically equally partitioned by dispatch function
+        global_idx = torch.tensor([j for partition in global_partition_lst for j in partition])
+        batch.reorder(global_idx)
+        global_balance_stats = log_seqlen_unbalance(
+            seqlen_list=global_seqlen_lst, partitions=global_partition_lst, prefix=logging_prefix
+        )
+        metrics.update(global_balance_stats)
+
+    def fit(self):
+        """
+        The training loop of PPO.
+        The driver process only need to call the compute functions of the worker group through RPC to construct the PPO dataflow.
+        The light-weight advantage computation is done on the driver process.
+        """
+        self.logger = Tracker(loggers=self.config.trainer.logger, config=self.config.to_dict())
+        self.global_step = 0
+        val_metrics: Optional[Dict[str, Any]] = None
+
+        # load checkpoint before doing anything
+        self._load_checkpoint()
+
+        # perform validation before training
+        # currently, we only support validation using the reward_function.
+        if self.val_reward_fn is not None and self.config.trainer.val_before_train:
+            val_metrics = self._validate()
+            self.logger.log(data=val_metrics, step=self.global_step)
+            if self.config.trainer.val_only:
+                return
+
+        for _ in tqdm(range(self.config.trainer.total_epochs), desc="Epoch", position=0):
+            for batch_dict in tqdm(self.train_dataloader, desc="Running step", position=1):
+                self.global_step += 1
+                if self.global_step > self.training_steps:
+                    break
+
+                metrics, timing_raw = {}, {}
+                batch: DataProto = DataProto.from_single_dict(batch_dict)
+
+                # pop those keys for generation
+                if "multi_modal_data" in batch.non_tensor_batch.keys():
+                    print(batch.non_tensor_batch['multi_modal_data'].shape)
+                    gen_batch = batch.pop(
+                        batch_keys=["input_ids", "attention_mask", "position_ids"],
+                        non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data"],
+                    )
+                    gen_batch.meta_info.update({
+                        "min_pixels": self.config.data.min_pixels,
+                        "max_pixels": self.config.data.max_pixels,
+                    })
+                else:
+                    gen_batch = batch.pop(
+                        batch_keys=["input_ids", "attention_mask", "position_ids"],
+                        non_tensor_batch_keys=["raw_prompt_ids"],
+                    )
+
+                with timer("step", timing_raw):
+                    # generate a batch
+                    with timer("gen", timing_raw):  # wg: worker group
+                        gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)
+
+                    if self.config.algorithm.adv_estimator == "remax":
+                        with timer("gen_max", timing_raw):
+                            gen_baseline_batch = deepcopy(gen_batch)
+                            gen_baseline_batch.meta_info["temperature"] = 0
+                            gen_baseline_batch.meta_info["n"] = 1
+                            gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
+
+                            batch = batch.union(gen_baseline_output)
+                            reward_baseline_tensor, _ = ray.get(self.reward_fn.compute_reward.remote(batch))
+                            reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
+
+                            batch.pop(batch_keys=list(gen_baseline_output.batch.keys()))
+                            batch.batch["reward_baselines"] = reward_baseline_tensor
+                            del gen_baseline_batch, gen_baseline_output
+
+                    batch.non_tensor_batch["uid"] = np.array(
+                        [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object
+                    )
+                    # repeat to align with repeated responses in rollout
+                    batch = batch.repeat(repeat_times=self.config.worker.rollout.n, interleave=True)
+                    batch = batch.union(gen_batch_output)
+
+                    # balance the number of valid tokens on each dp rank.
+                    # Note that this breaks the order of data inside the batch.
+                    # Please take care when you implement group based adv computation such as GRPO and rloo
+                    self._balance_batch(batch, metrics=metrics)
+
+                    # compute global_valid tokens
+                    batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
+
+                    # compute reward
+                    with timer("reward", timing_raw):
+                        reward_ref = self.reward_fn.compute_reward.remote(batch)
+
+                    # recompute old_log_probs
+                    with timer("old", timing_raw):
+                        old_log_probs = self.actor_rollout_wg.compute_log_probs(batch)
+                        batch = batch.union(old_log_probs)
+
+                    # compute ref_log_probs
+                    if self.use_reference_policy:
+                        with timer("ref", timing_raw):
+                            ref_log_probs = self.ref_policy_wg.compute_ref_log_probs(batch)
+                            batch = batch.union(ref_log_probs)
+
+                    # compute values
+                    if self.use_critic:
+                        with timer("values", timing_raw):
+                            values = self.critic_wg.compute_values(batch)
+                            batch = batch.union(values)
+
+                    with timer("adv", timing_raw):
+                        # get token level scores
+                        reward_tensor, reward_metrics = ray.get(reward_ref)
+                        batch.batch["token_level_scores"] = reward_tensor
+                        reward_metrics = {f"reward/{k}": v for k, v in reduce_metrics(reward_metrics).items()}
+                        metrics.update(reward_metrics)
+
+                        # apply kl penalty if available
+                        if not self.config.algorithm.use_kl_loss and self.use_reference_policy:
+                            # apply kl penalty to reward
+                            batch, kl_metrics = apply_kl_penalty(batch, self.kl_ctrl, self.config.algorithm.kl_penalty)
+                            metrics.update(kl_metrics)
+                        else:
+                            batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
+
+                        # compute advantages, executed on the driver process
+                        batch = compute_advantage(
+                            batch,
+                            adv_estimator=self.config.algorithm.adv_estimator,
+                            gamma=self.config.algorithm.gamma,
+                            lam=self.config.algorithm.lam,
+                        )
+
+                    # update critic
+                    if self.use_critic:
+                        with timer("update_critic", timing_raw):
+                            critic_output = self.critic_wg.update_critic(batch)
+
+                        critic_metrics = reduce_metrics(critic_output.non_tensor_batch)
+                        metrics.update(critic_metrics)
+
+                    # update actor
+                    if self.config.trainer.critic_warmup <= self.global_step:
+                        with timer("update_actor", timing_raw):
+                            actor_output = self.actor_rollout_wg.update_actor(batch)
+
+                        actor_metrics = reduce_metrics(actor_output.non_tensor_batch)
+                        metrics.update(actor_metrics)
+
+                    # validate
+                    if (
+                        self.val_reward_fn is not None
+                        and self.config.trainer.val_freq > 0
+                        and self.global_step % self.config.trainer.val_freq == 0
+                    ):
+                        with timer("validation", timing_raw):
+                            val_metrics = self._validate()
+
+                        metrics.update(val_metrics)
+
+                    if self.config.trainer.save_freq > 0 and self.global_step % self.config.trainer.save_freq == 0:
+                        with timer("save_checkpoint", timing_raw):
+                            self._save_checkpoint()
+
+                # collect metrics
+                num_gpus = self.resource_pool_manager.get_num_gpus()
+                metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
+                metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
+                metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, num_gpus=num_gpus))
+
+                self.logger.log(data=metrics, step=self.global_step)
+
+        # perform validation after training
+        if self.val_reward_fn is not None:
+            if (
+                val_metrics is None
+                or self.config.trainer.val_freq <= 0
+                or self.global_step % self.config.trainer.val_freq != 0
+            ):
+                val_metrics = self._validate()
+                self.logger.log(data=val_metrics, step=self.global_step)
+
+            print(f"Final validation metrics: {convert_dict_to_str(val_metrics)}")
+
+        if self.config.trainer.save_freq <= 0 or self.global_step % self.config.trainer.save_freq != 0:
+            self._save_checkpoint()

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/utils/checkpoint/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+from .checkpoint_manager import CHECKPOINT_TRACKER, remove_obsolete_ckpt
+
+
+__all__ = ["CHECKPOINT_TRACKER", "remove_obsolete_ckpt"]

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/utils/checkpoint/checkpoint_manager.py

@@ -0,0 +1,160 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import os
+import random
+import re
+import shutil
+import tempfile
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from filelock import FileLock
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+
+CHECKPOINT_TRACKER = "latest_global_step.txt"
+
+
+class BaseCheckpointManager(ABC):
+    """
+    A checkpoint manager that saves and loads
+    - model
+    - optimizer
+    - lr_scheduler
+    - extra_states
+    in a SPMD way.
+
+    We save
+    - sharded model states and optimizer states
+    - full lr_scheduler states
+    - huggingface tokenizer and config for ckpt merge
+    """
+
+    def __init__(
+        self,
+        model: FSDP,
+        optimizer: torch.optim.Optimizer,
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+        processing_class: Union[PreTrainedTokenizer, ProcessorMixin],
+    ):
+        self.model = model
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        self.processing_class = processing_class
+
+        assert isinstance(self.model, FSDP)
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
+
+    @abstractmethod
+    def load_checkpoint(self, *args, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def save_checkpoint(self, *args, **kwargs):
+        raise NotImplementedError
+
+    @staticmethod
+    def local_mkdir(path: str) -> str:
+        if not os.path.isabs(path):
+            working_dir = os.getcwd()
+            path = os.path.join(working_dir, path)
+
+        # Using hash value of path as lock file name to avoid long file name
+        lock_filename = f"ckpt_{hash(path) & 0xFFFFFFFF:08x}.lock"
+        lock_path = os.path.join(tempfile.gettempdir(), lock_filename)
+
+        try:
+            with FileLock(lock_path, timeout=60):
+                os.makedirs(path, exist_ok=True)
+        except Exception as e:
+            print(f"Warning: Failed to acquire lock for {path}: {e}")
+            os.makedirs(path, exist_ok=True)  # even if the lock is not acquired, try to create the directory
+
+        return path
+
+    @staticmethod
+    def get_rng_state() -> Dict[str, Any]:
+        rng_state = {
+            "cpu": torch.get_rng_state(),
+            "cuda": torch.cuda.get_rng_state(),
+            "numpy": np.random.get_state(),
+            "random": random.getstate(),
+        }
+        return rng_state
+
+    @staticmethod
+    def load_rng_state(rng_state: Dict[str, Any]):
+        torch.set_rng_state(rng_state["cpu"])
+        torch.cuda.set_rng_state(rng_state["cuda"])
+        np.random.set_state(rng_state["numpy"])
+        random.setstate(rng_state["random"])
+
+
+def find_latest_ckpt_path(path: Optional[str] = None, directory_format: str = "global_step_{}") -> Optional[str]:
+    if path is None:
+        return None
+
+    tracker_file = get_checkpoint_tracker_filename(path)
+    if not os.path.exists(tracker_file):
+        print("Checkpoint tracker file does not exist: %s", tracker_file)
+        return None
+
+    with open(tracker_file, "rb") as f:
+        iteration = int(f.read().decode())
+
+    ckpt_path = os.path.join(path, directory_format.format(iteration))
+    if not os.path.exists(ckpt_path):
+        print("Checkpoint does not exist: %s", ckpt_path)
+        return None
+
+    print("Found checkpoint: %s", ckpt_path)
+    return ckpt_path
+
+
+def get_checkpoint_tracker_filename(root_path: str) -> str:
+    """
+    Tracker file rescords the latest chckpoint during training to restart from.
+    """
+    return os.path.join(root_path, CHECKPOINT_TRACKER)
+
+
+def remove_obsolete_ckpt(path: str, global_step: int, save_limit: int = -1, directory_format: str = "global_step_{}"):
+    """
+    Remove the obsolete checkpoints that exceed the save_limit.
+    """
+    if save_limit <= 0:
+        return
+
+    if not os.path.exists(path):
+        return
+
+    pattern = re.escape(directory_format).replace(r"\{\}", r"(\d+)")
+    ckpt_folders = []
+    for folder in os.listdir(path):
+        if match := re.match(pattern, folder):
+            step = int(match.group(1))
+            if step < global_step:
+                ckpt_folders.append((step, folder))
+
+    ckpt_folders.sort(reverse=True)
+    for _, folder in ckpt_folders[save_limit - 1 :]:
+        folder_path = os.path.join(path, folder)
+        shutil.rmtree(folder_path, ignore_errors=True)
+        print(f"Removed obsolete checkpoint: {folder_path}")

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/utils/checkpoint/fsdp_checkpoint_manager.py

@@ -0,0 +1,113 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import os
+from typing import Optional, Union
+
+import torch
+import torch.distributed as dist
+from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
+
+from .checkpoint_manager import BaseCheckpointManager
+
+
+class FSDPCheckpointManager(BaseCheckpointManager):
+    """
+    A checkpoint manager that saves and loads
+    - model
+    - optimizer
+    - lr_scheduler
+    - extra_states
+    in a SPMD way.
+
+    We save
+    - sharded model states and optimizer states
+    - full lr_scheduler states
+    - huggingface tokenizer and config for ckpt merge
+    """
+
+    def __init__(
+        self,
+        model: FSDP,
+        optimizer: torch.optim.Optimizer,
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+        processing_class: Union[PreTrainedTokenizer, ProcessorMixin],
+    ):
+        super().__init__(model, optimizer, lr_scheduler, processing_class)
+
+    def load_checkpoint(self, path: Optional[str] = None):
+        if path is None:
+            return
+
+        # every rank download its own checkpoint
+        model_path = os.path.join(path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt")
+        optim_path = os.path.join(path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt")
+        extra_path = os.path.join(path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt")
+        print(f"[rank-{self.rank}]: Loading model from {os.path.abspath(model_path)}.")
+        print(f"[rank-{self.rank}]: Loading optimizer from {os.path.abspath(optim_path)}.")
+        print(f"[rank-{self.rank}]: Loading extra_state from {os.path.abspath(extra_path)}.")
+        model_state_dict = torch.load(model_path, weights_only=False)
+        optim_state_dict = torch.load(optim_path, weights_only=False)
+        extra_state_dict = torch.load(extra_path, weights_only=False)
+
+        state_dict_options = StateDictOptions(cpu_offload=True)
+        set_state_dict(
+            model=self.model,
+            optimizers=self.optimizer,
+            model_state_dict=model_state_dict,
+            optim_state_dict=optim_state_dict,
+            options=state_dict_options,
+        )
+        self.lr_scheduler.load_state_dict(extra_state_dict["lr_scheduler"])
+
+        # recover random state
+        if "rng" in extra_state_dict:
+            self.load_rng_state(extra_state_dict["rng"])
+
+    def save_checkpoint(self, path: str):
+        path = self.local_mkdir(path)
+        dist.barrier()
+
+        # every rank will save its own model and optim shard
+        state_dict_options = StateDictOptions(cpu_offload=True)
+        model_state_dict, optim_state_dict = get_state_dict(self.model, self.optimizer, options=state_dict_options)
+        extra_state_dict = {
+            "lr_scheduler": self.lr_scheduler.state_dict(),
+            "rng": self.get_rng_state(),
+        }
+        model_path = os.path.join(path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt")
+        optim_path = os.path.join(path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt")
+        extra_path = os.path.join(path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt")
+
+        print(f"[rank-{self.rank}]: Saving model to {os.path.abspath(model_path)}.")
+        print(f"[rank-{self.rank}]: Saving optimizer to {os.path.abspath(optim_path)}.")
+        print(f"[rank-{self.rank}]: Saving extra_state to {os.path.abspath(extra_path)}.")
+        torch.save(model_state_dict, model_path)
+        torch.save(optim_state_dict, optim_path)
+        torch.save(extra_state_dict, extra_path)
+
+        # wait for everyone to dump to local
+        dist.barrier()
+
+        if self.rank == 0:
+            hf_path = os.path.join(path, "huggingface")
+            os.makedirs(hf_path, exist_ok=True)
+            assert isinstance(self.model._fsdp_wrapped_module, PreTrainedModel)
+            self.model._fsdp_wrapped_module.config.save_pretrained(hf_path)
+            self.model._fsdp_wrapped_module.generation_config.save_pretrained(hf_path)
+            self.processing_class.save_pretrained(hf_path)
+
+        dist.barrier()

LLaVA-NeXT/llava-critic-r1/EasyR1/verl/utils/dataset.py

@@ -0,0 +1,235 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+
+import math
+import os
+from collections import defaultdict
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from jinja2 import Template
+from PIL import Image
+from PIL.Image import Image as ImageObject
+from PIL.Image import Image as PILImage
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer, ProcessorMixin
+from PIL import Image, JpegImagePlugin
+
+from ..models.transformers.qwen2_vl import get_rope_index
+from . import torch_functional as VF
+
+
+def collate_fn(features: List[Dict[str, Any]]) -> Dict[str, Any]:
+    tensors = defaultdict(list)
+    non_tensors = defaultdict(list)
+    for feature in features:
+        for key, value in feature.items():
+            if isinstance(value, torch.Tensor):
+                tensors[key].append(value)
+            else:
+                non_tensors[key].append(value)
+
+    for key, value in tensors.items():
+        tensors[key] = torch.stack(value, dim=0)
+
+    for key, value in non_tensors.items():
+        non_tensors[key] = np.array(value, dtype=object)
+
+    return {**tensors, **non_tensors}
+
+
+
+def process_image(image: Union[Dict[str, Any], ImageObject, str], min_pixels: int, max_pixels: int) -> ImageObject:
+    if isinstance(image, dict):
+        if 'path' in image:
+            image = Image.open(image['path']).copy()
+        else:
+            image = Image.open(BytesIO(image["bytes"]))
+    elif isinstance(image, bytes):
+        image = Image.open(BytesIO(image))
+    elif isinstance(image, list):
+        image = Image.open(image[0]).copy()
+        # image = image[0]
+    elif isinstance(image, PILImage):
+        image = image.copy()
+    else:
+        # image = image
+        raise TypeError(f'Unsupported image type: {type(image)}')
+        
+    if (image.width * image.height) > max_pixels:
+        resize_factor = math.sqrt(max_pixels / (image.width * image.height))
+        width, height = int(image.width * resize_factor), int(image.height * resize_factor)
+        image = image.resize((width, height))
+
+    if (image.width * image.height) < min_pixels:
+        resize_factor = math.sqrt(min_pixels / (image.width * image.height))
+        width, height = int(image.width * resize_factor), int(image.height * resize_factor)
+        image = image.resize((width, height))
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    return image
+
+def process_video(
+    video: str, min_pixels: Optional[int], max_pixels: Optional[int], video_fps: float, return_fps: bool = False
+) -> Union[list[ImageObject], tuple[list[ImageObject], list[float]]]:
+    vision_info = {"video": video, "min_pixels": min_pixels, "max_pixels": max_pixels, "fps": video_fps}
+    return fetch_video(vision_info, return_video_sample_fps=return_fps)
+
+class RLHFDataset(Dataset):
+    """
+    We assume the dataset contains a column that contains prompts and other information
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer: PreTrainedTokenizer,
+        processor: Optional[ProcessorMixin],
+        prompt_key: str = "prompt",
+        answer_key: str = "answer",
+        image_key: str = "images",
+        max_prompt_length: int = 1024,
+        truncation: str = "error",
+        format_prompt: Optional[str] = None,
+        max_pixels: Optional[int] = None,
+        min_pixels: Optional[int] = None,
+        filter_overlong_prompts: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.processor = processor
+        self.prompt_key = prompt_key
+        self.answer_key = answer_key
+        self.image_key = image_key
+        self.max_prompt_length = max_prompt_length
+        self.truncation = truncation
+        self.max_pixels = max_pixels
+        self.min_pixels = min_pixels
+        self.filter_overlong_prompts = filter_overlong_prompts
+
+        if "@" in data_path:
+            data_path, data_split = data_path.split("@")
+        else:
+            data_split = "train"
+
+        if os.path.isdir(data_path):
+            # when we use dataset builder, we should always refer to the train split
+            self.dataset = load_dataset("parquet", data_dir=data_path, split="train")
+        elif os.path.isfile(data_path):
+            self.dataset = load_dataset("parquet", data_files=data_path, split="train")
+        else:
+            # load remote dataset from huggingface hub
+            self.dataset = load_dataset(data_path, split=data_split)
+
+        self.format_prompt = None
+        if format_prompt:
+            with open(format_prompt, encoding="utf-8") as f:
+                self.format_prompt = f.read()
+
+        if self.filter_overlong_prompts:
+            self.dataset = self.dataset.filter(self._filter_overlong_prompts, desc="Filtering overlong prompts")
+
+    def _build_messages(self, example: Dict[str, Any]) -> List[Dict[str, Any]]:
+        prompt_str: str = example[self.prompt_key]
+        if self.format_prompt:
+            format_prompt = Template(self.format_prompt.strip())
+            prompt_str = format_prompt.render(content=prompt_str)
+
+        if self.image_key in example:
+            # https://huggingface.co/docs/transformers/en/tasks/image_text_to_text
+            content_list = []
+            for i, content in enumerate(prompt_str.split("<image>")):
+                if i != 0:
+                    content_list.append({"type": "image"})
+
+                if content:
+                    content_list.append({"type": "text", "text": content})
+
+            return [{"role": "user", "content": content_list}]
+        else:
+            return [{"role": "user", "content": prompt_str}]
+
+    def _filter_overlong_prompts(self, example: Dict[str, Any]) -> bool:
+        messages = self._build_messages(example)
+        processing_class = self.processor if self.processor is not None else self.tokenizer
+        return (
+            len(processing_class.apply_chat_template(messages, add_generation_prompt=True)) <= self.max_prompt_length
+        )
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        example: dict = self.dataset[index]
+        messages = self._build_messages(example)
+
+        if self.image_key in example:
+            prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+            raw_image_data = example.pop(self.image_key)
+            if not isinstance(raw_image_data, (list, tuple)):
+                        raw_image_data = [raw_image_data]
+
+            images = [
+                process_image(image, min_pixels=self.min_pixels, max_pixels=self.max_pixels)
+                for image in raw_image_data
+            ]
+            model_inputs = self.processor(images, [prompt], add_special_tokens=False, return_tensors="pt")
+            input_ids = model_inputs.pop("input_ids")[0]
+            attention_mask = model_inputs.pop("attention_mask")[0]
+            example["multi_modal_data"] = {"image": raw_image_data}
+        else:
+            prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+            model_inputs = self.tokenizer([prompt], add_special_tokens=False, return_tensors="pt")
+            input_ids = model_inputs.pop("input_ids")[0]
+            attention_mask = model_inputs.pop("attention_mask")[0]
+
+        if self.processor is not None and self.processor.image_processor.__class__.__name__ == "Qwen2VLImageProcessor":
+            # qwen2vl mrope
+            position_ids = get_rope_index(
+                self.processor,
+                input_ids=input_ids,
+                image_grid_thw=model_inputs.get("image_grid_thw"),
+                attention_mask=attention_mask,
+            )  # (3, seq_length)
+        else:
+            position_ids = torch.clip(attention_mask.cumsum(dim=0) - 1, min=0, max=None)  # (seq_length,)
+
+        input_ids, attention_mask, position_ids = VF.postprocess_data(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            max_length=self.max_prompt_length,
+            pad_token_id=self.tokenizer.pad_token_id,
+            left_pad=True,
+            truncation=self.truncation,
+        )
+        raw_prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+        if len(raw_prompt_ids) > self.max_prompt_length:
+            if self.truncation == "left":
+                raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length :]
+            elif self.truncation == "right":
+                raw_prompt_ids = raw_prompt_ids[: self.max_prompt_length]
+            elif self.truncation == "error":
+                raise RuntimeError(f"Prompt length {len(raw_prompt_ids)} is longer than {self.max_prompt_length}.")
+
+        example["input_ids"] = input_ids
+        example["attention_mask"] = attention_mask
+        example["position_ids"] = position_ids
+        example["raw_prompt_ids"] = raw_prompt_ids
+        example["ground_truth"] = example.pop(self.answer_key)
+        return example