[lite] [fix] Fix a performance issue and several minor issues before release

arealite/api/cli_args.py

@@ -195,7 +195,10 @@ class TrainEngineConfig:
     gradient_checkpointing: bool = field(
         default=True, metadata={"help": "Enable gradient checkpointing"}
     )
-    dtype: str = field(default="float16", metadata={"help": "Parameter dtype."})
+    dtype: str = field(default="bfloat16", metadata={"help": "Parameter dtype."})
+    grad_reduce_dtype: str = field(
+        default="float32", metadata={"help": "Gradient reduce dtype."}
+    )
     optimizer: Optional[OptimizerConfig] = field(
         default=None, metadata={"help": "Optimizer configuration"}
     )
@@ -262,7 +265,7 @@ class PPOActorConfig(TrainEngineConfig):
         default=1.0, metadata={"help": "Lambda parameter for GAE"}
     )
     adv_norm: bool = field(
-        default=True, metadata={"help": "Enable advantage normalization"}
+        default=True, metadata={"help": "Enable advantage normalization globally"}
     )
 
     # KL Control
@@ -275,7 +278,9 @@ class PPOActorConfig(TrainEngineConfig):
     )
     use_decoupled_loss: bool = field(
         default=False,
-        metadata={"help": "Use the decoupled loss. recompute_logprob must be True."},
+        metadata={
+            "help": "Use the decoupled loss. Implicitly enable recompute_logprob."
+        },
     )
     behav_imp_weight_cap: Optional[float] = field(
         default=None,
@@ -329,7 +334,7 @@ class SGLangConfig:
     schedule_policy: str = "lpm"
     schedule_conservativeness: float = 1.0
     cpu_offload_gb: int = 0
-    dtype: str = "float16"
+    dtype: str = "bfloat16"
     kv_cache_dtype: str = "auto"
     # logging
     log_level: str = "warning"
@@ -407,31 +412,8 @@ def build_args(
             dist_init_addr=dist_init_addr,
             **args,
         )
-        sglang_version = pkg_version.get_version("sglang")
-        if sglang_version:
-            version_less_than_0_4_4 = (
-                pkg_version.compare_versions(sglang_version, "0.4.4") < 0
-            )
-            version_less_than_0_4_3 = (
-                pkg_version.compare_versions(sglang_version, "0.4.3") < 0
-            )
-        elif pkg_version.is_available("sglang"):
-            version_less_than_0_4_4 = pkg_version.is_version_less("sglang", "0.4.4")
-            version_less_than_0_4_3 = pkg_version.is_version_less("sglang", "0.4.3")
-        else:
-            raise ValueError(
-                "A installed SGLang package or a specific SGLang version should be provided to build SGLang server cmd."
-            )
-        if version_less_than_0_4_4:
-            args.pop("log_requests_level")
-        if version_less_than_0_4_3:
-            args.pop("enable_nccl_nvls")
-            args.pop("triton_attention_num_kv_splits")
-            args.pop("cuda_graph_bs")
-            args.pop("enable_memory_saver")
-            args.pop("allow_auto_truncate")
-            args.pop("file_storage_path")
-
+        if not pkg_version.is_version_greater_or_equal("sglang", "0.4.9.post2"):
+            raise RuntimeError("Needs sglang>=0.4.9.post2 to run the code.")
         return args
 
 
@@ -466,7 +448,7 @@ class InferenceEngineConfig:
         default="round_robin",
         metadata={"help": "Request scheduling policy", "choices": ["round_robin"]},
     )
-    setup_timeout: float = field(default=90.0)
+    setup_timeout: float = field(default=120.0)
     request_timeout: float = field(
         default=3600, metadata={"help": "Timeout for HTTP requests."}
     )

arealite/api/engine_api.py

@@ -174,6 +174,7 @@ def prepare_batch(
         self,
         dataloader: StatefulDataLoader,
         workflow: "RolloutWorkflow",
+        should_accept: Callable | None = None,
     ):
         """Asynchronously submit and wait until a full batch is ready."""
         raise NotImplementedError()

arealite/api/workflow_api.py

@@ -1,10 +1,30 @@
-from typing import TYPE_CHECKING, Any, Dict
+import asyncio
+import itertools
+import queue
+import threading
+import time
+import traceback
+from typing import TYPE_CHECKING, Any, Callable, Dict, List
 
