[RL] reinforce learning benchmark framework (#10619)

* 增加rl目录下的init.py文件

* 增加copy rightb并格式化

* add reinforcement learning benchmark framework

* update copyright

* update

* add reinforce learning benchmark framework

* add elapsed time metrics in api serve framework

* fix some bug

* add reinforce learning framework with scripts file

* remove vllm quant type

* reformat rl benchmark code and scripts

* add readme file

* add fastdeploy engine support for api serve

llm/benchmark/rl/README.md

@@ -0,0 +1,188 @@
+# Large Language Model Throughput Testing Framework
+
+This framework is designed to test the **throughput performance** of Large Language Models (LLMs) across different deployment methods: **offline inference** (using PaddlePaddle and PyTorch) and **online inference** (via API). It specifically focuses on evaluating the model's ability to handle **batch queries**, measuring throughput in tokens per second under various configurations and batch sizes.
+
+## Features
+
+  * **Diverse Deployment Method Support:** Tests LLMs deployed via online API, and offline inference with PaddlePaddle and PyTorch.
+  * **Batch Query Throughput Calculation:** Accurately measures throughput (tokens/s) for concurrent queries, providing insights into the model's performance under load.
+  * **Detailed Time Logging:** Records the total time for each batch processing operation.
+
+## How It Works
+
+This framework operates by sending batched requests to specified endpoints (API or local inference scripts) and collecting performance data on how the model generates responses for multiple queries simultaneously.
+
+1.  **Input Queries:** You provide a set of questions as test input, typically from a `.parquet` or text file.
+2.  **Batch Processing:** The framework groups these questions into batches of a specified `rollout_input_batch_size`.
+3.  **Generate Responses:** For each query within a batch, the framework requests the model to generate `rollout_n` responses.
+4.  **Time Measurement:** The total time from sending a batch of questions to receiving all corresponding responses is recorded.
+5.  **Token Statistics:** The total number of tokens generated across all responses in a batch is summed up.
+6.  **Throughput Calculation:** Throughput (tokens/s) is calculated by dividing the total tokens generated by the total time taken for the batch to complete.
+
+## Usage
+
+This section details how to run throughput tests for each deployment method using the provided shell scripts.
+
+### Data Preparation
+
+To run the tests, you'll first need to download and extract the `rl_data.tar.gz` archive, which contains the GSM8K dataset in a suitable format for testing.
+
+```
+cd llm/benchmark/rl
+wget https://paddle-qa.bj.bcebos.com/paddlenlp/rl_data.tar.gz     
+tar -zxvf rl_data.tar.gz 
+```
+
+Extract the contents of the archive. This will create a data folder containing the GSM8K dataset.
+
+### Online API Inference
+
+This script tests the throughput of a remote LLM API.
+
+**Configuration (`api_serve.sh`):**
+
+```bash
+output_dir="api_serve_results"
+
+python api_serve.py \
+    --openai_urls "your_url1" "your_url2"\
+    --api_keys "key1" "key2" \
+    --model "Qwen2.5-7B-Instruct-1M" \
+    --tokenizer "Qwen/Qwen2.5-7B-Instruct-1M" \
+    --input_file ./data/gsm8k/instruct/train.parquet \
+    --output_dir ${output_dir} \
+    --use_fastdeploy true \
+    --rollout_input_batch_size 8 \
+    --rollout_n 8 \
+    --top_p 1.0 \
+    --temperature 0.7 \
+    --max_dec_len 8192 \
+    --limit_rows 512
+```
+
+  * **`--openai_urls`**: URLs of the API endpoints to test.
+  * **`--api_keys`**: API keys for authentication (if required).
+  * **`--model`**: Name of the model being tested.
+  * **`--tokenizer`**: Path or name of the tokenizer.
+  * **`--input_file`**: Path to the input dataset file.
+  * **`--output_dir`**: Directory to save output results.
+  * **`--use_fastdeploy`**: Use FastDeploy if true, otherwise use vLLM (default: true).
+  * **`--rollout_input_batch_size`**: The batch size for API requests.
+  * **`--rollout_n`**: Number of responses to generate for each input query.
+  * **`--max_dec_len`**: Maximum decoding length for responses.
+  * **`--limit_rows`**: Limit the number of input rows processed.
+
+**Run command:**
+
+```bash
+bash scripts/api_serve.sh
+```
+
+### Offline PaddlePaddle Inference
+
+This script tests the throughput of an LLM using offline PaddlePaddle inference, potentially with distributed processing.
+
+**Configuration (`paddle_infer.sh`):**
+
+```bash
+unset PADDLE_TRAINERS_NUM
+unset PADDLE_ELASTIC_JOB_ID
+unset PADDLE_TRAINER_ENDPOINTS
+unset DISTRIBUTED_TRAINER_ENDPOINTS
+unset FLAGS_START_PORT
+unset PADDLE_ELASTIC_TIMEOUT
+
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP":$PYTHONPATH
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP/llm":$PYTHONPATH
+
+export FLAGS_set_to_1d=False
+export NVIDIA_TF32_OVERRIDE=0
+export FLAGS_dataloader_use_file_descriptor=False
+export HF_DATASETS_DOWNLOAD_TIMEOUT=1
+export FLAGS_gemm_use_half_precision_compute_type=False
+export FLAGS_force_cublaslt_no_reduced_precision_reduction=True
+
+export FLAGS_custom_allreduce=0
+export FLAGS_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+output_dir="pdpd_bf16_offline"
+
+python -u -m paddle.distributed.launch --log_dir ${output_dir}/logs --gpus ${CUDA_VISIBLE_DEVICES} paddle_infer.py \
+  --actor_model_name_or_path your_model_name \
+  --max_src_len 2048 \
+  --min_dec_len 32 \
+  --max_dec_len 30720 \
+  --top_p 1.0 \
+  --temperature 1.0 \
+  --rollout_input_batch_size 4 \
+  --rollout_n 8 \
+  --rollout_max_num_seqs 24 \
+  --rollout_quant_type "" \
+  --tensor_parallel_degree 4 \
+  --limit_rows 640 \
+  --input_file file.parquet \
+  --output_dir ${output_dir} > ./paddleinfer.log 2>&1
+```
+
+  * **`CUDA_VISIBLE_DEVICES`**: Specifies the GPUs to be used.
+  * **`paddle.distributed.launch`**: Initiates a distributed PaddlePaddle training/inference job.
+  * **`--actor_model_name_or_path`**: Path to the pre-trained model.
+  * **`--max_src_len`**: Maximum source sequence length.
+  * **`--rollout_input_batch_size`**: The batch size for inference.
+  * **`--rollout_n`**: Number of responses to generate for each input query.
+  * **`--tensor_parallel_degree`**: Degree of tensor parallelism for distributed inference.
+  * **`--input_file`**: Path to the input dataset file.
+  * **`--output_dir`**: Directory to save output results and logs.
+
+### Offline PyTorch Inference
+
+This script tests the throughput of an LLM using offline PyTorch inference.
+
+**Configuration (`torch_infer.sh`):**
+
+```bash
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+
+output_dir="vllm_bf16_offline_flashattn"
+
+python torch_infer.py \
+    --actor_model_name_or_path Qwen/Qwen2.5-7B-Instruct-1M \
+    --max_src_len 2048 \
+    --min_dec_len 32 \
+    --max_dec_len 30720 \
+    --top_p 1.0 \
+    --temperature 1.0 \
+    --rollout_input_batch_size 4 \
+    --rollout_n 8 \
+    --tensor_parallel_degree 4 \
+    --limit_rows 640 \
+    --input_file ./data/gsm8k/instruct/train.parquet \
+    --output_dir ${output_dir} \
+    --gpu_memory_utilization 0.8 > ./torchinferflashattn.log 2>&1
+```
+  * **`--gpu_memory_utilization`**: Fraction of GPU memory to be reserved for the model.
+
+-----
+
+## Output Results
+
+The `output_dir` contains the following files:  
+
+**1. Statistics Files**  
+• `dispersed_stats.csv`  
+
+  Per-batch request length and throughput statistics. Fields:  
+  `batch_index, rollout_lengths, min_length, max_length, avg_length, completion_time, throughput_tokens_per_sec`  
+
+• `global_stats.csv`  
+
+  Aggregated global metrics. Fields:  
+  `batch_index, min_response_tokens, max_response_tokens, avg_response_tokens, total_response_tokens, completion_time, throughput_tokens_per_sec`  
+
+**2. Detailed Records**  
+• `rollout_details.jsonl`  
+
+  Raw per-request outputs (JSON Lines format), including input/output text.
+

llm/benchmark/rl/api_serve.py

@@ -0,0 +1,391 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import asyncio
+import csv
+import json
+import math
+import time
+from dataclasses import dataclass, field
+from itertools import cycle
+from pathlib import Path
+from typing import List, Tuple
+
+import pandas as pd
+from openai import AsyncOpenAI
+from tqdm import tqdm
+from utils import RangeSet
+
+from paddlenlp.transformers import AutoTokenizer
+from paddlenlp.utils.log import logger
+
+
+@dataclass
+class RequestPayload:
+    prompt: str = ""
+    num_responses: int = 8
+    idx: int = 0
+
+
+@dataclass
+class ResponsePayload:
+    idx: int = 0
+    question: str = ""
+    question_token_length: int = 0
+    responses: List[str] = field(default_factory=list)
+    elapsed_times: List[float] = field(default_factory=list)
+    token_lengths: List[int] = field(default_factory=list)
+    total_length: int = 0
+
+
+class ApiTask:
+    def __init__(self, args):
+        self.args = args
+        self.model = args.model
+        self.tokenzizer = AutoTokenizer.from_pretrained(self.args.tokenizer, use_fast=True)
+        self.clients = cycle(
+            AsyncOpenAI(base_url=url, api_key=api) for url, api in zip(args.openai_urls, args.api_keys)
+        )
+        self._max_concurrency = args.max_concurrency
+        self.semaphore = asyncio.Semaphore(args.max_concurrency)
+        self.output_dir = Path(self.args.output_dir)
+
+        self.global_stats_path = self.output_dir / "global_stats.csv"
+        self.dispersed_stats_path = self.output_dir / "dispersed_stats.csv"
+        self.rollout_details_path = self.output_dir / "rollout_details.jsonl"
+        self.status_file_path = self.output_dir / "status.txt"
+
+        self._load_status()
+
+    def tokenize(self, response: ResponsePayload) -> ResponsePayload:
+        question = response.question
+        responses = response.responses
+        response.question_token_length = len(self.tokenzizer(question).input_ids)
+
+        for i, resp in enumerate(responses):
+            tokens = self.tokenzizer(resp).input_ids
+            length = len(tokens)
+            response.token_lengths.append(length)
+            response.total_length += length
+
+        return response
+
+    def get_active_tasks_count(self) -> int:
+        return self._max_concurrency - self.semaphore._value
+
+    def get_client(self) -> AsyncOpenAI:
+        # Returns an AsyncOpenAI client instance
+        return next(self.clients)
+
+    def _save_status(self, batch_index):
+        """Save current processing status to file"""
+        self.processed_set.add(batch_index)
+        content = self.processed_set.to_file_format()
+        with open(self.status_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def _load_status(self):
+        """Load processing status from file"""
+        try:
+            with open(self.status_file_path, "r", encoding="utf-8") as f:
+                content = f.read().strip()
+                self.processed_set = RangeSet.from_file(content)
+                logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
+        except FileNotFoundError:
+            self.processed_set = RangeSet([])
+
+    def process_data(self, file_path: str) -> pd.DataFrame:
+        logger.info(f"Processing data from {file_path}...")
+        start_time = time.time()
+        df = pd.read_parquet(file_path)
+        if self.args.limit_rows != -1:
+            df = df.iloc[: self.args.limit_rows]
+        logger.info(f"Loaded {len(df)} samples in {time.time() - start_time:.2f}s")
+        return df
+
+    def batch_process(self, dataframe: pd.DataFrame):
+        batch_prompts = []
+        for idx, prompt in enumerate(dataframe[self.args.prompt_key]):
+            batch_prompts.append(
+                RequestPayload(prompt=prompt[0]["content"], idx=idx, num_responses=self.args.rollout_n)
+            )
+            if len(batch_prompts) == self.args.rollout_input_batch_size:
+                yield batch_prompts
+                batch_prompts = []
+
+    async def fastdeploy_call(self, request: RequestPayload) -> Tuple[str, float]:
+        client = self.get_client()
+        try:
+            async with self.semaphore:
+                start_time = time.time()
+                response = await client.chat.completions.create(
+                    model=self.model,
+                    messages=[{"role": "user", "content": request.prompt}],
+                    temperature=self.args.temperature,
+                    top_p=self.args.top_p,
+                    max_tokens=self.args.max_response_length,
+                    n=1,
+                    stream=True,
+                    timeout=60*60,
+                    metadata={
+                        "training": True,
+                        "raw_request": False,
+                    }
+                ) 
+                # Streaming text is stored in a list of chunks
+                chunks = []
+                # Streaming responses
+                async for chunk in response:
+                    delta = chunk.choices[0].delta
+                    if delta and delta.content:
+                        chunks.append(delta.content)
+                text = "".join(chunks)
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                logger.debug("Streaming response took %.2f seconds", elapsed_time)
+                return text, round(elapsed_time, 2)
+
+        except Exception as e:
+            logger.error("Error while streaming: %s", e)
+            raise ValueError(e)
+
+    async def vllm_call(self, request: RequestPayload) -> Tuple[str, float]:
+        client = self.get_client()
+        try:
+            async with self.semaphore:
+                start_time = time.time()
+                response = await client.completions.create(
+                    model=self.model,
+                    prompt=request.prompt,
+                    temperature=self.args.temperature,
+                    top_p=self.args.top_p,
+                    max_tokens=self.args.max_dec_len,
+                    n=1,
+                    stream=True,
+                )
+                # Streaming text is stored in a list of chunks
+                chunks = []
+                # Streaming responses
+                async for chunk in response:
+                    if chunk.choices and chunk.choices[0].text:
+                        chunks.append(chunk.choices[0].text)
+                text = "".join(chunks)
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                logger.debug("Streaming response took %.2f seconds", elapsed_time)
+                return text, round(elapsed_time, 2)
+
+        except Exception as e:
+            logger.error("Error while streaming: %s", e)
+            raise ValueError(e)
+
+    async def group_call(self, request: RequestPayload) -> ResponsePayload:
+        """Performs n complete token generation rollouts for the given query."""
+        if self.args.use_fastdeploy == "true":
+            call = self.fastdeploy_call
+        else:
+            call = self.vllm_call
+
+        tasks = [call(request) for _ in range(request.num_responses)]
+
+        result = ResponsePayload()
+        result.idx = request.idx
+        result.question = request.prompt
+        for task, elapsed_time in await asyncio.gather(*tasks):
+            result.responses.append(task)
+            result.elapsed_times.append(elapsed_time)
+        return result
+
+    async def batch_call(self, requests: List[RequestPayload]) -> Tuple[List[ResponsePayload], int]:
+        """Batch execution requests"""
+        start_time = time.time()
+        batch_results = await asyncio.gather(*[self.group_call(request) for request in requests])
+        end_time = time.time()
+        batch_elapsed_time = end_time - start_time
+        logger.debug("total batch took %.2f seconds", batch_elapsed_time)
+        return batch_results, batch_elapsed_time
+
+    def dispersed_stats(self, responses: List[ResponsePayload], batch_elapsed_time: float, batch_index):
+        batch_group_pd = pd.DataFrame(responses)
+
+        dispersed_stats_dict = {
+            "batch_index": batch_index,
+            "rollout_lengths": batch_group_pd["token_lengths"].to_list(),
+            "min_length": batch_group_pd["token_lengths"].apply(lambda x: min(x)).tolist(),
+            "max_length": batch_group_pd["token_lengths"].apply(lambda x: max(x)).tolist(),
+            "avg_length": batch_group_pd["token_lengths"].apply(lambda x: sum(x) / len(x)).tolist(),
+            "completion_time": batch_elapsed_time,
+            "throughput_tokens_per_sec": batch_group_pd["token_lengths"].apply((lambda x: sum(x))).sum()
+            / batch_elapsed_time,
+        }
+
+        return dispersed_stats_dict
+
+    def global_stats(self, responses: List[ResponsePayload], batch_elapsed_time: float, batch_index):
+        dispersed_stats_dict = self.dispersed_stats(responses, batch_elapsed_time, batch_index)
+
+        total_response_tokens = 0
+        for lengths in dispersed_stats_dict["rollout_lengths"]:
+            total_response_tokens += sum(lengths)
+
+        global_stats_dict = {}
+        global_stats_dict["batch_index"] = dispersed_stats_dict["batch_index"]
+        global_stats_dict["min_response_tokens"] = min(dispersed_stats_dict["min_length"])
+        global_stats_dict["max_response_tokens"] = max(dispersed_stats_dict["max_length"])
+        global_stats_dict["avg_response_tokens"] = total_response_tokens / (
+            self.args.rollout_n * self.args.rollout_input_batch_size
+        )
+        global_stats_dict["total_response_tokens"] = total_response_tokens
+        global_stats_dict["group_max_response_tokens"] = dispersed_stats_dict["max_length"]
+        global_stats_dict["completion_time"] = dispersed_stats_dict["completion_time"]
+        global_stats_dict["throughput_tokens_per_sec"] = dispersed_stats_dict["throughput_tokens_per_sec"]
+
+        return global_stats_dict, dispersed_stats_dict
+
+    def execute(self):
+        dataframe = self.process_data(self.args.input_file)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(self.global_stats_path, "a", newline="") as global_f, open(
+            self.dispersed_stats_path, "a", newline=""
+        ) as dispersed_f, open(self.rollout_details_path, "a", encoding="utf-8") as jsonl_f:
+            global_writer = csv.writer(global_f)
+            dispersed_writer = csv.writer(dispersed_f)
+
+            if self.processed_set.processed_count <= 0:
+                global_writer.writerow(
+                    [
+                        "batch_index",
+                        "min_response_tokens",
+                        "max_response_tokens",
+                        "avg_response_tokens",
+                        "total_response_tokens",
+                        "group_max_response_tokens",
+                        "completion_time",
+                        "throughput_tokens_per_sec",
+                    ]
+                )
+                dispersed_writer.writerow(
+                    [
+                        "batch_index",
+                        "rollout_lengths",
+                        "min_length",
+                        "max_length",
+                        "avg_length",
+                        "completion_time",
+                        "throughput_tokens_per_sec",
+                    ]
+                )
+
+            for batch_index, input_ids in tqdm(
+                enumerate(self.batch_process(dataframe)),
+                total=math.ceil(len(dataframe) / self.args.rollout_input_batch_size),
+                mininterval=0.1,
+            ):
+                if self.processed_set.contains(batch_index):
+                    continue
+
+                batch_results, batch_elapsed_time = asyncio.run(self.batch_call(input_ids))
+
+                for i in range(len(batch_results)):
+                    batch_results[i] = self.tokenize(batch_results[i])
+
+                global_stats_dict, dispersed_stats_dict = self.global_stats(
+                    batch_results, batch_elapsed_time, batch_index
+                )
+
+                global_writer.writerow(
+                    [
+                        batch_index,
+                        global_stats_dict["min_response_tokens"],
+                        global_stats_dict["max_response_tokens"],
+                        round(global_stats_dict["avg_response_tokens"], 2),
+                        global_stats_dict["total_response_tokens"],
+                        global_stats_dict["group_max_response_tokens"],
+                        round(global_stats_dict["completion_time"], 2),
+                        round(global_stats_dict["throughput_tokens_per_sec"], 2),
+                    ]
+                )
+
+                dispersed_writer.writerow(
+                    [
+                        batch_index,
+                        dispersed_stats_dict["rollout_lengths"],
+                        dispersed_stats_dict["min_length"],
+                        dispersed_stats_dict["max_length"],
+                        dispersed_stats_dict["avg_length"],
+                        round(dispersed_stats_dict["completion_time"], 2),
+                        round(dispersed_stats_dict["throughput_tokens_per_sec"], 2),
+                    ]
+                )
+
+                record = [
+                    {
+                        "batch_index": batch_index,
+                        "prompt_text": result.question,
+                        "rollouts": [
+                            {"response": res, "token_length": token_length}
+                            for res, token_length in zip(result.responses, result.token_lengths)
+                        ],
+                        "total_time": dispersed_stats_dict["completion_time"],
+                        "throughput_tokens_per_sec": dispersed_stats_dict["throughput_tokens_per_sec"],
+                    }
+                    for result in batch_results
+                ]
+
+                jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+                global_f.flush()
+                dispersed_f.flush()
+                jsonl_f.flush()
+                self._save_status(batch_index)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process prompts with OpenAI clients.")
+    parser.add_argument("--openai_urls", type=str, nargs="+", required=True, help="List of OpenAI service URLs")
+    parser.add_argument(
+        "--api_keys", type=str, nargs="+", default=None, help="List of API keys (default: 'NONE' for each service)"
+    )
+    parser.add_argument("--model", type=str, required=True, help="Model name (e.g., Qwen2.5-7B-Instruct-1M)")
+    parser.add_argument(
+        "--tokenizer", type=str, required=True, help="Tokenizer name (e.g., Qwen/Qwen2.5-7B-Instruct-1M)"
+    )
+    parser.add_argument("--rollout_input_batch_size", type=int, default=4, help="Batch size for requests")
+    parser.add_argument("--rollout_n", type=int, default=8, help="Number of responses per request")
+    parser.add_argument(
+        "--prompt_key", type=str, default="prompt", help="Key in the DataFrame for prompts (default: 'prompt')"
+    )
+    parser.add_argument("--input_file", type=str, required=True, help="Path to the input Parquet file")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./api_infer_results",
+        help="Directory for output CSV files (default: './api_infer_results')",
+    )
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p sampling parameter for text generation")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature parameter for text generation")
+    parser.add_argument("--max_dec_len", type=int, default=1024 * 2, help="Maximum response length (in tokens)")
+    parser.add_argument("--max_concurrency", type=int, default=1000, help="Maximum concurrent connections")
+    parser.add_argument(
+        "--limit_rows", type=int, default=-1, help="Maximum number of rows to read from the dataset (-1 means all)"
+    )
+    parser.add_argument("--use_fastdeploy", type=str.lower, choices=["true", "false"], default="true", help="Engine selection (true=FastDeploy, false=vLLM, default: true)")
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    task = ApiTask(args)
+    task.execute()

