From 555a539373c38eb20861dc9215834ee894a5dc9b Mon Sep 17 00:00:00 2001
From: lihanwen <602590163@qq.com>
Date: Thu, 18 Jun 2026 14:59:22 +0800
Subject: [PATCH] feat: add preprocessing scripts for Open-R1 and
 TinyLLaVA-Video datasets, and implement reward scoring functions

---
 examples/data_preprocess/openr1mm.py          |  98 +++++++++
 .../data_preprocess/tinyllava_video_r1.py     | 191 ++++++++++++++++++
 .../run_qwen3_5_2b_openr1_fsdp.sh             | 161 +++++++++++++++
 .../grpo_trainer/run_qwen3_5_2b_video_fsdp.sh | 163 +++++++++++++++
 verl/workers/engine_workers.py                |   1 +
 5 files changed, 614 insertions(+)
 create mode 100644 examples/data_preprocess/openr1mm.py
 create mode 100644 examples/data_preprocess/tinyllava_video_r1.py
 create mode 100644 examples/grpo_trainer/run_qwen3_5_2b_openr1_fsdp.sh
 create mode 100644 examples/grpo_trainer/run_qwen3_5_2b_video_fsdp.sh

diff --git a/examples/data_preprocess/openr1mm.py b/examples/data_preprocess/openr1mm.py
new file mode 100644
index 00000000000..30c0c1e30a2
--- /dev/null
+++ b/examples/data_preprocess/openr1mm.py
@@ -0,0 +1,98 @@
+# 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.
+"""
+Preprocess the lmms-lab/multimodal-open-r1-8k-verified dataset to parquet format.
+
+Images are kept as raw bytes (no decode, no resize).
+"""
+
+import argparse
+import os
+
+import datasets
+
+from verl.utils.hdfs_io import copy, makedirs
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--local_save_dir", default="~/data/openr1mm", help="The save directory for the preprocessed dataset."
+    )
+    parser.add_argument("--hdfs_dir", default=None)
+    args = parser.parse_args()
+
+    data_source = "lmms-lab/multimodal-open-r1-8k-verified"
+    dataset = datasets.load_dataset(data_source)
+
+    instruction = (
+        "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 enclosed within <answer> </answer> tags."
+    )
+
+    def make_map_fn(split):
+        def process_fn(example, idx):
+            problem = example.pop("problem")
+            solution = example.pop("solution")
+            img = example.pop("image")
+
+            prompt_content = f"<image>\n{problem}\n\n{instruction}"
+
+            # Keep image as raw bytes dict to avoid lossy re-encoding.
+            # The Qwen VL processor handles resize at runtime.
+            if isinstance(img, dict) and "bytes" in img:
+                image_data = img
+            elif isinstance(img, bytes):
+                image_data = {"bytes": img}
+            else:
+                image_data = img
+
+            data = {
+                "data_source": data_source,
+                "prompt": [{"role": "user", "content": prompt_content}],
+                "images": [image_data],
+                "ability": "math",
+                "reward_model": {"style": "rule", "ground_truth": solution},
+                "extra_info": {
+                    "split": split,
+                    "index": idx,
+                    "question": problem,
+                    "answer": solution,
+                },
+            }
+            return data
+
+        return process_fn
+
+    full_dataset = dataset["train"]
+    full_dataset = full_dataset.cast_column("image", datasets.Image(decode=False))
+    split_dataset = full_dataset.train_test_split(test_size=0.1, seed=42)
+
+    train_dataset = split_dataset["train"].map(function=make_map_fn("train"), with_indices=True, num_proc=8)
+    test_dataset = split_dataset["test"].map(function=make_map_fn("test"), with_indices=True, num_proc=8)
+
+    columns = ["data_source", "prompt", "images", "ability", "reward_model", "extra_info"]
+    train_dataset = train_dataset.select_columns(columns)
+    test_dataset = test_dataset.select_columns(columns)
+
+    local_save_dir = os.path.expanduser(args.local_save_dir)
+    os.makedirs(local_save_dir, exist_ok=True)
+
+    train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet"))
+    test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet"))
+
+    if args.hdfs_dir is not None:
+        makedirs(args.hdfs_dir)
+        copy(src=local_save_dir, dst=args.hdfs_dir)
diff --git a/examples/data_preprocess/tinyllava_video_r1.py b/examples/data_preprocess/tinyllava_video_r1.py
new file mode 100644
index 00000000000..d89310178bb
--- /dev/null
+++ b/examples/data_preprocess/tinyllava_video_r1.py
@@ -0,0 +1,191 @@
+# 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.
+"""
+Preprocess Zhang199/TinyLLaVA-Video-R1-training-data to verl parquet format.
+
+  - Prompt: "<video>{problem}" (problem already has inline options)
+  - Video: absolute file path
+  - Label: solution ("<answer>X</answer>")
+  - Uses verl-standard inline instruction (think/answer format).
+
+Usage:
+    # Step 1: Download
+    export HF_ENDPOINT=https://hf-mirror.com
+    hf download Zhang199/TinyLLaVA-Video-R1-training-data \\
+        --repo-type dataset --local-dir ~/data/tinyllava-video-r1
+
+    # Step 2: Extract videos
+    unzip ~/data/tinyllava-video-r1/NextQA.zip -d ~/data/tinyllava-video-r1/
+
+    # Step 3: Preprocess
+    python examples/data_preprocess/tinyllava_video_r1.py \\
+        --data_dir ~/data/tinyllava-video-r1 \\
+        --local_save_dir ~/data/tinyllava_video_r1
+"""
+
+import argparse
+import json
+import os
+import sys
+from typing import Optional
+
+import datasets
+
+from verl.utils.hdfs_io import copy, makedirs
+
+DATA_SOURCE = "Zhang199/TinyLLaVA-Video-R1-training-data"
+
+# Inline instruction: think/answer format for video QA.
+INSTRUCTION = (
+    "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 a single option letter (e.g., A, B, C, D, E) "
+    "enclosed within <answer> </answer> tags."
+)
+
+
+def build_prompt_text(problem: str) -> str:
+    """Build prompt with video placeholder: "<video>{problem}".
+
+    The JSONL problem field already contains inline options:
+      "What animal is shown?\nOptions:\nA. owl.\nB. sheeps.\n..."
+    """
+    return f"<video>\n{problem}\n\n{INSTRUCTION}"
+
+
+def make_map_fn(
+    data_source: str,
+    video_dir: str,
+    split: str,
+    video_fps: Optional[float] = None,
+    video_max_frames: Optional[int] = None,
+):
+    """Factory function following verl geo3k/openr1mm closure pattern."""
+
+    def process_fn(example, idx):
+        problem = example["problem"]
+        solution = example["solution"]  # already "<answer>X</answer>"
+
+        # Resolve video path from video_dir + video_filename.
+        # JSONL paths are "./NextQA/NExTVideo/..." → strip "./" prefix.
+        video_rel = example["video_filename"].lstrip("./")
+        video_path = os.path.join(video_dir, video_rel)
+        if not os.path.exists(video_path):
+            print(f"[WARN] Video file not found: {video_path}", file=sys.stderr)
+
+        prompt_content = build_prompt_text(problem)
+
+        # Video sampling params (fps=1, max_frames=32)
+        video_entry = {"video": video_path}
+        if video_fps is not None:
+            video_entry["fps"] = video_fps
+        if video_max_frames is not None:
+            video_entry["max_frames"] = video_max_frames
+
+        return {
+            "data_source": data_source,
+            "prompt": [{"role": "user", "content": prompt_content}],
+            "videos": [video_entry],
+            "ability": "video_qa",
+            "reward_model": {"style": "rule", "ground_truth": solution},
+            "extra_info": {
+                "split": split,
+                "index": idx,
+                "question": problem,
+                "answer": solution,
+                "video_path": video_path,
+            },
+        }
+
+    return process_fn
+
+
+def load_jsonl(path: str) -> list[dict]:
+    data = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                data.append(json.loads(line))
+    return data
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Preprocess TinyLLaVA-Video-R1 to verl parquet.")
+    parser.add_argument("--data_dir", type=str, default=None, help="Downloaded dataset directory.")
+    parser.add_argument("--local_save_dir", default="~/data/tinyllava_video_r1", help="Output directory.")
+    parser.add_argument("--hdfs_dir", default=None)
+    parser.add_argument("--video_fps", type=float, default=1, help="Video sampling FPS (default: 1)")
+    parser.add_argument("--video_max_frames", type=int, default=32, help="Max frames per video (default: 32)")
+    args = parser.parse_args()
+
+    if not args.data_dir:
+        parser.error("--data_dir is required")
+
+    # ---- Load ----
+    jsonl_path = os.path.join(args.data_dir, "nextqa_0-30s.jsonl")
+    if not os.path.exists(jsonl_path):
+        print(f"[ERROR] Not found: {jsonl_path}")
+        sys.exit(1)
+
+    print(f"Loading: {jsonl_path}")
+    data = load_jsonl(jsonl_path)
+    print(f"  {len(data)} samples")
+
+    # Sanity check
+    s0 = data[0]
+    print(f"  First sample problem: {s0['problem'][:80]}...")
+    print(f"  First sample video:   {s0['video_filename']}")
+    print(f"  First sample solution: {s0['solution']}")
+
+    # ---- Video directory ----
+    # JSONL video_filename includes "NextQA/" prefix (e.g. "./NextQA/NExTVideo/..."),
+    # so video_dir must be the dataset root, not dataset_root/NextQA.
+    video_dir = args.data_dir
+
+    # ---- Convert + 90/10 split (same as openr1mm.py) ----
+    full_dataset = datasets.Dataset.from_list(data)
+    split_dataset = full_dataset.train_test_split(test_size=0.1, seed=42)
+
+    train_map_fn = make_map_fn(
+        DATA_SOURCE, video_dir, "train", video_fps=args.video_fps, video_max_frames=args.video_max_frames
+    )
+    test_map_fn = make_map_fn(
+        DATA_SOURCE, video_dir, "test", video_fps=args.video_fps, video_max_frames=args.video_max_frames
+    )
+    train_dataset = split_dataset["train"].map(function=train_map_fn, with_indices=True, num_proc=4)
+    test_dataset = split_dataset["test"].map(function=test_map_fn, with_indices=True, num_proc=4)
+
+    columns = ["data_source", "prompt", "videos", "ability", "reward_model", "extra_info"]
+    train_dataset = train_dataset.select_columns(columns)
+    test_dataset = test_dataset.select_columns(columns)
+
+    # ---- Save ----
+    save_dir = os.path.expanduser(args.local_save_dir)
+    os.makedirs(save_dir, exist_ok=True)
+
+    train_out = os.path.join(save_dir, "train.parquet")
+    test_out = os.path.join(save_dir, "test.parquet")
+
+    print(f"\nSaving train ({len(train_dataset)} samples) → {train_out}")
+    train_dataset.to_parquet(train_out)
+    print(f"Saving test  ({len(test_dataset)} samples) → {test_out}")
+    test_dataset.to_parquet(test_out)
+
+    if args.hdfs_dir:
+        makedirs(args.hdfs_dir)
+        copy(src=save_dir, dst=args.hdfs_dir)
+
+    print(f"\nDone! Train: {train_out}, Test: {test_out}")
diff --git a/examples/grpo_trainer/run_qwen3_5_2b_openr1_fsdp.sh b/examples/grpo_trainer/run_qwen3_5_2b_openr1_fsdp.sh
new file mode 100644
index 00000000000..21af97dce77
--- /dev/null
+++ b/examples/grpo_trainer/run_qwen3_5_2b_openr1_fsdp.sh
@@ -0,0 +1,161 @@
+#!/usr/bin/env bash
+# GRPO | Qwen3.5-2B | FSDP training | Open-R1 multimodal math dataset
+# dependency: GPU vllm==0.18.0, transformers@<cc7ab9be>
+# dependency: NPU vllm==0.18.0, vllm-ascend@<54879467>, transformers@<cc7ab9be>
+
+set -xeuo pipefail
+
+########################### user-adjustable ###########################
+# DEVICE is auto-detected by probing torch_npu; override only for special cases.
+DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)}
+INFER_BACKEND=${INFER_BACKEND:-vllm}
+PROJECT_NAME=${PROJECT_NAME:-GRPO-OpenR1MM}
+EXPERIMENT_NAME=${EXPERIMENT_NAME:-Qwen3.5-2B-GRPO-OpenR1MM}
+NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-}
+NNODES=${NNODES:-1}
+
+GEN_TP=${GEN_TP:-4}
+SP_SIZE=${SP_SIZE:-1}
+FSDP_SIZE=${FSDP_SIZE:-}
+ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.3}
+
+RAY_DATA_HOME=${RAY_DATA_HOME:-/data/verl}
+export RAY_TMPDIR=${RAY_TMPDIR:-/data/tmp/ray}
+MODEL_PATH=${MODEL_PATH:-"${HOME}/verl/models/Qwen3.5-2B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${PROJECT_NAME}/${EXPERIMENT_NAME}"}
+LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/logs/${PROJECT_NAME}/${EXPERIMENT_NAME}"}
+TRAIN_FILE=${TRAIN_FILE:-"${HOME}/verl/data/openr1mm/train.parquet"}
+TEST_FILE=${TEST_FILE:-"${HOME}/verl/data/openr1mm/test.parquet"}
+WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+########################### end user-adjustable ###########################
+
+########################### derived defaults ###########################
+n_devices_per_node=${NDEVICES_PER_NODE:-8}
+fsdp_size=${FSDP_SIZE:-8}
+
+case "${DEVICE}" in
+    gpu)
+        ;;
+    npu)
+        export HCCL_CONNECT_TIMEOUT=1500
+        export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+        export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+        export HCCL_EXEC_TIMEOUT=3600
+        n_devices_per_node=8
+        fsdp_size=8
+        ;;
+    *)
+        echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2
+        exit 1
+        ;;
+esac
+
+start_time=$(date +%Y%m%d)_$(date +%H%M%S)
+mkdir -p "${LOG_DIR}" "${RAY_TMPDIR}"
+
+########################### parameter arrays ###########################
+
+DATA=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+    data.train_files="${TRAIN_FILE}"
+    data.val_files="${TEST_FILE}"
+    data.train_batch_size=64
+    data.max_prompt_length=2048
+    data.max_response_length=1536
+    data.filter_overlong_prompts=True
+    data.filter_overlong_prompts_workers=8
+    data.truncation='error'
+    data.image_patch_size=16
+    data.image_key=images
+    data.shuffle=False
+)
+
+REWARD=(
+    reward.custom_reward_function.path=verl/utils/reward_score/openr1mm.py
+    reward.custom_reward_function.name=compute_score
+)
+
+MODEL=(
+    actor_rollout_ref.model.path=${MODEL_PATH}
+    actor_rollout_ref.model.use_remove_padding=False
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+)
+
+ACTOR=(
+    actor_rollout_ref.actor.optim.lr=5e-7
+    actor_rollout_ref.actor.ppo_mini_batch_size=16
+    actor_rollout_ref.actor.use_dynamic_bsz=False
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1
+    actor_rollout_ref.actor.use_kl_loss=True
+    actor_rollout_ref.actor.entropy_coeff=0
+    actor_rollout_ref.actor.kl_loss_coef=0.02
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl
+    actor_rollout_ref.actor.use_torch_compile=False
+    actor_rollout_ref.actor.strategy=fsdp2
+    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size}
+    actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True
+    actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True
+    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True
+    actor_rollout_ref.actor.fsdp_config.offload_policy=True
+    actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE}
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+)
+
+REF=(
+    actor_rollout_ref.ref.strategy=fsdp2
+    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=False
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1
+    actor_rollout_ref.ref.fsdp_config.param_offload=False
+    actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True
+    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True
+    actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE}
+    actor_rollout_ref.ref.use_torch_compile=False
+    actor_rollout_ref.ref.fsdp_config.offload_policy=False
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.name=${INFER_BACKEND}
+    actor_rollout_ref.rollout.prompt_length=2048
+    actor_rollout_ref.rollout.response_length=1536
+    actor_rollout_ref.rollout.ignore_eos=False
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP}
+    actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL}
+    actor_rollout_ref.rollout.n=5
+    actor_rollout_ref.rollout.enable_chunked_prefill=False
+    actor_rollout_ref.rollout.max_num_batched_tokens=4096
+    actor_rollout_ref.rollout.free_cache_engine=True
+    actor_rollout_ref.rollout.enforce_eager=False
+    actor_rollout_ref.rollout.enable_prefix_caching=False
+    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=6144
+)
+
+TRAINER=(
+    trainer.critic_warmup=0
+    trainer.logger=['console','wandb']
+    trainer.project_name="${PROJECT_NAME}"
+    trainer.experiment_name="${EXPERIMENT_NAME}"
+    trainer.n_gpus_per_node=${n_devices_per_node}
+    trainer.nnodes=${NNODES}
+    trainer.balance_batch=False
+    trainer.default_local_dir="${CKPTS_DIR}"
+    trainer.val_before_train=False
+    trainer.save_freq=5
+    trainer.test_freq=5
+    trainer.total_epochs=15
+)
+
+########################### launch ###########################
+python3 -m verl.trainer.main_ppo \
+    "${DATA[@]}" \
+    "${MODEL[@]}" \
+    "${REWARD[@]}" \
+    "${ACTOR[@]}" \
+    "${REF[@]}" \
+    "${ROLLOUT[@]}" \
+    "${TRAINER[@]}" \
+    "$@" 2>&1 | tee "${LOG_DIR}/qwen3_5-2b-openr1mm-${start_time}.log"
diff --git a/examples/grpo_trainer/run_qwen3_5_2b_video_fsdp.sh b/examples/grpo_trainer/run_qwen3_5_2b_video_fsdp.sh
new file mode 100644
index 00000000000..5e067ea7c3c
--- /dev/null
+++ b/examples/grpo_trainer/run_qwen3_5_2b_video_fsdp.sh
@@ -0,0 +1,163 @@
+#!/usr/bin/env bash
+# GRPO | Qwen3.5-2B | FSDP training | TinyLLaVA-Video-R1 video QA dataset
+# dependency: GPU vllm==0.18.0, transformers@<cc7ab9be>
+# dependency: NPU vllm==0.18.0, vllm-ascend@<54879467>, transformers@<cc7ab9be>
+
+set -xeuo pipefail
+
+########################### user-adjustable ###########################
+# DEVICE is auto-detected by probing torch_npu; override only for special cases.
+DEVICE=${DEVICE:-$(python3 -c 'import torch_npu' 2>/dev/null && echo npu || echo gpu)}
+INFER_BACKEND=${INFER_BACKEND:-vllm}
+PROJECT_NAME=${PROJECT_NAME:-GRPO-video}
+EXPERIMENT_NAME=${EXPERIMENT_NAME:-Qwen3.5-2B-GRPO-video}
+NDEVICES_PER_NODE=${NDEVICES_PER_NODE:-}
+NNODES=${NNODES:-1}
+
+GEN_TP=${GEN_TP:-4}
+SP_SIZE=${SP_SIZE:-1}
+FSDP_SIZE=${FSDP_SIZE:-}
+ROLLOUT_GPU_MEM_UTIL=${ROLLOUT_GPU_MEM_UTIL:-0.1}
+
+RAY_DATA_HOME=${RAY_DATA_HOME:-/data/verl}
+export RAY_TMPDIR=${RAY_TMPDIR:-/data/tmp/ray}
+MODEL_PATH=${MODEL_PATH:-"${HOME}/verl/models/Qwen3.5-2B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${PROJECT_NAME}/${EXPERIMENT_NAME}"}
+LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/logs/${PROJECT_NAME}/${EXPERIMENT_NAME}"}
+TRAIN_FILE=${TRAIN_FILE:-"${HOME}/verl/data/tinyllava-video-r1/train.parquet"}
+TEST_FILE=${TEST_FILE:-"${HOME}/verl/data/tinyllava-video-r1/test.parquet"}
+WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+########################### end user-adjustable ###########################
+
+########################### derived defaults ###########################
+n_devices_per_node=${NDEVICES_PER_NODE:-8}
+fsdp_size=${FSDP_SIZE:-8}
+
+case "${DEVICE}" in
+    gpu)
+        ;;
+    npu)
+        export HCCL_CONNECT_TIMEOUT=1500
+        export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+        export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+        export HCCL_EXEC_TIMEOUT=3600
+        n_devices_per_node=8
+        fsdp_size=8
+        ;;
+    *)
+        echo "Unsupported DEVICE=${DEVICE}. Expected 'gpu' or 'npu'." >&2
+        exit 1
+        ;;
+esac
+
+start_time=$(date +%Y%m%d)_$(date +%H%M%S)
+mkdir -p "${LOG_DIR}" "${RAY_TMPDIR}"
+
+########################### parameter arrays ###########################
+
+DATA=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+    data.train_files="${TRAIN_FILE}"
+    data.val_files="${TEST_FILE}"
+    data.train_batch_size=128
+    data.max_prompt_length=2048
+    data.max_response_length=1024
+    data.filter_overlong_prompts=True
+    data.filter_overlong_prompts_workers=8
+    data.truncation='error'
+    data.image_patch_size=16
+    data.image_key=images
+    data.shuffle=False
+)
+
+REWARD=(
+    reward.custom_reward_function.path=verl/utils/reward_score/tinyllava_video_r1.py
+    reward.custom_reward_function.name=compute_score
+)
+
+MODEL=(
+    actor_rollout_ref.model.path=${MODEL_PATH}
+    actor_rollout_ref.model.use_remove_padding=False
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+)
+
+ACTOR=(
+    actor_rollout_ref.actor.optim.lr=1e-6
+    actor_rollout_ref.actor.ppo_mini_batch_size=16
+    actor_rollout_ref.actor.use_dynamic_bsz=False
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1
+    actor_rollout_ref.actor.use_kl_loss=True
+    actor_rollout_ref.actor.entropy_coeff=0.001
+    actor_rollout_ref.actor.kl_loss_coef=0.01
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl
+    actor_rollout_ref.actor.use_torch_compile=False
+    actor_rollout_ref.actor.strategy=fsdp2
+    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size}
+    actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True
+    actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True
+    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True
+    actor_rollout_ref.actor.fsdp_config.offload_policy=True
+    actor_rollout_ref.actor.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE}
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+)
+
+REF=(
+    actor_rollout_ref.ref.strategy=fsdp2
+    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=False
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1
+    actor_rollout_ref.ref.fsdp_config.param_offload=False
+    actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True
+    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True
+    actor_rollout_ref.ref.fsdp_config.ulysses_sequence_parallel_size=${SP_SIZE}
+    actor_rollout_ref.ref.use_torch_compile=False
+    actor_rollout_ref.ref.fsdp_config.offload_policy=False
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.name=${INFER_BACKEND}
+    actor_rollout_ref.rollout.prompt_length=2048
+    actor_rollout_ref.rollout.response_length=1024
+    actor_rollout_ref.rollout.ignore_eos=False
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP}
+    actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEM_UTIL}
+    actor_rollout_ref.rollout.n=5
+    actor_rollout_ref.rollout.enable_chunked_prefill=True
+    actor_rollout_ref.rollout.max_num_batched_tokens=4096
+    actor_rollout_ref.rollout.free_cache_engine=True
+    actor_rollout_ref.rollout.enforce_eager=False
+    actor_rollout_ref.rollout.enable_prefix_caching=False
+    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=6144
+    actor_rollout_ref.rollout.layered_summon=True
+)
+
+TRAINER=(
+    trainer.critic_warmup=0
+    trainer.logger=['console','wandb']
+    trainer.project_name="${PROJECT_NAME}"
+    trainer.experiment_name="${EXPERIMENT_NAME}"
+    trainer.n_gpus_per_node=${n_devices_per_node}
+    trainer.nnodes=${NNODES}
+    trainer.balance_batch=False
+    trainer.default_local_dir="${CKPTS_DIR}"
+    trainer.val_before_train=False
+    trainer.save_freq=5
+    trainer.test_freq=5
+    trainer.total_epochs=15
+)
+
+########################### launch ###########################
+python3 -m verl.trainer.main_ppo \
+    "${DATA[@]}" \
+    "${MODEL[@]}" \
+    "${REWARD[@]}" \
+    "${ACTOR[@]}" \
+    "${REF[@]}" \
+    "${ROLLOUT[@]}" \
+    "${TRAINER[@]}" \
+    "$@" 2>&1 | tee "${LOG_DIR}/qwen3_5-2b-video-${start_time}.log"
+
diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py
index c65b99e44d6..bc04c0e6464 100644
--- a/verl/workers/engine_workers.py
+++ b/verl/workers/engine_workers.py
@@ -278,6 +278,7 @@ def train_mini_batch(self, data: TensorDict) -> TensorDict:
             total_num_iterations = data.shape[0] // mini_batch_size_per_gpu * epochs
 
             for batch_idx, mini_batch_td in enumerate(dataloader):
+                maybe_fix_3d_position_ids(mini_batch_td)
                 # add global token num
                 if "input_ids" in mini_batch_td:
                     global_token_num = mini_batch_td["input_ids"].offsets().diff().tolist()  # (total_nnz,)