+import torch.distributed as dist
+import uvloop
 from tensordict import TensorDict
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from arealite.api.cli_args import InferenceEngineConfig
+from arealite.api.engine_api import InferenceEngine
+from arealite.api.io_struct import RolloutStat
+from arealite.utils.data import concat_padded_tensors
+from realhf.base import logging
 
 if TYPE_CHECKING:
     from arealite.api.engine_api import InferenceEngine
 
+logger = logging.getLogger("arealite.workflow_api")
+
+
+ROLLOUT_POLL_WAIT_TIME = 0.05
+
 
 class RolloutWorkflow:
 
@@ -16,3 +36,226 @@ async def arun_episode(
         See concrete example implementations under the `arealite/workflow` directory.
         """
         raise NotImplementedError()
+
+
+class WorkflowExecutor:
+
+    def __init__(
+        self,
+        config: InferenceEngineConfig,
+        inference_engine: "InferenceEngine",
+    ):
+        config.max_concurrent_rollouts = (
+            config.max_concurrent_rollouts or config.consumer_batch_size
+        )
+        self.config = config
+        self.exiting = threading.Event()
+        self.paused = threading.Event()
+        self.lock = threading.Lock()
+
+        self.inference_engine = inference_engine
+
+        qsize = config.queue_size or config.max_concurrent_rollouts * 16
+        self.input_queue = queue.Queue(maxsize=qsize)
+        self.output_queue = queue.Queue(maxsize=qsize)
+        self.result_cache: List[TensorDict] = []
+
+        self.rollout_stat = RolloutStat()
+
+    def initialize(self):
+        self.rollout_tasks: Dict[str, asyncio.Task] = {}
+        self.rollout_thread = threading.Thread(target=self._rollout_thread)
+        self.rollout_thread.start()
+
+    def destroy(self):
+        self.exiting.set()
+        self.rollout_thread.join()
+
+    def get_capacity(self):
+        if dist.is_initialized():
+            world_size = dist.get_world_size()
+        else:
+            world_size = 1
+
+        with self.lock:
+            max_concurrent_rollouts = max(
+                1, self.config.max_concurrent_rollouts // world_size
+            )
+            capacity = max_concurrent_rollouts - len(self.rollout_tasks)
+            # Staleness control
+            version = self.inference_engine.get_version()
+            ofp = self.config.max_head_offpolicyness
+            sample_cnt = self.rollout_stat.accepted + self.rollout_stat.running
+            consumer_bs = max(1, self.config.consumer_batch_size // world_size)
+            capacity = min(capacity, (ofp + version + 1) * consumer_bs - sample_cnt)
+        return capacity
+
+    def _rollout_thread(self):
+        """Thread that runs the rollout loop."""
+        try:
+            uvloop.run(self._rollout_thread_async())
+        except Exception:
+            traceback.print_exc()
+
+    async def _rollout_thread_async(self):
+        rollout_tasks = self.rollout_tasks
+        rid = 0
+        try:
+            while not self.exiting.is_set():
+                # Check capacity
+                capacity = self.get_capacity()
+                # Create new rollout task
+                self.lock.acquire()
+                while (
+                    capacity > 0
+                    and not self.paused.is_set()
+                    and self.input_queue.qsize() > 0
+                ):
+                    data, workflow = self.input_queue.get_nowait()
+                    logger.debug(f"Get data from puller: {data}")
+                    task = asyncio.create_task(
+                        workflow.arun_episode(self.inference_engine, data),
+                        name=str(rid),
+                    )
+                    rollout_tasks[str(rid)] = task
+                    self.rollout_stat.submitted += 1
+                    self.rollout_stat.running += 1
+                    if self.config.enable_rollout_tracing:
+                        logger.info(
+                            f"Submit rollout rid {rid}. "
+                            f"Submit: {self.rollout_stat.submitted}, "
+                            f"running: {self.rollout_stat.running}, "
+                            f"accepted: {self.rollout_stat.accepted}."
+                        )
+                    capacity -= 1
+                    rid += 1
+                tasks = list(rollout_tasks.values())
+                self.lock.release()
+
+                # Wait for rollout completion
+                done = []
+                if tasks:
+                    done, _ = await asyncio.wait(
+                        tasks,
+                        timeout=ROLLOUT_POLL_WAIT_TIME,
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                # Collect done results
+                for task in done:
+                    traj = await task
+                    traj: TensorDict
+                    task_rid = task.get_name()
+                    with self.lock:
+                        rollout_tasks.pop(task_rid)
+                        self.rollout_stat.accepted += 1
+
+                        self.rollout_stat.running -= 1
+                        if self.config.enable_rollout_tracing:
+                            logger.info(
+                                f"Finish rollout {task_rid}. "
+                                f"Submit: {self.rollout_stat.submitted}, "
+                                f"running: {self.rollout_stat.running}, "
+                                f"accepted: {self.rollout_stat.accepted}."
+                            )
+                    try:
+                        self.output_queue.put_nowait(traj)
+                    except queue.Full:
+                        raise RuntimeError(
+                            "Output queue full. Please increase queue_size."
+                        )
+
+                await asyncio.sleep(1)
+        except Exception:
+            traceback.print_exc()
+        finally:
+            # Cancel remaining tasks
+            with self.lock:
+                for task in rollout_tasks.values():
+                    if not task.done():
+                        task.cancel()
+                        try:
+                            await task
+                        except asyncio.CancelledError:
+                            pass
+
+    def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None:
+        try:
+            self.input_queue.put_nowait((data, workflow))
+        except queue.Full:
+            raise RuntimeError("Input queue full. Please increase queue_size.")
+
+    def wait(
+        self,
+        count: int,
+        timeout: float | None = None,
+        should_accept: Callable | None = None,
+    ) -> TensorDict:
+        tik = time.perf_counter()
+        accepted = len(self.result_cache)
+        timeout = timeout or float(7 * 24 * 3600)
+        while (
+            accepted < count
+            and not self.exiting.is_set()
+            and time.perf_counter() - tik < timeout
+        ):
+            try:
+                result = self.output_queue.get(timeout=ROLLOUT_POLL_WAIT_TIME)
+                if should_accept is None or should_accept(result):
+                    self.result_cache.append(result)
+                    accepted += 1
+                else:
+                    with self.lock:
+                        self.rollout_stat.accepted -= 1
+            except queue.Empty:
+                pass
+        if self.exiting.is_set():
+            raise RuntimeError("Rollout engine is exiting, cannot wait for results.")
+        if accepted < count:
+            raise TimeoutError(
+                f"Timed out waiting for {count} rollouts, " f"only received {accepted}."
+            )
+        results, self.result_cache = (
+            self.result_cache[:count],
+            self.result_cache[count:],
+        )
+        return concat_padded_tensors(results)
+
+    def rollout_batch(
+        self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow"
+    ) -> TensorDict:
+        """Submit a batch of requests to the inference engine and wait for the results."""
+        for item in data:
+            self.submit(item, workflow)
+        return self.wait(count=len(data))
+
+    def prepare_batch(
+        self,
+        dataloader: StatefulDataLoader,
+        workflow: "RolloutWorkflow",
+        should_accept: Callable | None = None,
+    ):
+        if not hasattr(self, "data_generator"):
+            self.data_generator = itertools.cycle(dataloader)
+        assert dataloader.batch_size is not None
+        while True:
+            # Submit at least two batches to allow maximum overlap
+            if (
+                self.get_capacity() + dataloader.batch_size > 0
+                and self.input_queue.qsize() + dataloader.batch_size
+                < self.input_queue.maxsize
+            ):
+                data = next(self.data_generator)
+                for item in data:
+                    self.submit(item, workflow=workflow)
+            try:
+                return self.wait(
+                    dataloader.batch_size, timeout=1, should_accept=should_accept
+                )
+            except TimeoutError:
+                pass
+
+    def pause(self):
+        self.paused.set()
+
+    def resume(self):
+        self.paused.clear()

arealite/engine/fsdp_engine.py

@@ -55,7 +55,7 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None):
         # Simple auto wrap policy
         self.mixed_precision_policy = MixedPrecisionPolicy(
             param_dtype=getattr(torch, self.config.dtype),
-            reduce_dtype=torch.float32,
+            reduce_dtype=getattr(torch, self.config.grad_reduce_dtype),
             cast_forward_inputs=True,
         )
         self.device_mesh = create_fsdp_device_mesh(self.world_size, self.world_size)

arealite/engine/ppo/actor.py

@@ -128,7 +128,7 @@ def compute_advantages(self, data: TensorDict) -> None:
         advantages = torch.stack(advantages_reversed[::-1], dim=1)
 
         # Optionally perform advantage normalization.
-        if self.adv_norm:
+        if self.adv_norm or self.group_adv_norm:
             if self.group_adv_norm:
                 adv_list = []
                 for i in range(0, bs, self.group_size):