llm/benchmark/rl/paddle_infer.py

@@ -0,0 +1,351 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import csv
+import json
+import math
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import paddle
+import pandas as pd
+from tqdm import tqdm
+from utils import RangeSet
+
+from paddlenlp.rl.trainer import process_row
+from paddlenlp.rl.utils import TrainingArguments
+from paddlenlp.rl.utils.infer_utils import get_policy_predictor, infer_guard
+from paddlenlp.trainer import PdArgumentParser, set_seed
+from paddlenlp.transformers import AutoModelForCausalLM, AutoTokenizer
+from paddlenlp.trl.llm_utils import init_dist_env
+from paddlenlp.utils.log import logger
+
+
+@contextmanager
+def switch_level_context(level="ERROR"):
+    original_level = logger.logLevel
+    logger.set_level(level)
+
+    try:
+        yield
+    finally:
+        logger.set_level(original_level)
+
+
+def chunk(all_input_ids, size):
+    if size <= 0:
+        raise ValueError("Size must be greater than 0")
+    return [all_input_ids[i : i + size] for i in range(0, len(all_input_ids), size)]
+
+
+@dataclass
+class DumpyTrainingArguments(TrainingArguments):
+    actor_model_name_or_path: str = field(
+        default="Qwen/Qwen2.5-7B-Instruct-1M", metadata={"help": "pretrained model name or path"}
+    )
+    input_file: str = field(
+        default="kk/instruct/3-7ppl/combined.parquet", metadata={"help": "input the Parquet file path"}
+    )
+    limit_rows: int = field(
+        default=-1, metadata={"help": "Maximum number of rows to read from the dataset (-1 means all)"}
+    )
+    output_dir: str = field(default="./pt_infer_results", metadata={"help": "output directory path"})
+    rollout_input_batch_size: int = field(default=32, metadata={"help": "batch size for inference engine inputs"})
+    rollout_n: int = field(default=8, metadata={"help": "number of rollouts (for performance testing)"})
+    log_interval: int = field(default=1, metadata={"help": "logging interval (in batches)"})
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.dtype = "bfloat16"
+        if self.rollout_quant_type.lower() in ["bfloat16", "bf16"]:
+            self.rollout_quant_type = ""
+        elif self.rollout_quant_type.lower() in ["wint8", "weight_only_int8"]:
+            self.rollout_quant_type = "weight_only_int8"
+
+
+class DumpyInferenceTask:
+    def __init__(self, args):
+        self.args = args
+        self.output_dir = Path(args.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        # ​​Initialize output file path​
+        self.global_stats_path = self.output_dir / "global_stats.csv"
+        self.dispersed_stats_path = self.output_dir / "dispersed_stats.csv"
+        self.rollout_details_path = self.output_dir / "rollout_details.jsonl"
+        self.status_file_path = self.output_dir / "status.txt"
+
+        self._init_environment()
+        self._load_status()
+        self._load_model()
+        self._prepare_tokenizer()
+
+    def _init_environment(self):
+        self.tensor_parallel_rank, self.tensor_parallel_degree = init_dist_env()
+        self.use_fusemt = True
+        self.amp_dtype = self.args.dtype
+
+    def _load_status(self):
+        """Load processing status from file"""
+        try:
+            with open(self.status_file_path, "r", encoding="utf-8") as f:
+                content = f.read().strip()
+                self.processed_set = RangeSet.from_file(content)
+                logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
+        except FileNotFoundError:
+            self.processed_set = RangeSet([])
+
+    def _save_status(self, batch_index):
+        """Save current processing status to file"""
+        self.processed_set.add(batch_index)
+        content = self.processed_set.to_file_format()
+        with open(self.status_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def _load_model(self):
+        logger.info(f"Loading model from {self.args.actor_model_name_or_path}...")
+        start_time = time.time()
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.args.actor_model_name_or_path,
+            dtype=self.args.dtype,
+            low_cpu_mem_usage=True,
+            tensor_parallel_degree=self.tensor_parallel_degree,
+            tensor_parallel_rank=self.tensor_parallel_rank,
+        )
+        self.model.eval()
+        self.model.to(device=paddle.CUDAPinnedPlace())
+        logger.info(f"Model loaded in {time.time() - start_time:.2f}s")
+
+    def _prepare_tokenizer(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.actor_model_name_or_path, use_fast=True)
+        self.tokenizer.padding_side = "left"
+
+    def process_data(self, file_path: str) -> pd.DataFrame:
+        logger.info(f"Processing data from {file_path}...")
+        start_time = time.time()
+        df = pd.read_parquet(file_path)
+        if self.args.limit_rows != -1:
+            df = df.iloc[: self.args.limit_rows]
+        logger.info(f"Loaded {len(df)} samples in {time.time() - start_time:.2f}s")
+        return df
+
+    @switch_level_context(level="ERROR")
+    def run_inference(self, input_ids, batch_index=0):
+        set_seed(42)
+        start_time = time.time()
+        output_ids = get_policy_predictor().predict(
+            input_ids=input_ids,
+            repeat_num=self.args.rollout_n,
+            temperature=self.args.temperature,
+            top_p=self.args.top_p,
+        )
+        if self.args.world_size > 1:
+            paddle.distributed.barrier()
+        end_time = time.time()
+        if self.args.should_log:
+            statistics = self.postprocess_data(input_ids, output_ids, batch_index=batch_index)
+            statistics["total_time"] = end_time - start_time
+            return statistics
+        return None
+
+    def postprocess_data(self, input_ids, output_ids, batch_index=0):
+        # Process prompts
+        global_prompt_tokens = []
+        global_prompt_tokens_len = []
+        group_prompt_texts = []
+        for row in input_ids:
+            row_ids = process_row(row, remove_value=self.tokenizer.pad_token_id, remove_side="both").tolist()
+            global_prompt_tokens.append(row_ids)
+            global_prompt_tokens_len.append(len(row_ids))
+            group_prompt_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Process responses
+        response_tokens = []
+        response_tokens_len = []
+        response_texts = []
+        for row in output_ids:
+            row_ids = process_row(row, remove_value=self.tokenizer.pad_token_id, remove_side="both").tolist()
+            response_tokens.append(row_ids)
+            response_tokens_len.append(len(row_ids))
+            response_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Group processing
+        group_response_tokens = chunk(response_tokens, self.args.rollout_n)
+        group_response_tokens_len = chunk(response_tokens_len, self.args.rollout_n)
+        group_response_texts = chunk(response_texts, self.args.rollout_n)
+
+        # Calculate statistics
+        global_response_tokens_total = sum(response_tokens_len)
+        global_response_tokens_len_min = min(response_tokens_len)
+        global_response_tokens_len_max = max(response_tokens_len)
+        global_response_tokens_len_mean = round(global_response_tokens_total / len(response_tokens_len), 2)
+
+        group_response_tokens_len_mean = [round(sum(x) / len(x), 2) for x in group_response_tokens_len]
+        group_response_tokens_len_max = [max(x) for x in group_response_tokens_len]
+        group_response_tokens_len_min = [min(x) for x in group_response_tokens_len]
+
+        return {
+            # Global prompt stats
+            "batch_index": batch_index,
+            "global_prompt_tokens_len": global_prompt_tokens_len,
+            "group_prompt_texts": group_prompt_texts,
+            # Group response stats
+            "group_response_tokens": group_response_tokens,
+            "group_response_tokens_len": group_response_tokens_len,
+            "group_response_texts": group_response_texts,
+            "group_response_tokens_len_min": group_response_tokens_len_min,
+            "group_response_tokens_len_max": group_response_tokens_len_max,
+            "group_response_tokens_len_mean": group_response_tokens_len_mean,
+            # Global response stats
+            "global_response_tokens_total": global_response_tokens_total,
+            "global_response_tokens_len_min": global_response_tokens_len_min,
+            "global_response_tokens_len_max": global_response_tokens_len_max,
+            "global_response_tokens_len_mean": global_response_tokens_len_mean,
+        }
+
+    def batch_process(self, dataframe: pd.DataFrame):
+        all_input_ids = []
+
+        def _pad_batch(input_ids_list):
+            return self.tokenizer.pad({"input_ids": input_ids_list}, padding=True, return_tensors="pd")["input_ids"]
+
+        for idx, row in dataframe.iterrows():
+            prompt = row["prompt"].tolist()[0]["content"]
+            input_ids = self.tokenizer(prompt, add_special_tokens=False, return_attention_mask=False)["input_ids"]
+            all_input_ids.append(input_ids)
+
+            if len(all_input_ids) >= self.args.rollout_input_batch_size:
+                yield _pad_batch(all_input_ids)
+                all_input_ids = []
+
+        if all_input_ids:
+            yield _pad_batch(all_input_ids)
+
+    def execute(self):
+        """Main Execution Pipeline"""
+        # Data Preparation​
+        dataframe = self.process_data(self.args.input_file)
+
+        # Create Output Directory
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(self.global_stats_path, "a", newline="") as global_f, open(
+            self.dispersed_stats_path, "a", newline=""
+        ) as dispersed_f, open(self.rollout_details_path, "a", encoding="utf-8") as jsonl_f:
+
+            if self.args.should_log:
+                # Initialize the CSV writer
+                global_writer = csv.writer(global_f)
+                dispersed_writer = csv.writer(dispersed_f)
+                if self.processed_set.processed_count <= 0:
+                    global_writer.writerow(
+                        [
+                            "batch_index",
+                            "min_response_tokens",
+                            "max_response_tokens",
+                            "avg_response_tokens",
+                            "total_response_tokens",
+                            "group_max_response_tokens",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+                    dispersed_writer.writerow(
+                        [
+                            "batch_index",
+                            "rollout_lengths",
+                            "min_length",
+                            "max_length",
+                            "avg_length",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+
+            with infer_guard(self):
+                for batch_index, input_ids in tqdm(
+                    enumerate(self.batch_process(dataframe)),
+                    total=math.ceil(len(dataframe) / self.args.rollout_input_batch_size),
+                    disable=not self.args.should_log,
+                ):
+                    if self.processed_set.contains(batch_index):
+                        continue
+
+                    statistics = self.run_inference(input_ids, batch_index=batch_index)
+
+                    if self.args.should_log:
+                        # Write global statistics
+                        total_time = round(statistics["total_time"], 2)
+                        total_tokens = statistics["global_response_tokens_total"]
+                        throughput = round(total_tokens / total_time if total_time > 0 else 0, 2)
+
+                        global_writer.writerow(
+                            [
+                                batch_index,
+                                statistics["global_response_tokens_len_min"],
+                                statistics["global_response_tokens_len_max"],
+                                statistics["global_response_tokens_len_mean"],
+                                total_tokens,
+                                statistics["group_response_tokens_len_max"],
+                                total_time,
+                                throughput,
+                            ]
+                        )
+
+                        dispersed_writer.writerow(
+                            [
+                                batch_index,
+                                statistics["group_response_tokens_len"],
+                                statistics["group_response_tokens_len_min"],
+                                statistics["group_response_tokens_len_max"],
+                                statistics["group_response_tokens_len_mean"],
+                                total_time,
+                                throughput,
+                            ]
+                        )
+
+                        # Write detailed records (one row per query)
+                        prompt_text = statistics["group_prompt_texts"]
+                        response_tokens = statistics["group_response_tokens"]
+                        response_texts = statistics["group_response_texts"]
+
+                        record = {
+                            "batch_index": batch_index,
+                            "prompt_text": prompt_text,
+                            "rollouts": [
+                                {"token_ids": tokens, "text": text, "length": len(tokens)}
+                                for tokens, text in zip(response_tokens, response_texts)
+                            ],
+                            "total_time": total_time,
+                            "throughput_tokens_per_sec": throughput,
+                        }
+                        jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+                        global_f.flush()
+                        dispersed_f.flush()
+                        jsonl_f.flush()
+                        self._save_status(batch_index)
+
+
+def main():
+    parser = PdArgumentParser((DumpyTrainingArguments,))
+    args = parser.parse_args_into_dataclasses()[0]
+    task = DumpyInferenceTask(args)
+    task.execute()
+
+
+if __name__ == "__main__":
+    main()

llm/benchmark/rl/scripts/api_serve.sh

@@ -0,0 +1,16 @@
+output_dir="your_output_dir"
+
+python api_serve.py \
+    --openai_urls "your_url1" "your_url2"\
+    --api_keys "key1" "key2" \
+    --model "Qwen2.5-7B-Instruct-1M" \
+    --tokenizer "Qwen/Qwen2.5-7B-Instruct-1M" \
+    --use_fastdeploy true \
+    --input_file your_file \
+    --output_dir ${output_dir} \
+    --rollout_input_batch_size 8 \
+    --rollout_n 8 \
+    --top_p 1.0 \
+    --temperature 0.7 \
+    --max_dec_len 8192 \
+    --limit_rows 512

llm/benchmark/rl/scripts/paddle_infer.sh

@@ -0,0 +1,39 @@
+unset PADDLE_TRAINERS_NUM
+unset PADDLE_ELASTIC_JOB_ID
+unset PADDLE_TRAINER_ENDPOINTS
+unset DISTRIBUTED_TRAINER_ENDPOINTS
+unset FLAGS_START_PORT
+unset PADDLE_ELASTIC_TIMEOUT
+
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP":$PYTHONPATH
+export PYTHONPATH="your_paddlenlp_path/PaddleNLP/llm":$PYTHONPATH
+
+export FLAGS_set_to_1d=False
+export NVIDIA_TF32_OVERRIDE=0
+export FLAGS_dataloader_use_file_descriptor=False
+export HF_DATASETS_DOWNLOAD_TIMEOUT=1
+export FLAGS_gemm_use_half_precision_compute_type=False
+export FLAGS_force_cublaslt_no_reduced_precision_reduction=True
+
+export FLAGS_custom_allreduce=0
+export FLAGS_mla_use_tensorcore=0
+export FLAGS_cascade_attention_max_partition_size=2048
+
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+output_dir="pdpd_bf16_offline"
+
+python -u -m paddle.distributed.launch --log_dir ${output_dir}/logs --gpus ${CUDA_VISIBLE_DEVICES} paddle_infer.py \
+  --actor_model_name_or_path your_model_name \
+  --max_src_len 2048 \
+  --min_dec_len 32 \
+  --max_dec_len 30720 \
+  --top_p 1.0 \
+  --temperature 1.0 \
+  --rollout_input_batch_size 4 \
+  --rollout_n 8 \
+  --rollout_max_num_seqs 24 \
+  --rollout_quant_type "" \
+  --tensor_parallel_degree 4 \
+  --limit_rows 640 \
+  --input_file file.parquet \
+  --output_dir ${output_dir} > ./paddleinfer.log 2>&1 

llm/benchmark/rl/scripts/torch_infer.sh

@@ -0,0 +1,18 @@
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+
+output_dir="vllm_bf16_offline_flashattn"
+
+nohup python torch_infer.py \
+    --actor_model_name_or_path your_model_name \
+    --max_src_len 2048 \
+    --min_dec_len 32 \
+    --max_dec_len 30720 \
+    --top_p 1.0 \
+    --temperature 1.0 \
+    --rollout_input_batch_size 4 \
+    --rollout_n 8 \
+    --tensor_parallel_degree 4 \
+    --limit_rows 640 \
+    --input_file file.parquet \
+    --output_dir ${output_dir} \
+    --gpu_memory_utilization 0.8 > ./torchinferflashattn.log 2>&1 &

llm/benchmark/rl/torch_infer.py

@@ -0,0 +1,339 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import csv
+import json
+import math
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import pandas as pd
+import tqdm
+from transformers import AutoTokenizer, logging
+from utils import RangeSet
+from vllm import LLM, SamplingParams
+from vllm.inputs import TokensPrompt
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--actor_model_name_or_path",
+        type=str,
+        default="Qwen/Qwen2.5-7B-Instruct-1M",
+        help="Name or path of the pretrained model",
+    )
+    parser.add_argument("--input_file", type=str, default="./combined.parquet", help="Path to input parquet file")
+    parser.add_argument("--output_dir", type=str, default="./pt_infer_results", help="Path to output directory")
+    parser.add_argument(
+        "--rollout_input_batch_size", type=int, default=2, help="Batch size for each inference engine input"
+    )
+    parser.add_argument(
+        "--rollout_n", type=int, default=2, help="Number of repeated inference runs (for performance testing)"
+    )
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type for computation")
+    parser.add_argument("--tensor_parallel_degree", type=int, default=1, help="Degree of model parallelism")
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p sampling parameter for text generation")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature parameter for text generation")
+    parser.add_argument("--should_log", type=bool, default=True, help="Whether to print logging information")
+    parser.add_argument("--max_src_len", type=int, default=1024 * 2, help="Maximum prompt length (in tokens)")
+    parser.add_argument("--max_dec_len", type=int, default=1024 * 2, help="Maximum response length (in tokens)")
+    parser.add_argument("--min_dec_len", type=int, default=0, help="Minimum response length (in tokens)")
+    parser.add_argument(
+        "--limit_rows", type=int, default=-1, help="Maximum number of rows to read from the dataset (-1 means all)"
+    )
+    parser.add_argument("--gpu_memory_utilization", type=float, default=0.9, help="Percentage of GPU memory usage")
+    args = parser.parse_args()
+
+    return args
+
+
+def chunk(all_input_ids, size):
+    if size <= 0:
+        raise ValueError("Size must be greater than 0")
+    return [all_input_ids[i : i + size] for i in range(0, len(all_input_ids), size)]
+
+
+class DumpyInferenceTask:
+    def __init__(self, args):
+        self.args = args
+        self.output_dir = Path(args.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.global_stats_path = self.output_dir / "global_stats.csv"
+        self.dispersed_stats_path = self.output_dir / "dispersed_stats.csv"
+        self.rollout_details_path = self.output_dir / "rollout_details.jsonl"
+        self.status_file_path = self.output_dir / "status.txt"
+
+        self._load_status()
+        self._load_model()
+        self._prepare_tokenizer()
+
+    def _load_status(self):
+        """Load processing status from file"""
+        try:
+            with open(self.status_file_path, "r", encoding="utf-8") as f:
+                content = f.read().strip()
+                self.processed_set = RangeSet.from_file(content)
+                logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
+        except FileNotFoundError:
+            self.processed_set = RangeSet([])
+
+    def _save_status(self, batch_index):
+        """Save current processing status to file"""
+        self.processed_set.add(batch_index)
+        content = self.processed_set.to_file_format()
+        with open(self.status_file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    def _load_model(self):
+        logger.info(f"Loading model from {self.args.actor_model_name_or_path}...")
+        start_time = time.time()
+        self.model = LLM(
+            model=self.args.actor_model_name_or_path,
+            dtype=self.args.dtype,
+            enable_sleep_mode=False,
+            tensor_parallel_size=self.args.tensor_parallel_degree,
+            # distributed_executor_backend="external_launcher",
+            enforce_eager=False,
+            gpu_memory_utilization=self.args.gpu_memory_utilization,
+            disable_custom_all_reduce=True,
+            disable_mm_preprocessor_cache=True,
+            skip_tokenizer_init=False,
+            max_model_len=self.args.max_src_len + self.args.max_dec_len,
+            disable_log_stats=True,
+            max_num_batched_tokens=self.args.max_src_len + self.args.max_dec_len,
+            enable_chunked_prefill=True,
+            enable_prefix_caching=True,
+        )
+        self.sampling_params = SamplingParams(
+            top_p=self.args.top_p,
+            temperature=self.args.temperature,
+            min_tokens=self.args.min_dec_len,
+            max_tokens=self.args.max_dec_len,
+            n=self.args.rollout_n,
+            seed=42,
+        )
+        logger.info(f"Model loaded in {time.time() - start_time:.2f}s")
+
+    def _prepare_tokenizer(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.actor_model_name_or_path, use_fast=True)
+        self.tokenizer.padding_side = "left"
+
+    def process_data(self, file_path: str) -> pd.DataFrame:
+        logger.info(f"Processing data from {file_path}...")
+        start_time = time.time()
+        df = pd.read_parquet(file_path)
+        if self.args.limit_rows != -1:
+            df = df.iloc[: self.args.limit_rows]
+        logger.info(f"Loaded {len(df)} samples in {time.time() - start_time:.2f}s")
+        return df
+
+    def run_inference(self, prompts, batch_index=0):
+        input_ids = [self.tokenizer(prompt, add_special_tokens=False)["input_ids"] for prompt in prompts]
+        token_prompt_ids = [TokensPrompt(prompt_token_ids=prompt) for prompt in input_ids]
+
+        start_time = time.time()
+        request_outputs = self.model.generate(
+            prompts=token_prompt_ids, sampling_params=self.sampling_params, use_tqdm=True
+        )
+        end_time = time.time()
+
+        batch_token_ids = []
+        for output in request_outputs:
+            for each_output in output.outputs:
+                batch_token_ids.append(each_output.token_ids)
+
+        if self.args.should_log:
+            statistics = self.postprocess_data(input_ids, batch_token_ids, batch_index=batch_index)
+            statistics["total_time"] = end_time - start_time
+            return statistics
+        return None
+
+    def postprocess_data(self, input_ids, output_ids, batch_index=0):
+        global_prompt_tokens = []
+        global_prompt_tokens_len = []
+        group_prompt_texts = []
+        for row in input_ids:
+            row_ids = row
+            global_prompt_tokens.append(row_ids)
+            global_prompt_tokens_len.append(len(row_ids))
+            group_prompt_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Process responses
+        response_tokens = []
+        response_tokens_len = []
+        response_texts = []
+        for row in output_ids:
+            row_ids = row
+            response_tokens.append(row_ids)
+            response_tokens_len.append(len(row_ids))
+            response_texts.append(self.tokenizer.decode(row_ids, skip_special_tokens=False))
+
+        # Group processing
+        group_response_tokens = chunk(response_tokens, self.args.rollout_n)
+        group_response_tokens_len = chunk(response_tokens_len, self.args.rollout_n)
+        group_response_texts = chunk(response_texts, self.args.rollout_n)
+
+        # Calculate statistics
+        global_response_tokens_total = sum(response_tokens_len)
+        global_response_tokens_len_min = min(response_tokens_len)
+        global_response_tokens_len_max = max(response_tokens_len)
+        global_response_tokens_len_mean = round(global_response_tokens_total / len(response_tokens_len), 2)
+
+        group_response_tokens_len_mean = [round(sum(x) / len(x), 2) for x in group_response_tokens_len]
+        group_response_tokens_len_max = [max(x) for x in group_response_tokens_len]
+        group_response_tokens_len_min = [min(x) for x in group_response_tokens_len]
+
+        return {
+            # Global prompt stats
+            "batch_index": batch_index,
+            "global_prompt_tokens_len": global_prompt_tokens_len,
+            "group_prompt_texts": group_prompt_texts,
+            # Group response stats
+            "group_response_tokens": group_response_tokens,
+            "group_response_tokens_len": group_response_tokens_len,
+            "group_response_texts": group_response_texts,
+            "group_response_tokens_len_min": group_response_tokens_len_min,
+            "group_response_tokens_len_max": group_response_tokens_len_max,
+            "group_response_tokens_len_mean": group_response_tokens_len_mean,
+            # Global response stats
+            "global_response_tokens_total": global_response_tokens_total,
+            "global_response_tokens_len_min": global_response_tokens_len_min,
+            "global_response_tokens_len_max": global_response_tokens_len_max,
+            "global_response_tokens_len_mean": global_response_tokens_len_mean,
+        }
+
+    def batch_process(self, dataframe: pd.DataFrame):
+        all_prompts = []
+
+        for idx, row in dataframe.iterrows():
+            prompt = row["prompt"].tolist()[0]["content"]
+            all_prompts.append(prompt)
+
+            if len(all_prompts) >= self.args.rollout_input_batch_size:
+                yield all_prompts
+                all_prompts = []
+
+        if all_prompts:
+            yield all_prompts
+
+    def execute(self):
+        dataframe = self.process_data(self.args.input_file)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(self.global_stats_path, "a", newline="") as global_f, open(
+            self.dispersed_stats_path, "a", newline=""
+        ) as dispersed_f, open(self.rollout_details_path, "a", encoding="utf-8") as jsonl_f:
+
+            if self.args.should_log:
+                global_writer = csv.writer(global_f)
+                dispersed_writer = csv.writer(dispersed_f)
+                if self.processed_set.processed_count <= 0:
+                    global_writer.writerow(
+                        [
+                            "batch_index",
+                            "min_response_tokens",
+                            "max_response_tokens",
+                            "avg_response_tokens",
+                            "total_response_tokens",
+                            "group_max_response_tokens",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+                    dispersed_writer.writerow(
+                        [
+                            "batch_index",
+                            "rollout_lengths",
+                            "min_length",
+                            "max_length",
+                            "avg_length",
+                            "completion_time",
+                            "throughput_tokens_per_sec",
+                        ]
+                    )
+
+            for batch_index, prompts in tqdm.tqdm(
+                enumerate(self.batch_process(dataframe)),
+                total=math.ceil(len(dataframe) / self.args.rollout_input_batch_size),
+                disable=not self.args.should_log,
+            ):
+                if self.processed_set.contains(batch_index):
+                    continue
+
+                statistics = self.run_inference(prompts, batch_index=batch_index)
+
+                if self.args.should_log:
+                    total_time = round(statistics["total_time"], 2)
+                    total_tokens = statistics["global_response_tokens_total"]
+                    throughput = round(total_tokens / total_time if total_time > 0 else 0, 2)
+
+                    global_writer.writerow(
+                        [
+                            batch_index,
+                            statistics["global_response_tokens_len_min"],
+                            statistics["global_response_tokens_len_max"],
+                            statistics["global_response_tokens_len_mean"],
+                            total_tokens,
+                            statistics["group_response_tokens_len_max"],
+                            total_time,
+                            throughput,
+                        ]
+                    )
+
+                    dispersed_writer.writerow(
+                        [
+                            batch_index,
+                            statistics["group_response_tokens_len"],
+                            statistics["group_response_tokens_len_min"],
+                            statistics["group_response_tokens_len_max"],
+                            statistics["group_response_tokens_len_mean"],
+                            total_time,
+                            throughput,
+                        ]
+                    )
+
+                    prompt_text = statistics["group_prompt_texts"]
+                    response_tokens = statistics["group_response_tokens"]
+                    response_texts = statistics["group_response_texts"]
+
+                    record = {
+                        "batch_index": batch_index,
+                        "prompt_text": prompt_text,
+                        "rollouts": [
+                            {"token_ids": tokens, "text": text, "length": len(tokens)}
+                            for tokens, text in zip(response_tokens, response_texts)
+                        ],
+                        "total_time": total_time,
+                        "throughput_tokens_per_sec": throughput,
+                    }
+                    jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+                    global_f.flush()
+                    dispersed_f.flush()
+                    jsonl_f.flush()
+                    self._save_status(batch_index)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    print(args)
+    task = DumpyInferenceTask(args)
+    task.execute()

llm/benchmark/rl/utils.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class RangeSet:
+    """Manage processed line ranges with efficient storage and querying"""
+
+    ranges: List[tuple]
+
+    def add(self, number: int):
+        """Add a number to the range set and merge adjacent ranges"""
+        new_ranges = []
+        added = False
+        for start, end in sorted(self.ranges):
+            if number < start - 1:
+                if not added:
+                    new_ranges.append((number, number))
+                    added = True
+                new_ranges.append((start, end))
+            elif number == start - 1:
+                new_ranges.append((number, end))
+                added = True
+            elif number <= end:
+                new_ranges.append((start, end))
+                added = True
+            else:
+                new_ranges.append((start, end))
+        if not added:
+            new_ranges.append((number, number))
+        self.ranges = self.merge_ranges(new_ranges)
+
+    @staticmethod
+    def merge_ranges(ranges: List[tuple]) -> List[tuple]:
+        """Merge overlapping or adjacent ranges"""
+        if not ranges:
+            return []
+        sorted_ranges = sorted(ranges)
+        merged = [sorted_ranges[0]]
+        for current in sorted_ranges[1:]:
+            last = merged[-1]
+            if current[0] <= last[1] + 1:
+                merged[-1] = (last[0], max(last[1], current[1]))
+            else:
+                merged.append(current)
+        return merged
+
+    def contains(self, number: int) -> bool:
+        """Check if a number exists in any range"""
+        for start, end in self.ranges:
+            if start <= number <= end:
+                return True
+        return False
+
+    def to_file_format(self) -> str:
+        """Serialize ranges to compact string format"""
+        return ",".join(f"{start}-{end}" if start != end else str(start) for start, end in self.ranges)
+
+    @classmethod
+    def from_file(cls, content: str) -> "RangeSet":
+        """Deserialize from string format"""
+        if not content:
+            return cls(ranges=[])
+        ranges = []
+        for part in content.split(","):
+            if "-" in part:
+                start, end = map(int, part.split("-"))
+                ranges.append((start, end))
+            else:
+                num = int(part)
+                ranges.append((num, num))
+        return cls(ranges=ranges)
+
+    @property
+    def processed_count(self) -> int:
+        """Total number of processed items"""
+        return sum(end - start + 1 for start, end in self.ranges)