PullRequest: 994 修复奖励数据处理和对齐工具，新增日志和指标计算功能

Merge branch chucai.dzq/align-onpolicy of [email protected]:inclusionAI/AReaL.git into asystem/gh
https://code.alipay.com/inclusionAI/AReaL/pull_requests/994

Reviewed-by: 峯回 <[email protected]>


* align onploicy, modify dataload and rlvr

Author: 楚财

areal/api/io_struct.py

@@ -17,6 +17,41 @@
     from transformers import AutoProcessor
 
 
+@dataclass
+class LLMRequest:
+    rid: str = field(default_factory=lambda: str(uuid.uuid4()))
+    text: str | None = None
+    input_ids: list[int] = field(default_factory=list)
+    gconfig: GenerationHyperparameters = field(
+        default_factory=GenerationHyperparameters
+    )
+    metadata: dict[str, Any] = field(default_factory=dict)
+    model_id: str | None = None
+
+
+@dataclass
+class LLMResponse:
+    # outputs
+    input_tokens: list[int] = field(default_factory=list)
+    output_tokens: list[int] = field(default_factory=list)
+    output_logprobs: list[float] = field(default_factory=list)
+    output_version: int = field(default_factory=int)
+    stop_reason: Literal["length", "stop", "interrupt", "abort"] = "stop"
+
+    # statistics
+    latency: float = float("inf")
+    ttft: float = float("inf")  # Time to first token
+    itl: list[float] = field(default_factory=list)  # List of inter-token latencies
+
+    @property
+    def input_len(self) -> int:
+        return len(self.input_tokens)
+
+    @property
+    def output_len(self) -> int:
+        return len(self.output_tokens)
+
+
 @dataclass
 class ModelRequest:
     rid: str = field(default_factory=lambda: str(uuid.uuid4()))
@@ -208,6 +243,7 @@ class SaveLoadMeta:
     naive_distributed: bool = False
     global_step: int | None = None
 
+
 @dataclass
 class RolloutStat:
     submitted: int = 0

areal/examples/configs/my001/on_policy.yaml

@@ -19,9 +19,9 @@ train_dataset:
   type: "rl"
 
 scheduler:
-#  endpoint: "http://asystem-scheduler.asystem-my001-swift.svc.sigma-my001.ml01.sgp-ml.local:8081"
-  functioncall_service_domain: "http://110.75.237.19:8080"
+  #  endpoint: "http://asystem-scheduler.asystem-my001-swift.svc.sigma-my001.ml01.sgp-ml.local:8081"
   endpoint: "http://asystem-scheduler.asystem-cluster-prod-1.svc:8081"
+  functioncall_service_domain: "http://110.75.237.19:8080"
   reward_model_path: "/storage/jiulin.jl/Skywork-Reward-V2-Qwen3-8B"
   reward_model_service_url: "http://reward-model-service.asystem-test.svc.sigma-my001.ml01.sgp-ml.local:30000/classify"
 
@@ -43,7 +43,6 @@ 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

@@ -1,19 +1,21 @@
+import concurrent.futures
 import json
 import os
 import pprint
-import sys
-import asyncio
 import shutil
-import concurrent.futures
+import sys
 from concurrent.futures import ThreadPoolExecutor
 
 from datasets import load_dataset
 from torchdata.stateful_dataloader import StatefulDataLoader
 
+from realhf.api.core.data_api import load_hf_tokenizer
+
 from areal.api.cli_args import (
     SchedulingStrategy,
     load_expr_config,
 )
+from areal.api.engine_api import WeightUpdateMeta
 from areal.api.io_struct import AllocationMode, FinetuneSpec
 from areal.extension.asystem.api.cli_args import GRPOConfig
 from areal.extension.asystem.ascheduler import AsystemScheduler
@@ -23,23 +25,14 @@
     RemoteHybridInferenceWorker,
 )
 from areal.extension.asystem.remote_hybrid_train_worker import RemoteHybridTrainWorker
-from areal.extension.asystem.util import ShuffleSampler, wait_future_ordered
+from areal.extension.asystem.utils.align_tools import summarize_rewards
+from areal.extension.asystem.utils.util import ShuffleSampler, wait_future_ordered
 from areal.utils import logging, stats_tracker
