PullRequest: 970 onpolicy sync training

Merge branch chucai.dzq/update-weights-debug of [email protected]:inclusionAI/AReaL.git into asystem/gh
https://code.alipay.com/inclusionAI/AReaL/pull_requests/970

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


* onpolicy sync training

Author: 楚财

areal/api/io_struct.py

@@ -206,7 +206,7 @@ class SaveLoadMeta:
     processor: Optional["AutoProcessor"] = None
     base_model_path: str | None = None
     naive_distributed: bool = False
-
+    global_step: int | None = None
 
 @dataclass
 class RolloutStat:

areal/examples/configs/my001/on_policy.yaml

@@ -1,9 +1,9 @@
-experiment_name: mini-model
-trial_name: test
-allocation_mode: "sglang:d2t4p1+d2t1p4"
+experiment_name: mini-model-rebased
+trial_name: on-policy
+allocation_mode: "sglang:d8t4p1+d8t1p4"
 seed: 42
-total_train_epochs: 1
-total_train_steps: 2
+total_train_epochs: 10
+total_train_steps: 1145
 weight_update_type: "astate"
 enable_colocate_mode: false
 async_training: false
@@ -15,7 +15,7 @@ train_dataset:
   path: "/storage/dataset/nlp/areal/moe_lite_math_0527_merge_train_areal.jsonl"
   shuffle: true
   max_length: 1024
-  batch_size: 8
+  batch_size: 64
   type: "rl"
 
 scheduler:
@@ -37,13 +37,13 @@ stats_logger:
     path: "/home/admin/logs/tfevent/asystem"
 
 gconfig:
-  n_samples: 1
+  n_samples: 8
   min_new_tokens: 0
   # NOTE!!
   # 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: 256
-  max_tokens: 1280
+  max_new_tokens: 15360
+  max_tokens: 16383
   greedy: false
   temperature: 1.0
   top_k: 1000000
@@ -111,7 +111,7 @@ rollout:
       env_vars:
         # if use ling max v2, need to specify USE_MAX_V2 = 1
         USE_MAX_V2: 1
-      image: /storage/openpsi/images/hybrid-engine-13680179-20250923154343.sif
+      image: /storage/openpsi/images/hybrid-engine-13680179-20251015181317.sif
 
 actor: &actor_ref
   experiment_name: ${experiment_name}
@@ -177,7 +177,7 @@ actor: &actor_ref
       distributed_backend: "nccl"
       distributed_timeout_minutes: 600
       enable_one_logger: false
-      expert_model_parallel_size: 1
+      expert_model_parallel_size: 8
       ffn_hidden_size: 5120
       first_k_dense_replace: 1
       global_batch_size: 512
@@ -293,7 +293,7 @@ actor: &actor_ref
         CUDA_LAUNCH_BLOCKING: 1
         # if use ling max v2, need to specify USE_MAX_V2 = 1
         USE_MAX_V2: 1
-      image: /storage/openpsi/images/hybrid-engine-13680179-20250923154343.sif
+      image: /storage/openpsi/images/hybrid-engine-13680179-20251015181317.sif
 
 ref:
   <<: *actor_ref
@@ -306,5 +306,5 @@ recover:
   latest_disable_save_hf: true
   periodic_disable_save_hf: false
   latest_save_interval: 1
-  periodic_save_interval: 20
+  periodic_save_interval: 2
   fileroot: "${storage_prefix}/experiments"

areal/examples/grpo_trainer.py

