[bugfix] clip eos_token in judger and fix dataflow retry logic (#1189)

* fix dapo_math judger and fix dataflow logic

* fix dapo_math judger and fix dataflow logic

* fix typo

* rm data

* add comments

* fix comments

* fix

* fix

* fix

* fix ci env

* fix

Author: YanhuiDua

ci/scripts/CI_ENV.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 export QWEN3_VL_MOE_PATH=${CI_SHARE_MODEL}/Qwen3-VL-30B-A3B-Instruct_MOE
-export QWEN3_VL_DENSE_PATH=${CI_SHARE_MODEL}/Qwen3-VL-8B-Instruct_DENSE
+export QWEN3_VL_DENSE_PATH=${CI_SHARE_MODEL}/Qwen3-VL-4B-Instruct
 export INTERN_VL_1B_PATH=${CI_SHARE_MODEL}/InternVL3_5-1B-HF
 export VIDEO_ROOT=${CI_SHARE_DATA}/images
 export QWEN3_4B_PATH=${CI_SHARE_MODEL}/Qwen3-4B-Instruct-2507
@@ -16,6 +16,7 @@ export INTERNS1_DENSE_PATH=${CI_SHARE_MODEL}/intern-s1-mini
 export ROLLOUT_MODEL_PATH=${CI_SHARE_MODEL}/Qwen3-8B
 export ALPACA_PATH=${CI_SHARE_DATA}/alpaca
 export INTERNS1_DATA_META=${CI_SHARE_DATA}/vlm_ci_data.json
+export ROLLOUT_DAPO_DATA_PATH=${CI_SHARE_DATA}/rl_test_judger_dapo_math_data.jsonl
 export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0
 export XTUNER_DETERMINISTIC=true
 export XTUNER_USE_LMDEPLOY=1

examples/v1/config/rl_qwen25_7B_dapo.py

@@ -86,8 +86,12 @@
 dataloader_config = DataloaderConfig(pack_max_length=pack_max_length, collator="fake_collator", pack_level="none")
 
 # 3. judger
+from xtuner.v1.utils.rl_test_utils import get_eos_token
+eos_token_id = get_eos_token(model_path)
+eos_token_str = tokenizer.convert_ids_to_tokens(eos_token_id)
 dapomath_judger_config = DapoMathJudgerConfig(
-    judger_name = "dapo_math", 
+    judger_name="dapo_math", 
+    eos_token=eos_token_str,
     enable_overlong_buffer = True, 
     max_response_len =max_response_length, 
     overlong_buffer_len=4096, 

tests/ray/test_evaluator.py

@@ -3,7 +3,6 @@
 import ray
 from transformers import AutoTokenizer
 
-
 from xtuner.v1.ray.config.worker import RolloutConfig
 from xtuner.v1.ray.judger.controller import JudgerConfig
 from xtuner.v1.ray.base import AcceleratorResourcesConfig, AutoAcceleratorWorkers
@@ -110,6 +109,19 @@ def custom_compute_metric(samples):
         custom_correctness = ray.get(custom_evaluator.run.remote())
         self.assertEqual(correctness['accuracy'], custom_correctness['custom_accuracy'])
 
-
+    @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")
+    def test_lmdeploy_evaluator_with_failed_response(self):
+        evaluator_cfg = EvaluatorConfig(
+            dataset_cfg=self.eval_dataset_cfg,
+            tokenizer=self.tokenizer,
+            max_concurrent=1,
+            eval_sample_ratio=1,  # generate 5 samples
+            sample_params=SampleParams(temperature=2.5),  # invalid temperature to trigger error
+            max_retry_times=1,
+        )
+        evaluator = Evaluator.remote(evaluator_cfg, self.test_env)
+        correctness = ray.get(evaluator.run.remote())
+        self.assertEqual(len(correctness), 0)
+        
 if __name__ == '__main__':
     unittest.main()

tests/ray/test_judger.py

@@ -17,6 +17,7 @@
 MODEL_PATH = os.environ["ROLLOUT_MODEL_PATH"]
 DATA_PATH = os.environ["ROLLOUT_DATA_PATH"]
 VERL_ROLLOUT_DATA_PATH = os.environ["VERL_ROLLOUT_DATA_PATH"]
+DAPO_DATA_PATH = os.environ.get("ROLLOUT_DAPO_DATA_PATH")
 
 FAKE_JUDGER_INPUT_ITEM = RLDataFlowItem(
     uid = RLUIDItem(action_id=uuid4().int,
@@ -67,6 +68,39 @@ def construct_judger_data(data_path):
             dataitem.append(data_item)
     return dataitem
 
+def construct_dapo_judger_data(data_path):
+    data_item_list = []
+    save_reward = []
+    with open(data_path, 'r', encoding='utf-8') as f:
+        lines = f.readlines()
+        for i in range(0, len(lines), 7):
+            group = ''.join(lines[i:i+7]).strip()
+            if group:
+                try:
+                    item = json.loads(group)
+                    data_item = RLDataFlowItem(
+                        uid = RLUIDItem(
+                            action_id=uuid4().int,
+                            observation_id=uuid4().int
+                            ),
+                        data = RLDatasetItem(
+                            messages=[{
+                                'role': 'user', 
+                                'content': ""
+                            }],
+                            reward_model={"ground_truth": item["label"]},
+                            data_source={"dapo_math": 1.0}
+                        ),
+                        env = RLEnvDataItem(
+                            rollout=RLRolloutResponseItem(response=item['response'])
+                        )
+                    )
+                    data_item_list.append(data_item)
+                    save_reward.append(item["reward"])
+                except Exception as e:
+                    print(f"Error parsing group starting at line {i+12}: {e}")
+    return data_item_list, save_reward
+
 class TestJudgerController(unittest.TestCase):
 
     @classmethod
@@ -101,6 +135,33 @@ def test_gsm8k_judger(self):
         self.assertEqual(res2[0].reward["score"], 1.0)
         self.assertEqual(res2[1].reward["score"], 1.0)
 
+    def test_dapo_judger(self):
+        from xtuner.v1.ray.judger.dapo_math import DapoMathJudgerConfig
+        from xtuner.v1.utils.rl_test_utils import get_eos_token_from_model_path
+        from transformers import AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+        eos_token_str = get_eos_token_from_model_path(MODEL_PATH, tokenizer)
+
+        dapo_judger_config = DapoMathJudgerConfig(
+            judger_name="dapo_math", 
+            eos_token=eos_token_str,
+            enable_overlong_buffer=True, 
+            max_response_len=32768, 
+            overlong_buffer_len=4096, 
+            overlong_penalty_factor=1.0, 
+            tokenizer=tokenizer
+        
+        )
+        judger_cfg = JudgerConfig(
+            reward_judger_configs=[dapo_judger_config]
+        )
+        judger_controller = JudgerController.remote(judger_cfg)
+        judger_data, save_reward = construct_dapo_judger_data(DAPO_DATA_PATH)
+        group_data = ray.get(judger_controller.run.remote(judger_data)) 
+        reward = [data.reward["score"] for data in group_data]
+        avg_score = np.mean(reward)
+        self.assertLessEqual(float(np.abs(avg_score - np.mean(save_reward))), 0.001)
+
     def test_gsm8k_multi_judger(self):
         from xtuner.v1.ray.judger.gsm8k import GSM8KJudgerConfig
         # 支持一个GSM8KJudgerConfig创建多个实例

tests/ray/test_rollout.py

@@ -104,6 +104,30 @@ def setUp(self):
     def tearDown(self):
         ray.shutdown()
 
+    @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")
+    def test_lmdeploy_dataflow_with_failed_response(self):
+        failed_dataflow_cfg = DataFlowConfig(
+            env="test",
+            max_concurrent=1,
+            prompt_repeat_k=2,
+            global_batch_size=1,
+            enable_partial_rollout=0,
+            max_retry_times=1,
+        )
+        self.test_env = SingleTurnEnvironment.remote(
+            "test_env",
+            self.pg,
+            rollout_cfg=self.rollout_cfg,
+        )
+        self.test_flow = DataFlow.remote("test_env",
+                                        failed_dataflow_cfg,
+                                        self.replay_buffer_cfg,
+                                        self.test_env
+                                        )
+        sample_params = SampleParams(temperature=2.5)  # invalid temperature to trigger error
+        responses = ray.get(self.test_flow.run.remote(num=1, sample_params=sample_params), timeout=300)
+        self.assertEqual(len(responses),0)
+  
     @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")
     def test_lmdeploy_generate(self):
         sample_params = SampleParams(temperature=0.0)

xtuner/v1/ray/base/accelerator.py

@@ -72,7 +72,9 @@ def __init__(self, **kwargs):
         available_memory = available_resources.get("memory", 0)
         available_gpus = available_resources.get("GPU", 0)
 
-        assert kwargs["num_workers"] <= available_gpus, "Not enough available GPUS in Ray cluster."
+        assert kwargs["num_workers"] <= available_gpus, (
+            f"Not enough available GPUS in Ray cluster, available_gpus is {available_gpus} but xtuner needs {kwargs['num_workers']}."
+        )
         # TODO: manage single controller's cpu resource to replace "10" here
         assert (kwargs["num_cpus_per_worker"] * kwargs["num_workers"]) + 10 <= available_cpus, (
             f"Not enough available CPUs in Ray cluster, available_cpus is {available_cpus} but xtuner needs {kwargs['num_cpus_per_worker'] * kwargs['num_workers'] + 10}."

xtuner/v1/ray/dataflow/flow.py

@@ -8,7 +8,7 @@
 from tqdm.auto import tqdm
 from typing_extensions import Annotated
 
-from xtuner.v1.data_proto.rl_data import RLDataFlowItem
+from xtuner.v1.data_proto.rl_data import RLDataFlowItem, check_dataflow_item
 from xtuner.v1.ray.environment import SingleTurnEnvironment
 from xtuner.v1.ray.rollout.controller import SampleParams
 from xtuner.v1.ray.utils import create_task
@@ -141,29 +141,35 @@ async def worker_task(self, group_samples_for_retry: Optional[List[RLDataFlowIte
             Optional[List[RLDataFlowItem]]: The group of samples if the task
             fails and needs to be retried, otherwise None.
         """
-        if group_samples_for_retry is not None:
-            for data_item in group_samples_for_retry:
-                data_item.extra_info.retry_times += 1
-
-        group_data_items = group_samples_for_retry
         try:
-            # 该函数中所有的数据结构都是RLDataFlowItem
             # step 1: sample
-            with timer("sample", self.timer_dict):
-                group_data_items = await self.replay_buffer.sample.remote(  # type: ignore[attr-defined]
-                    self.env,
-                    self.config.enable_partial_rollout,
-                    self.config.prompt_repeat_k,
-                )
-                self.send_samples_count += 1
-                self.logger.debug(
-                    f"[ROLLOUT] Get 1 sample and dataflow have sent {self.send_samples_count} to rollout_controller"
-                )
+            # TODO(@duanyanhui): More fine-grained control over group data generation:
+            # Pass n to the inference engine to ensure that the same data is processed by the same server, improving efficiency
+            # Resend only the failed prompts in a group when retrying worker_task to avoid wasted computation resources."
+            if group_samples_for_retry is None or len(group_samples_for_retry) == 0:
+                with timer("sample", self.timer_dict):
+                    group_data_items = await self.replay_buffer.sample.remote(  # type: ignore[attr-defined]
+                        self.env,
+                        self.config.enable_partial_rollout,
+                        self.config.prompt_repeat_k,
+                    )
+                    self.send_samples_count += 1
+                    self.logger.debug(
+                        f"[ROLLOUT] Get 1 sample and dataflow have sent {self.send_samples_count} to rollout_controller"
+                    )
+            else:
+                group_data_items = group_samples_for_retry
+                for data_item in group_samples_for_retry:
+                    data_item.extra_info.retry_times += 1
+
             # step 2: env generate
             with timer("generate", self.timer_dict):
                 group_data_items = await self.env_controller.run.remote(  # type: ignore[attr-defined]
                     group_data_items, sample_params=self.sample_params, extra_params=self.extra_params
                 )
+                # 需要在这里处理check_dataflow_item，因为要保留group_data_items的data信息，作为retry的输入
+                if not check_dataflow_item(group_data_items):
+                    return group_data_items
 
             # step 3: filter
             with timer("post_process", self.timer_dict):
@@ -175,8 +181,6 @@ async def worker_task(self, group_samples_for_retry: Optional[List[RLDataFlowIte
 
         except Exception as e:
             self.logger.error(f"Worker task failed with exception: {e}. Returning meta for retry.", exc_info=True)
-            for sample in group_data_items:  # type: ignore[union-attr]
-                sample.extra_info.retry_times += 1
             return group_data_items
 
     async def concurrent_task_runner(self):
@@ -204,7 +208,10 @@ async def concurrent_task_runner(self):
         with tqdm(total=self.target_batch_size, desc="rollout_controller for training samples") as pbar:
             update_step = max(1, int(self.target_batch_size * 0.1))
             next_update_threshold = update_step
-            while self.finished_samples_count < self.target_batch_size:
+            while (
+                self.finished_samples_count < self.target_batch_size
+                and self.failed_samples_count < self.target_batch_size
+            ):
                 if self.finished_samples_count >= next_update_threshold:
                     pbar.n = self.finished_samples_count
                     pbar.refresh()
@@ -227,27 +234,36 @@ async def concurrent_task_runner(self):
                     if result is not None:
                         if result[0].extra_info.retry_times < self.config.max_retry_times:
                             # If the retry count is less than max_retry_times, retry the task
+                            self.logger.info(
+                                f"Retrying task for {result[0].data}. Retry count: {result[0].extra_info.retry_times}"
+                            )
                             retry_task = create_task(self.worker_task(group_samples_for_retry=result))
                             pending_tasks.add(retry_task)
                         else:
-                            self.logger.error(f"Max retry reached for {result[0]['prompt_id']}. Not retrying.")
                             self.failed_samples_count += 1
-
+                            self.logger.error(
+                                f"Max retry reached for {result[0].data}. Not retrying. Current failed count: {self.failed_samples_count}"
+                            )
                 self.finished_samples_count = ray.get(self.replay_buffer.get_finished_samples.remote())
                 waiting_tasks = pending_tasks
 
             pbar.n = self.finished_samples_count
             pbar.refresh()
 
-        self.logger.info("Target batch size reached. Pausing env controller.")
+        if self.finished_samples_count == self.target_batch_size:
+            self.logger.info("Target batch size reached. Pausing env controller.")
+        if self.failed_samples_count == self.target_batch_size:
+            self.logger.info("Max failed samples reached. Pausing env controller.")
+
         ray.get(self.env_controller.pause.remote())
 
         if waiting_tasks:
             await asyncio.wait_for(asyncio.gather(*waiting_tasks, return_exceptions=True), timeout=10)
 
-        self.unfinished_samples_count = ray.get(self.replay_buffer.get_unfinished_samples.remote())
-        self.logging_replaybuffer_state()
-        self.logging_timing_perf()
+        if self.finished_samples_count == self.target_batch_size:
+            self.unfinished_samples_count = ray.get(self.replay_buffer.get_unfinished_samples.remote())
+            self.logging_replaybuffer_state()
+            self.logging_timing_perf()
 
     async def run(
         self,

xtuner/v1/ray/dataflow/replay_buffer.py

@@ -318,7 +318,7 @@ def get(self, global_batch_size: int) -> List[List[RLDataFlowItem]]:
         """
         samples = []
         if len(self._returned) < global_batch_size:
-            raise ValueError("Not enough finished samples in replay buffer")
+            self.logger.error("Not enough finished samples in replay buffer")
             return []
         else:
             target_finished_list = self._returned[:global_batch_size]

xtuner/v1/ray/environment/single_turn_env.py

@@ -7,7 +7,6 @@
 from xtuner.v1.data_proto.rl_data import (
     RLDataFlowItem,
     RLJudgerResponseItem,
-    check_dataflow_item,
     update_dataflow_item,
 )
 from xtuner.v1.ray.environment.base_env import BaseEnvironment
@@ -92,6 +91,4 @@ async def run(
         if self.judger_controller:
             judger_responses: RLJudgerResponseItem = await self.judger_controller.run.remote(group_data_items)
             group_data_items = update_dataflow_item(group_data_items, "env.judger", judger_responses)
-        if not check_dataflow_item(group_data_items):
-            return []
         return group_data_items