-from areal.utils.hf_utils import load_hf_tokenizer
 from areal.utils.stats_logger import StatsLogger
-from areal.api.engine_api import WeightUpdateMeta
 
 logger = logging.getLogger("Trainer")
 
 
-def custom_collate_fn(batch):
-    all_keys = set().union(*(d.keys() for d in batch))
-    collated_batch = {}
-    for key in all_keys:
-        collated_batch[key] = [d.get(key) for d in batch]
-    return collated_batch
-
-
 def clear_dir(path):
     if os.path.exists(path):
         for filename in os.listdir(path):
@@ -54,13 +47,9 @@ def main(args):
     config, _ = load_expr_config(args, GRPOConfig)
     config: GRPOConfig
 
-    if config.gconfig.max_tokens is None:
-        logger.info(
-            "config.gconfig.max_tokens is None, set it to max_new_tokens + max_prompt_len"
-        )
-        config.gconfig.max_tokens = (
-            config.gconfig.max_new_tokens + config.train_dataset.max_length
-        )
+    config.gconfig.max_tokens = (
+        config.gconfig.max_new_tokens + config.train_dataset.max_length
+    )
 
     if config.enable_colocate_mode:
         config.rollout.engine_config["enable_memory_saver"] = True
@@ -122,22 +111,9 @@ def main(args):
         train_dataset = dataset["train"]
         train_dataset = train_dataset.filter(
             lambda x: len(tokenizer.encode(x["prompt"]))
-                      <= config.train_dataset.max_length
+            <= config.train_dataset.max_length
         )
 
-        def process(sample):
-            messages = [
-                {
-                    "role": "user",
-                    "content": sample["prompt"]
-                    .replace("<role>HUMAN</role>", "")
-                    .replace("<role>ASSISTANT</role>", ""),
-                }
-            ]
-            return {"messages": messages}
-
-        train_dataset = train_dataset.map(process).remove_columns(["prompt"])
-
         dataloader = StatefulDataLoader(
             train_dataset,
             batch_size=config.train_dataset.batch_size,
@@ -221,16 +197,21 @@ def process(sample):
         if config.actor.hybrid_engine.wrap_policy.kl_ctl > 0:
             ref = TrainController(
                 RemoteHybridTrainWorker,
-                config.actor,
+                config.ref,
                 scheduler,
             )
 
         allocation_mode = AllocationMode.from_str(config.allocation_mode)
 
         def init_train_and_rollout_controller_helper(actor, rollout):
             logger.info("initializing trainer controller and rollout controller")
-            actor.initialize(role="actor", alloc_mode=allocation_mode, ft_spec=ft_spec,
-                             group_size=config.gconfig.n_samples, )
+            actor.initialize(
+                role="actor",
+                alloc_mode=allocation_mode,
+                ft_spec=ft_spec,
+                group_size=config.gconfig.n_samples,
+                enable_colocate_mode=config.enable_colocate_mode,
+            )
             rollout.initialize(role="rollout", alloc_mode=allocation_mode)
 
         if config.enable_colocate_mode:
@@ -254,7 +235,9 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                     )
 
                 wait_future_ordered(futures)
-            logger.info(f"initialized all controllers in colocation mode {config.enable_colocate_mode}")
+            logger.info(
+                f"initialized all controllers in colocation mode {config.enable_colocate_mode}"
+            )
         else:
             with ThreadPoolExecutor(max_workers=3) as executor:
                 futures = [
@@ -268,7 +251,10 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                         storage_prefix=config.storage_prefix,
                     ),
                     executor.submit(
-                        rollout.initialize, role="rollout", alloc_mode=allocation_mode
+                        rollout.initialize,
+                        role="rollout",
+                        alloc_mode=allocation_mode,
+                        enable_colocate_mode=config.enable_colocate_mode,
                     ),
                 ]
                 if ref is not None:
@@ -324,15 +310,19 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                             )
                             clear_dir(weight_update_config.path)
                         else:
-                            with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+                            with concurrent.futures.ThreadPoolExecutor(
+                                max_workers=2
+                            ) as executor:
                                 upload_future = executor.submit(
                                     actor.upload_weights, weight_update_config
                                 )
                                 update_future = executor.submit(
                                     rollout.update_weights, weight_update_config
                                 )
                                 wait_future_ordered([upload_future, update_future])