@@ -418,7 +418,7 @@ def submit_job(self, job: Job) -> dict[str, Any]:
             )
 
     def wait_for_jobs(
-        self, role: str, submitted_jobs: dict[str, str], timeout: float = 300.0
+        self, role: str, submitted_jobs: dict[str, str], timeout: float = 1200.0
     ) -> dict[str, Worker]:
         """
         等待作业启动并返回服务器信息
@@ -532,9 +532,6 @@ def _parse_ports_list(self, container_statuses: list) -> list:
         return ports_list
 
     def stop_job(self, job_uid: str):
-        # hack
-        return
-        """停止作业"""
         logger.info(f"Stopping job with UID: {job_uid}")
 
         try:

areal/extension/asystem/ascheduler/__init__.py

@@ -13,7 +13,6 @@
 from areal.scheduler.rpc.serialization import deserialize_value, serialize_value
 from areal.utils import logging
 from areal.utils.http import response_retryable
-from areal.extension.asystem.utils.async_utils import run_async_with_loop
 
 logger = logging.getLogger("RPCClient")
 

areal/extension/asystem/ascheduler/rpc_client.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import time
 from concurrent.futures import ThreadPoolExecutor
 
 from areal.api.alloc_mode import AllocationMode
@@ -59,7 +60,7 @@ async def _async_initialize(self, job: Job, *args, **kwargs):
 
         # Wait for workers to be ready
         self.logger.info("Waiting for workers to be ready...")
-        self.workers = self.scheduler.get_workers(role=job.role)
+        self.workers = self.scheduler.get_workers(role=job.role, timeout=1200)
         self.logger.info(f"Workers ready: {[w.id for w in self.workers]}")
 
         # Get engine class path for dynamic import on workers
@@ -92,6 +93,8 @@ async def _async_initialize(self, job: Job, *args, **kwargs):
             self.scheduler.async_call_engine(worker.id, "initialize", init_config)
             for worker, init_config in zip(self.workers, init_configs)
         ]
+        import time
+        time.sleep(60)
         await asyncio.gather(*tasks)
         self.logger.info("All engines are initialized...")
 
@@ -179,16 +182,16 @@ def _build_engine_initialize_config(
             main_server_addrs = [
                 f"{worker.ip}:{worker.engine_ports[0]}"
                 for worker in self.workers[
-                    index : index + self.alloc_mode.gen_instance_size
-                ]
+                              index: index + self.alloc_mode.gen_instance_size
+                              ]
                 if worker.engine_ports
             ]
             free_addrs = [
                 [
                     f"{worker.ip}:{port}"
                     for worker in self.workers[
-                        index : index + self.alloc_mode.gen_instance_size
-                    ]
+                                  index: index + self.alloc_mode.gen_instance_size
+                                  ]
                     for port in worker.engine_ports[1:]
                 ]
             ]
@@ -207,7 +210,74 @@ def _build_engine_initialize_config(
 
         return init_configs
 
-    async def update_weights(self, meta: WeightUpdateMeta) -> None:
+    def update_weights(self, meta: WeightUpdateMeta) -> None:
+        """Update weights - thread-safe for ThreadPoolExecutor calls."""
         self.logger.info("begin update_weights")
-        execute_parallel_tasks(self.workers, self.scheduler, "update_weights", meta)
+        self._execute_async_task_on_workers("update_weights", meta)
         self.logger.info("finish update_weights")
+
+    def set_version(self, version: int) -> None:
+        self._version = version
+        self.logger.info("begin set_version")
+        self._execute_async_task_on_workers("set_version", version)
+        self.logger.info("finish set_version")
+
+    def notify_event(self, event: str, global_step: int) -> None:
+        """Notify workers about training start/end events.
+
+        Args:
+            event: "train_start" or "train_end"
+            global_step: Current global step
+        """
+        self.logger.info(f"begin notify_event global_step: {global_step}")
+        self._execute_async_task_on_workers("notify_event", event, global_step)
+        self.logger.info(f"finished notify_event global_step: {global_step}")
+        return None
+
+    def _execute_async_task_on_workers(self, method_name: str, *args, **kwargs):
+        def _run_async_in_thread():
+            """Run async code in a thread-safe manner."""
+            # Always create a new event loop for this thread to avoid conflicts
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            try:
+                async def _async_exec_func():
+                    try:
+                        self.logger.info(f"Executing {method_name} on {len(self.workers)} workers")
+                        tasks = [
+                            self.scheduler.async_call_engine(
+                                worker.id, method_name, *args, **kwargs, _should_bcast=False
+                            )
+                            for worker in self.workers
+                        ]
+                        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+                        # Check for exceptions in results
+                        for i, result in enumerate(results):
+                            if isinstance(result, Exception):
+                                self.logger.error(
+                                    f"Worker {self.workers[i].id} failed to execute {method_name}: {result}")
+                            else:
+                                self.logger.info(f"Worker {self.workers[i].id} successfully executed {method_name}")
+
+                        # Re-raise if any exceptions occurred
+                        for result in results:
+                            if isinstance(result, Exception):
+                                raise result
+
+                        return results
+                    except Exception as e:
+                        self.logger.error(f"Failed to execute {method_name} on workers: {e}")
+                        raise e
+
+                return loop.run_until_complete(_async_exec_func())
+            finally:
+                # Always close the loop we created
+                if not loop.is_closed():
+                    loop.close()
+                # Clear the event loop for this thread
+                asyncio.set_event_loop(None)
+
+        return _run_async_in_thread()
+