PullRequest: 941 新增GSM8K奖励函数并修改训练配置和流程

Merge branch test/tmp1102 of [email protected]:inclusionAI/AReaL.git into asystem/gh
https://code.alipay.com/inclusionAI/AReaL/pull_requests/941

Reviewed-by: 楚财 <[email protected]>


* .
* .
* .
* .
* .
* .

Author: 峯回

areal/controller/rollout_controller.py

@@ -335,7 +335,7 @@ def _commit_one_to_runner(self):
 
         # Choose worker via round-robin
         worker = self._choose_worker()
-
+        self.logger.info(f"Submit rollout to worker {worker.id}, task_input: {task_input}")
         self.scheduler.call_engine(
             worker.id,
             "submit",

areal/examples/configs/my001/on_policy.yaml

@@ -41,6 +41,7 @@ gconfig:
   # Due to the limitations of sglang, max_new_tokens + max_prompt_len must be less than the model's context_len (set in the model's config.json),
   # and cannot be equal to it. See https://github.com/sgl-project/sglang/blob/f98366604b23e331422bf3c62d4e7410ae4fab87/python/sglang/srt/managers/tokenizer_manager.py#L638C9-L638C11
   max_new_tokens: 15360
+  max_tokens: 16383
   greedy: false
   temperature: 1.0
   top_k: 1000000

areal/examples/grpo_trainer.py

@@ -97,7 +97,7 @@ def __init__(self, config: dict[str, Any]):
             f"AsystemScheduler initialized for {self.run_name}. API URL: {self.api_url}"
         )
 
-    def batch_cleanup_jobs(self, signum):
+    def batch_cleanup_jobs(self, signum, frame):
         logger.info(f"signum {signum} received: handle_signals starts")
         for role, job_uid in self.submitted_jobs.items():
             try:

areal/extension/asystem/ascheduler/__init__.py

@@ -115,7 +115,7 @@ async def async_create_engine(self, worker_id, engine, *args, **kwargs):
                 last_exception = EngineCreationError(
                     worker_id, f"Connection error: {str(e)}"
                 )
-                logger.error(f"Connection error on attempt {attempt + 1}: {e}")
+                logger.warning(f"Connection error on attempt {attempt + 1}: {e}")
 
             except httpx.TimeoutException as e:
                 # Timeout errors are retryable
@@ -140,7 +140,7 @@ async def async_create_engine(self, worker_id, engine, *args, **kwargs):
             if last_exception is not None:
                 if attempt < max_retries - 1:
                     logger.warning(
-                        f"Retrying create_engine in 1 second... ({attempt + 1}/{max_retries})"
+                        f"Retrying create_engine in 5 second... ({attempt + 1}/{max_retries})"
                     )
                     await asyncio.sleep(5)
                     continue

areal/extension/asystem/ascheduler/rpc_client.py

@@ -47,9 +47,9 @@ if [[ -n "${PORT_LIST}" ]]; then
     # 获取第一个端口
     FIRST_PORT="${PORTS[0]}"
     # 添加到 WORKER_COMMAND
-    WORKER_COMMAND="/usr/bin/python -u -m areal.scheduler.rpc.async_rpc_server --worker-type ${WORKER_TYPE} --worker-index ${WORKER_INDEX} --port ${FIRST_PORT}"
+    WORKER_COMMAND="/usr/bin/python -u -m areal.scheduler.rpc.rpc_server --worker-type ${WORKER_TYPE} --worker-index ${WORKER_INDEX} --port ${FIRST_PORT}"
 else
-    WORKER_COMMAND="/usr/bin/python -u -m areal.scheduler.rpc.async_rpc_server --worker-type ${WORKER_TYPE} --worker-index ${WORKER_INDEX}"
+    WORKER_COMMAND="/usr/bin/python -u -m areal.scheduler.rpc.rpc_server --worker-type ${WORKER_TYPE} --worker-index ${WORKER_INDEX}"
 fi
 
 #log output to local worker dir

areal/extension/asystem/ascheduler/scripts/launch-worker.sh