-                            logger.info(f"{weight_update_config.type} update weight succeeded, step: {step}")
+                            logger.info(
+                                f"{weight_update_config.type} update weight succeeded, step: {step}"
+                            )
 
                     with (
                         stats_tracker.record_timing("rollout_step"),
@@ -360,11 +350,8 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                                         reward_fn="areal.extension.asystem.math_reward.reward_fn",
                                         gconfig=config.gconfig,
                                         tokenizer=config.tokenizer_path,
-                                        enable_thinking=False,
-                                        dump_dir=os.path.join(
-                                            f"{config.storage_prefix}/experiments/logs/root/{config.experiment_name}/{config.trial_name}",
-                                            "generated",
-                                        ),
+                                        exp_name=config.experiment_name,
+                                        trial_name=config.trial_name,
                                     ),
                                 )
                             else:
@@ -375,24 +362,27 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                                         reward_fn="areal.extension.asystem.math_reward.reward_fn",
                                         gconfig=config.gconfig,
                                         tokenizer=config.tokenizer_path,
-                                        enable_thinking=False,
-                                        dump_dir=os.path.join(
-                                            f"{config.storage_prefix}/experiments/logs/root/{config.experiment_name}/{config.trial_name}",
-                                            "generated",
-                                        ),
+                                        exp_name=config.experiment_name,
+                                        trial_name=config.trial_name,
                                     ),
                                 )
 
-                    #TODO: calc_training_data_metrics
-                    # with (stats_tracker.scope("training_data"), ):
-                    #     calc_training_data_metrics(rollout_res)
+                    # with (
+                    #     stats_tracker.scope("training_data"),
+                    # ):
+                    #     calc_training_data_metrics(batch.get_data())
                     #     calc_training_data_group_metrics(
-                    #         rollout_res, config.gconfig.n_samples
+                    #         batch.get_data, config.gconfig.n_samples
                     #     )
-                    #     calc_training_data_version_metrics(rollout_res, global_step)
-                    #
-                    logger.info(f"rollout batch res: {batch}, reward: {batch["rewards"]}")
-                    with (stats_tracker.record_timing("notify_rollout_end_event"), ):
+                    #     calc_training_data_version_metrics(batch.get_data, global_step)
+
+                    logger.info(
+                        "rollout batch reward summary: %s",
+                        summarize_rewards(batch["rewards"]),
+                    )
+                    with (
+                        stats_tracker.record_timing("notify_rollout_end_event"),
+                    ):
                         logger.info(
                             f"start to notify_rollout_end_event, step: {step}, epoch: {epoch}"
                         )
@@ -420,7 +410,9 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                         stats_tracker.record_timing("train_step"),
                         stats_tracker.scope("train"),
                     ):
-                        with (stats_tracker.record_timing("notify_train_start_event"),):
+                        with (
+                            stats_tracker.record_timing("notify_train_start_event"),
+                        ):
                             logger.info(
                                 f"start to notify_train_start_event, step: {step}, epoch: {epoch}"
                             )
@@ -429,12 +421,14 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                                 f"notify_train_start_event succeeded, step: {step}, epoch: {epoch}"
                             )
 
-                        with (stats_tracker.record_timing("train_distributed_batch"), ):
+                        with (
+                            stats_tracker.record_timing("train_distributed_batch"),
+                        ):
                             logger.info(f"start to train, step: {step}, epoch: {epoch}")
                             actor.train_batch(
                                 batch,
                                 loss_fn=lambda logits, batch_data: None,
-                                loss_weight_fn=lambda batch_data: None
+                                loss_weight_fn=lambda batch_data: None,
                             )
                             logger.info(
                                 f"train succeeded, step: {step}, epoch: {epoch}"
@@ -487,7 +481,9 @@ def init_train_and_rollout_controller_helper(actor, rollout):
                                     f"[Trainer] periodic_checkpoint recover save success, epoch:{epoch}, epoch_step: {step}, global_step:{global_step}"
                                 )
 
-                        with (stats_tracker.record_timing("notify_train_end_event"),):
+                        with (
+                            stats_tracker.record_timing("notify_train_end_event"),
+                        ):
                             logger.info(
                                 f"start to notify_train_end_event, step: {step}, epoch: {epoch}"
                             )