PullRequest: 958 trainer debug

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

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


* trainer debug

Author: 楚财

areal/controller/rollout_controller.py

@@ -403,7 +403,6 @@ def wait(self, count: int, timeout: float | None = None) -> DistributedBatch:
             capacity = self.get_capacity()
             # Submit pending tasks
             self.logger.info(f"Capacity: {capacity}, pending inputs: {len(self._pending_inputs)}")
-
             for _ in range(capacity):
                 if len(self._pending_inputs) == 0:
                     break
@@ -484,6 +483,7 @@ def rollout_batch(
             A concatenated batch of trajectory results
         """
         # Submit all requests
+
         for item in data:
             self.submit(
                 item,

areal/examples/configs/my001/on_policy.yaml

@@ -1,11 +1,12 @@
 experiment_name: mini-model
-trial_name: on-policy
-allocation_mode: "sglang:d8t4p1+d8t1p4"
+trial_name: test
+allocation_mode: "sglang:d2t4p1+d2t1p4"
 seed: 42
-total_train_epochs: 10
-total_train_steps: 1145
+total_train_epochs: 1
+total_train_steps: 2
 weight_update_type: "astate"
-enable_colocate_mode: true
+enable_colocate_mode: false
+async_training: false
 
 storage_prefix: "/storage/openpsi"
 
@@ -14,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: 64
+  batch_size: 8
   type: "rl"
 
 scheduler:
@@ -36,13 +37,13 @@ stats_logger:
     path: "/home/admin/logs/tfevent/asystem"
 
 gconfig:
-  n_samples: 8
+  n_samples: 1
   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: 15360
-  max_tokens: 16383
+  max_new_tokens: 256
+  max_tokens: 1280
   greedy: false
   temperature: 1.0
   top_k: 1000000
@@ -176,7 +177,7 @@ actor: &actor_ref
       distributed_backend: "nccl"
       distributed_timeout_minutes: 600
       enable_one_logger: false
-      expert_model_parallel_size: 8
+      expert_model_parallel_size: 1
       ffn_hidden_size: 5120
       first_k_dense_replace: 1
       global_batch_size: 512

areal/examples/grpo_trainer.py

@@ -13,6 +13,7 @@
 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")
 
@@ -172,9 +173,18 @@ def call_engine(self, worker_id, method, max_retries=3, *args, **kwargs):
             WorkerFailedError: If worker process has failed
             EngineCallError: If method call fails
         """
-        return asyncio.run(
-            self.async_call_engine(worker_id, method, max_retries, *args, **kwargs)
-        )
+        # 创建新的事件循环并运行异步任务
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            return loop.run_until_complete(
+                self.async_call_engine(worker_id, method, max_retries, *args, **kwargs)
+            )
+        finally:
+            try:
+                loop.close()
+            except Exception:
+                pass
 
     async def async_call_engine(
         self, worker_id, method, max_retries=3, *args, **kwargs
@@ -241,7 +251,7 @@ async def async_call_engine_with_serialized_data(
         for attempt in range(1, max_retries + 1):
             try:
                 logger.info(
-                    f"Async calling method '{method}' on worker '{worker_id}' (attempt {attempt})"
+                    f"Async calling method '{method}' on worker '{worker_id}' (attempt {attempt}), url: {url}"
                 )
 
                 response = await self._async_http_client.post(

areal/extension/asystem/ascheduler/rpc_client.py

@@ -8,6 +8,7 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from areal.api.alloc_mode import AllocationMode
+from areal.api.io_struct import WeightUpdateMeta
 from areal.api.cli_args import InferenceEngineConfig
 from areal.api.engine_api import InferenceEngine
 from areal.api.scheduler_api import Job, Scheduler
@@ -20,6 +21,7 @@
 from areal.extension.asystem.remote_hybrid_inference_worker import (
     RemoteHypidInferenceInitConfig,
 )
+from areal.extension.asystem.controller.util import execute_parallel_tasks
 from areal.utils import logging
 
 
@@ -204,3 +206,8 @@ def _build_engine_initialize_config(
             init_configs.append(init_config)
 
         return init_configs
+
+    async def update_weights(self, meta: WeightUpdateMeta) -> None:
+        self.logger.info("begin update_weights")
+        execute_parallel_tasks(self.workers, self.scheduler, "update_weights", meta)
+        self.logger.info("finish update_weights")

areal/extension/asystem/controller/rollout_controller.py

@@ -12,9 +12,10 @@
 from typing import Any
 from areal.extension.asystem.api.cli_args import TrainEngineConfig
 from areal.api.engine_api import TrainEngine
-from areal.api.io_struct import AllocationMode, FinetuneSpec
+from areal.api.io_struct import FinetuneSpec
 from areal.api.scheduler_api import Job, Scheduler
 from areal.controller.train_controller import TrainController as BaseTrainController
+from areal.extension.asystem.controller.util import execute_parallel_tasks, calc_metrics
 from areal.extension.asystem.remote_hybrid_train_worker import RemoteMegatronInitConfig
 from areal.utils import logging, stats_tracker
 from areal.controller.batch import DistributedBatch
@@ -23,55 +24,6 @@
 logger = logging.getLogger("TrainController")
 
 
-def _execute_parallel_tasks(workers, scheduler, method_name, *args):
-    """Execute tasks in parallel across all workers.
-    
-    This is a helper function to reduce code duplication when executing
-    the same method on all workers with identical parameters.
-    
-    Parameters
-    ----------
-    workers : list
-        List of worker objects
-    scheduler : Scheduler
-        Scheduler instance for async calls
-    method_name : str
-        Name of the method to call on each worker's engine
-    *args, **kwargs
-        Arguments to pass to the method
-        
-    Returns
-    -------
-    list
-        Results from all workers
-        
-    Raises
-    ------
-    RuntimeError
-        If any worker fails to execute the task
-    """
-    tasks = [
-        scheduler.async_call_engine(
-            worker.id, method_name, *args, _should_bcast=False
-        )
-        for worker in workers
-    ]
-    
-    try:
-        return asyncio.run(asyncio.gather(*tasks, return_exceptions=False))
-    except KeyboardInterrupt:
-        raise
-    except Exception as e:
-        raise RuntimeError(f"{method_name} failed, error: {e}")
-
-
-def _calc_metrics(batch_inputs):
-    # seqlen std
-    seqlens = [td["seqlen"].sum().item() for td in batch_inputs]
-    seqlen_std = torch.tensor(seqlens).float().std().item()
-    stats_tracker.scalar(**{"seqlen_std": seqlen_std})
-
-
 class TrainController(BaseTrainController):
     """ASystem-specific TrainController.
 
@@ -218,7 +170,7 @@ def train_batch(
         with (stats_tracker.record_timing("train_batch_data_split"), ):
             batches = input_.chunk_by_ffd(self.group_size, self.dp_size)
 
-        _calc_metrics(batches)
+        calc_metrics(batches)
 
         tasks = [
             self.scheduler.async_call_engine(
@@ -286,15 +238,17 @@ def compute_logp(self, input_: DistributedBatch) -> Tensor:
 
     def upload_weights(self, meta: WeightUpdateMeta):
         """Upload weights to the inference engine."""
-        _execute_parallel_tasks(self.workers, self.scheduler, "upload_weights", meta)
+        self.logger.info("begin upload_weights")
+        execute_parallel_tasks(self.workers, self.scheduler, "upload_weights", meta)
+        self.logger.info("finished upload_weights")
 
     def save(self, meta: SaveLoadMeta):
         """Save model weights (and optimizer states) for later use."""
-        _execute_parallel_tasks(self.workers, self.scheduler, "save", meta)
+        execute_parallel_tasks(self.workers, self.scheduler, "save", meta)
 
     def load(self, meta: SaveLoadMeta):
         """Load model weights and optimizer states from a file."""
-        _execute_parallel_tasks(self.workers, self.scheduler, "load", meta)
+        execute_parallel_tasks(self.workers, self.scheduler, "load", meta)
 
     def notify_event(self, event: str, global_step: int) -> None:
         """Notify workers about training start/end events.
@@ -303,5 +257,5 @@ def notify_event(self, event: str, global_step: int) -> None:
             event: "train_start" or "train_end"
             global_step: Current global step
         """
-        _execute_parallel_tasks(self.workers, self.scheduler, "notify_event", event, global_step)
+        execute_parallel_tasks(self.workers, self.scheduler, "notify_event", event, global_step)
         return None

areal/extension/asystem/controller/train_controller.py

@@ -0,0 +1,71 @@
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+import torch
+
+from areal.utils import stats_tracker
+from areal.extension.asystem.utils.async_utils import run_async_with_loop
+
+
+def execute_parallel_tasks(workers, scheduler, method_name, *args):
+    """Execute tasks in parallel across all workers.
+
+    This is a helper function to reduce code duplication when executing
+    the same method on all workers with identical parameters.
+
+    Parameters
+    ----------
+    workers : list
+        List of worker objects
+    scheduler : Scheduler
+        Scheduler instance for async calls
+    method_name : str
+        Name of the method to call on each worker's engine
+    *args, **kwargs
+        Arguments to pass to the method
+
+    Returns
+    -------
+    list
+        Results from all workers
+
+    Raises
+    ------
+    RuntimeError
+        If any worker fails to execute the task
+    """
+    logger.info(f"[DEBUG] execute_parallel_tasks called with method: {method_name}, workers: {[w.id for w in workers]}")
+    tasks = [
+        scheduler.async_call_engine(
+            worker.id, method_name, *args, _should_bcast=False
+        )
+        for worker in workers
+    ]
+
+    try:
+        logger.info(f"[DEBUG] Created {len(tasks)} async tasks")
+        # 创建新的事件循环并运行所有任务
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            logger.info(f"[DEBUG] Starting async execution")
+            result = loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=False))
+            logger.info(f"[DEBUG] Async execution completed successfully")
+            return result
+        finally:
+            try:
+                loop.close()
+            except Exception:
+                pass
+    except KeyboardInterrupt:
+        raise
+    except Exception as e:
+        logger.error(f"[DEBUG] execute_parallel_tasks failed: {str(e)}")
+        raise RuntimeError(f"{method_name} failed, error: {e}")
+
+
+def calc_metrics(batch_inputs):
+    # seqlen std
+    seqlens = [td["seqlen"].sum().item() for td in batch_inputs]
+    seqlen_std = torch.tensor(seqlens).float().std().item()
+    stats_tracker.scalar(**{"seqlen_std": seqlen_std})

areal/extension/asystem/controller/util.py

@@ -232,7 +232,7 @@ async def _rollout_thread_async(self):
                 ):
                     data, workflow = self.input_queue.get_nowait()
 
-                    # logger.info(f"Get data from puller: {data}")
+                    logger.info(f"_rollout_thread_async before arun_episode data: {data}")
                     task = asyncio.create_task(
                         (
                             workflow.arun_episode(self, data)
@@ -427,6 +427,8 @@ def get_capacity(self):
         return capacity
 
     def update_weights(self, meta):
+        logger.info(f"[DEBUG] update_weights called with meta: {meta}")
+        logger.info(f"[DEBUG] Available addresses: {self.addresses}")
         self._update_weights(meta)
         return True
 
@@ -460,7 +462,7 @@ def update_single_server(addr):
                 wait_future_ordered(futures)
 
             logger.info(
-                f"Loading weights done in {(time.time_ns() - load_timestamp) / 1e6:.2f} ms, updated version: {meta.model_version}"
+                f"Loading weights done in {(time.time_ns() - load_timestamp) / 1e6:.2f} ms"
             )
         elif meta.type == "nccl" or meta.type == "astate":
             load_timestamp = time.time_ns()
@@ -498,15 +500,15 @@ def update_single_server(addr):
                 wait_future_ordered(futures)
 
             logger.info(
-                f"Loading weights done in {(time.time_ns() - load_timestamp) / 1e6:.2f} ms, updated version: {meta.model_version}"
+                f"Loading weights done in {(time.time_ns() - load_timestamp) / 1e6:.2f} ms"
             )
         else:
             raise FrameworkError(
                 "FrameworkError",
                 "InferenceWorkerError",
                 f"Unknown weight update type {meta.type}",
             )
-        self.set_version(meta.model_version)
+        # self.set_version(meta.model_version)
 
     def pause(self):
         self.paused.set()

areal/extension/asystem/remote_hybrid_inference_worker.py

@@ -46,6 +46,7 @@ def __init__(self, config: TrainEngineConfig):
         self.megatron_addr = None
         self.global_step = self.config.global_step
         self.global_rank = 0
+        self._version: int = 0
 
         # initialization
         self.initialized = False
@@ -930,6 +931,11 @@ def _compute_logprobs(
 
         return None
 
+    def set_version(self, version: int):
+        self._version = version
+
+    def get_version(self) -> int:
+        return self._version
 
 def serialize_and_compress(data):
     serialized_data = cloudpickle.dumps(data)