@@ -38,8 +38,8 @@ async def reward_fn(
 
     format_rewards = []
 
-    query_id = kwargs.get("query_id")[0]
-    task = kwargs.get("task")[0]
+    query_id = kwargs.get("query_id")
+    task = kwargs.get("task")
     answers = [completion]
     query_id_strs = [query_id]
 
@@ -59,7 +59,7 @@ async def reward_fn(
     elif task == "ifeval":
         format_rewards = await ifeval_verify(id2info, answers, query_id_strs)
     elif task == "swe":
-        extra_info = kwargs.get("extra_info")[0]
+        extra_info = kwargs.get("extra_info")
         if extra_info and extra_info.get("provider", "functioncall") == "local":
             format_rewards = await local_swe_verify(id2info, answers, query_id_strs)
         else:
@@ -224,8 +224,8 @@ def extract_python_code(text, min_length=20, strict_syntax=False):
     async def main():
         answer = "<answer>\n28\n</answer>"
         data = {
-            "task": ["general"],
-            "query_id": ["general-42941"],
+            "task": "general",
+            "query_id": "general-42941",
             "prompt": [
                 "<role>HUMAN</role>33岁孩子不听话,如何处理父子之间矛盾?<role>ASSISTANT</role>"
             ],

areal/extension/asystem/math_reward.py

@@ -6,7 +6,7 @@
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from dataclasses import dataclass
 from queue import Empty, Full, Queue
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
 import aiohttp
 import requests
@@ -22,15 +22,14 @@
     RolloutStat,
     WeightUpdateMeta,
 )
+from areal.api.workflow_api import RolloutWorkflow
 from areal.extension.asystem.api.cli_args import RemoteHybridInferenceConfig
 from areal.extension.asystem.util import wait_future_ordered
 from areal.utils import logging, seeding
 from areal.utils.data import concat_padded_tensors, cycle_dataloader
 from areal.utils.errors import EngineError, FrameworkError
 from areal.utils.http import arequest_with_retry, get_default_connector
 
-if TYPE_CHECKING:
-    from areal.api.workflow_api import RolloutWorkflow
 logger = logging.getLogger(__name__)
 
 ROLLOUT_POLL_WAIT_TIME = 0.05
@@ -236,9 +235,7 @@ async def _rollout_thread_async(self):
                     # logger.info(f"Get data from puller: {data}")
                     task = asyncio.create_task(
                         (
-                            workflow.arun_episodes(self, data)
-                            if isinstance(data, list)
-                            else workflow.arun_episode(self, data)
+                            workflow.arun_episode(self, data)
                         ),
                         name=str(rid),
                     )
@@ -345,6 +342,7 @@ async def agenerate(self, req: ModelRequest) -> ModelResponse:
         start_time = time.perf_counter()
         accumulated_output_tokens = []
         accumulated_output_logprobs = []
+        accumulated_versions = []
 
         # Deal with rollout interruption
         stop_reason = ""
@@ -385,6 +383,9 @@ async def agenerate(self, req: ModelRequest) -> ModelResponse:
             # Update accumulated outputs
             accumulated_output_tokens.extend(output_tokens)
             accumulated_output_logprobs.extend(output_logprobs)
+            accumulated_versions.extend(
+                [self.get_version()] * len(output_logprobs)
+            )
 
             # Check if generation is complete
             finish_reason = meta_info["finish_reason"]
@@ -399,7 +400,7 @@ async def agenerate(self, req: ModelRequest) -> ModelResponse:
             input_tokens=req.input_ids,
             output_tokens=accumulated_output_tokens,
             output_logprobs=accumulated_output_logprobs,
-            output_version=self.get_version(),
+            output_versions=accumulated_versions,
             stop_reason=stop_reason,
             latency=latency,
             ttft=latency,  # Simplified for non-streaming
@@ -532,14 +533,13 @@ def update_weights_from_disk(self, addr, path: str):
 
     def submit(
         self,
-        data: list[dict[str, Any]] | dict[str, Any],
-        workflow: "RolloutWorkflow",
+        data: dict[str, Any],
+        workflow: RolloutWorkflow | None = None,
+        workflow_builder: Callable | None = None,
+        should_accept: Callable | None = None,
     ) -> None:
         try:
-            if not isinstance(data, list):
-                data = [data]
-            for d in data:
-                self.input_queue.put_nowait((d, workflow))
+            self.input_queue.put_nowait((data, workflow))
         except Full:
             raise FrameworkError(
                 "FrameworkError",
@@ -548,7 +548,7 @@ def submit(
             )
 
     def submit_batch(
-        self, data: list[dict[str, Any]], workflow: "RolloutWorkflow"
+        self, data: list[dict[str, Any]], workflow: RolloutWorkflow
     ) -> None:
         try:
             self.input_queue.put_nowait(data, workflow)
@@ -701,3 +701,11 @@ def notify_event(self, event: str, global_step: int) -> None:
         except Exception as e:
             raise EngineError("InferenceEngineError", "NotifyEventError", e)
         return None
+
+    def wait_quiet(
+        self, count: int, timeout: float | None = None, max_retries: int = 1,
+    ) -> dict[str, Any] | None:
+        try:
+            return self.wait(count, timeout=timeout)
+        except TimeoutError:
+            return "NO_RESULT"

areal/extension/asystem/remote_hybrid_inference_worker.py

@@ -53,8 +53,7 @@ def __init__(self, data_source, seed=42):
         self.shuffle_indices = get_shuffle_indices(size=len(data_source), seed=seed)
 
     def __iter__(self):
-        for idx in self.shuffle_indices:
-            yield from idx
+        return iter(self.shuffle_indices)
 
     def __len__(self):
         return len(self.data_source)

areal/extension/asystem/util.py

@@ -0,0 +1,4 @@
+def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs):
+    from areal.reward.math_parser import process_results
+
+    return int(process_results(completions, answer)[0])