[Docs] [Lite] AReaLite release documentation

docs/arealite/gsm8k_grpo.md

@@ -29,8 +29,8 @@ show you how these steps are done in details.
 
 ## Launching the Experiment
 
-As shown in [Quickstart Guide](../tutorial/quickstart.md), experiments in AReaLite are
-launched using standalone launchers with the following commands:
+As shown in the [quickstart guide](../tutorial/quickstart.md), experiments in AReaLite
+are launched using standalone launchers with the following commands:
 
 ```
 # Local Launcher
@@ -72,7 +72,7 @@ config: GRPOConfig
 ## Loading and Preprocessing Dataset
 
 We use the `datasets` and `torchdata` packages to load and preprocess the dataset into
-our dataloader. First, we download `openai/gsm8k` from Huggingface and split it by data
+our dataloader. First, we download `openai/gsm8k` from Hugging Face and split it by data
 parallel ranks, then map it to our desired format:
 
 ```python
@@ -109,12 +109,100 @@ details.
 
 ## Rollout
 
-The data lifecycle is controlled by an `RLVRWorkflow`, which defines how data progresses
-from prompts to complete rollout data containing all fields required for training. Our
-example shows a single-turn RLVR workflow with a math reward function. The core logic of
-the workflow is implemented in an async method `arun_episode`, which takes a prompt,
-generate answers with `RemoteSGLangEngine`, computes rewards, and populates additional
-fields to produce finalized training data.
+### Inference Engine: `RemoteSGLangEngine`
+
+In AReaLite, generation tasks are offloaded to remote inference servers, which operate
+on separate GPUs from those used for training. The `RemoteSGLangEngine` acts as a client
+that interacts with the servers. `RemoteSGLangEngine` runs in a SPMD manner on every
+training process, without occupying any GPUs.
+
+`RemoteSGLangEngine` provides two APIs, `agenerate` and `async_update_weights`. It is
+worth mentioning that, in asynchronous RL experiment in AReaLite, inference-side weight
+update could happen **in the middle of** generation of one prompt. With that being said,
+one output sequence could be generated by multiple versions of models. Let us glimpse
+into code of `agenerate` and `async_update_weights` for a better understanding.
+
+In `async_update_weights`, the engine first send `pause_generation` requests to all
+inference servers, notifying them a weight update is about to happen. Upon receiveing
+`pause_generation`, inference servers will immediately stop generating and respond with
+already generated tokens. Then, the engine sends `update_weights_from_distributed` (for
+NCCL update) or `update_weights_from_disk` (for disk update). After the update is
+finished, the engine sends `continue_generation` to inference server telling them to
+start working again.
+
+```python
+class RemoteSGLangEngine:
+    ...
+    def async_update_weights(self, meta: WeightUpdateMeta):
+        # `async_update_weights` is completely async.
+        # It submits task to a ProcessPoolExecutor and returns a future
+        for addr in self.addresses:
+            res = requests.post(f"http://{addr}/pause_generation")
+        if meta.type == "nccl":
+            future = self.executor.submit(
+                # a function that send `update_weights_from_distributed` request
+                update_weights_from_distributed,
+            )
+        elif meta.type == "disk":
+            ...
+
+        def callback(future):
+            for addr in self.addresses
+                requests.post(f"http://{addr}/continue_generation")
+
+        future.add_done_callback(callback)
+        return future
+```
+
+`agenerate` takes an `LLMRequest` with `input_ids` of **a single prompt** and generation
+hyperparameters, and returns the final generation result, an `LLMResponse` with
+`output_tokens` and other outputs. Since the generation could be interrupted,
+`agenerate` iteratively prepares payload, sends requests and receives responses until
+the generation finishes.
+
+```python
+class RemoteSGLangEngine:
+    ...
+    async def agenerate(self, req: LLMRequest):
+        payload = ... # prepare payload for request
+        # If request is from the same workflow, choose old server
+        # to allow KVCache reuse. Otherwise choose server in a round
+        # robin manner.
+        server_addr = self.choose_server(req)
+        stop_reason = None
+        # other outputs are omitted for simplicity
+        output_tokens = []
+        while (stop_reason != "stop" and len(output_tokens) < max_new_tokens):
+            # Request is interrupted, wait to avoid contention
+            if stop_reason is not None:
+                await asyncio.sleep(0.5)
+            # send request to remote sever
+            result = await arequest_with_retry(
+                addr=server_addr,
+                endpoint="/generate",
+                payload=payload,
+                method="POST"
+            )
+            output_tokens.extend(result["output_ids"])
+            # prepare payload for the next request
+            payload["input_ids"] += results["output_ids"]
+            payload["sample_params"]["max_new_tokens"] -= len(results["output_ids"])
+        return LLMResponse(
+            input_tokens=req.input_ids,
+            output_tokens=output_tokens,
+            ...
+        )
+
+```
+
+### `RLVRWorkflow` and `WorkflowExecutor`
+
+The rollout data lifecycle is controlled by an `RLVRWorkflow`, which defines how data
+progresses from prompts to complete rollout data containing all fields required for
+training. Our example shows a single-turn RLVR workflow with a math reward function. The
+core logic of the workflow is implemented in an async method `arun_episode`, which takes
+a prompt, generate answers with `RemoteSGLangEngine`, computes rewards, and populates
+additional fields to produce finalized training data.
 
 ```python
 class RLVRWorkflow(RolloutWorkflow):
@@ -158,26 +246,22 @@ workflow = RLVRWorkflow(
 )
 ```
 
-In AReaLite, generation tasks are offloaded to remote inference servers, which operate
-on separate GPUs from those used for training. The `RemoteSGLangEngine` acts as a client
-that interacts with the servers. `RemoteSGLangEngine` runs in a SPMD manner on every
-training process, without occupying any GPUs.
-
-`RemoteSGLangEngine` is responsible for managing the data streaming through rollout
+`WorkflowExecutor` is responsible for managing the data streaming through rollout
 workflows, and collates completed rollout data into batched training samples. When
 initializing, it launches a rollout thread that runs rollout workflows as `asyncio`
 tasks. The following code shows the simplified version of rollout thread implementation,
 which iteratively:
 
 - Checks available capacity. The capacity controls current number of rollout workflows
-  to limit concurrency and data off-policyness.
+  to limit concurrency and **data off-policyness** (The difference between the model
+  version used by generation and the model version updated by the trainer).
 - If there is capacity left and rollout is not paused for weight update, continuously
   obtains data from `input_queue` and creates `asyncio` tasks to run the workflows.
 - Waits for rollout workflows to finish.
 - Gathers data from finished workflows and puts them into `output_queue`
 
 ```python
-class RemoteSGLangEngine(InferenceEngine):
+class WorkflowExecutor:
     ...
     async def _rollout_thread_async(self):
         rid = 0
@@ -202,15 +286,15 @@ class RemoteSGLangEngine(InferenceEngine):
                     rid += 1
                 # Wait for rollout completion
                 tasks = list(rollout_tasks.values())
-                done = []
+                completed_tasks = []
                 if tasks:
-                    done, _ = await asyncio.wait(
+                    completed_tasks, _ = await asyncio.wait(
                         tasks,
                         timeout=ROLLOUT_POLL_WAIT_TIME,
                         return_when=asyncio.FIRST_COMPLETED,
                     )
                 # Collect done results, put the results into output queue
-                for task in done:
+                for task in completed_tasks:
                     traj = await task
                     task_rid = task.get_name()
                     rollout_tasks.pop(task_rid)
@@ -257,10 +341,11 @@ def prepare_batch(
             pass
 ```
 
-The usage of `RemoteSGLangEngine` in the training script is simple:
+The usage of `WorkflowExecutor` in the training script is simple:
 
 ```python
-rollout = RemoteSGLangEngine(config.rollout)
+inf_engine = RemoteSGLangEngine(config.inf_engine)
+rollout = WorkflowExecutor(config.rollout, inf_engine)
 rollout.initialize()
 eval_rollout = ...
 
@@ -348,25 +433,24 @@ weight_update_meta = WeightUpdateMeta.from_disk(config.saver)
 ```
 
 After a training step is finished, we transfer new weights from actor engine to remote
-inference servers with steps shown in the following code:
+inference servers:
+
+1. The rollout engine needs to stop sending generation requests to remote servers
+   (`rollout.pause()`) before weight update to avoid server-side congestion.
+1. Since we need to invoke weight update on the trainer engine and remote inference
+   servers at the same time, in the training script, we asynchronously send requests to
+   remote inference servers, and then immediately upload weights on the trainer engine.
 
 ```python
-# 1. Pause rollout on remote inference servers
 rollout.pause()
-# 2. Send requests to remote servers, tell them to update weights
 if dist.get_rank() == 0:
-    future = rollout.update_weights(weight_update_meta)
-# 3. Actor begins to transfer weights
+    future = rollout.async_update_weights(weight_update_meta)
 actor.upload_weights(weight_update_meta)
-# 4. Wait for remote servers to return after finishing updates
 if dist.get_rank() == 0:
     future.result()
-# 5. Synchronize rollout processes for model version update
 dist.barrier(device_ids=[actor.device.index])
 torch.cuda.synchronize()
-# 6. Resume rollout on remote inference servers
 rollout.resume()
-# 7. Set version, ensures versions on actor and rollout engine are identical
 actor.set_version(global_step + 1)
 rollout.set_version(global_step + 1)
 ```
@@ -398,7 +482,7 @@ for global_step in range(max_steps):
 
     rollout.pause()
     if dist.get_rank() == 0:
-        future = rollout.update_weights(weight_update_meta)
+        future = rollout.async_update_weights(weight_update_meta)
     actor.upload_weights(weight_update_meta)
     if dist.get_rank() == 0:
         future.result()

docs/customization/agent.md

@@ -8,10 +8,12 @@ You can find the complete implementation in `arealite/workflow/multi_turn.py`.
 
 ## Step 1: Define Your Workflow
 
-AReaLite gives you flexibility in how you design your agents. Instead of rigid `Agent`
-classes that might constrain your agent's capabilities, AReaLite captures all rollout
-behavior in a `RolloutWorkflow` class. This approach lets you customize your agent's
-behavior however you need.
+AReaLite gives you flexibility in how you design your agents to run **an episode**. **An
+episode** defines how your agent rollouts a complete training sample from an input
+prompt, using tools, reward functions, and (multi-turn) generation. Instead of rigid
+`Agent` classes that might constrain your agent's capabilities, AReaLite captures all
+rollout behavior in a `RolloutWorkflow` class. This approach allows you to customize
+your agent's behavior however you need.
 
 ```python
 # arealite/api/workflow_api.py
@@ -40,7 +42,7 @@ interact.
 > generated from that prompt—it's not batched. However, you can generate multiple
 > trajectories from a single prompt (for example, with GRPO or tree search).
 
-### Setting Up the Multi-Turn Math Workflow
+### Setting Up the Multi-turn Math Workflow
 
 Let's build a multi-turn rollout workflow for solving math problems. First, we'll define
 the `__init__` method to set up what we need during rollout:
@@ -82,10 +84,11 @@ class MultiTurnWorkflow(RolloutWorkflow):
         seq, logprobs, loss_mask, versions = [], [], [], []
         messages = data["messages"]
         # Run multi-turn rollout until we get the correct answer
-        t = reward = 0
+        turn_index = 0
+        reward = 0
         discount = 1.0
         rid = uuid.uuid4().hex
-        while reward == 0 and t < self.max_turns:
+        while reward == 0 and turn_index < self.max_turns:
             # Convert the conversation into input tokens
             input_ids = self.tokenizer.apply_chat_template(
                 messages,
@@ -111,7 +114,7 @@ class MultiTurnWorkflow(RolloutWorkflow):
 > **Note**: The `rid` field in `LLMRequest` is the request ID. Requests with the same ID
 > will reuse the LLM inference server's KV caches for better efficiency.
 
-### Handling Multi-Turn Conversations
+### Handling Multi-turn Conversations
 
 Next, we'll check if the current answer is correct using our `reward_fn`. This function
 should return 1 for correct answers and 0 otherwise. When the answer is wrong, we'll

docs/tutorial/quickstart.md

@@ -75,7 +75,6 @@ python3 -m arealite.launcher.ray examples/arealite/gsm8k_grpo.py \
     allocation_mode=sglang.d12p1t1+d4p1t1 \
     cluster.n_nodes=4 \
     cluster.n_gpus_per_node=4 \
-    ...
 
 # Launch with Slurm launcher. 16 nodes (8 GPUs each), 12 nodes for generation, 4 nodes for training
 python3 -m arealite.launcher.slurm examples/arealite/gsm8k_grpo.py \
@@ -85,7 +84,6 @@ python3 -m arealite.launcher.slurm examples/arealite/gsm8k_grpo.py \
     allocation_mode=sglang.d96p1t1+d32p1t1 \
     cluster.n_nodes=16 \
     cluster.n_gpus_per_node=8 \
-    ...
 ```
 
 Additional references:
@@ -99,9 +97,9 @@ Additional references:
 >
 > 1. Ensure `allocation_mode` matches your cluster configuration
 >    (`#GPUs == cluster.n_nodes * cluster.n_gpus_per_node`)
-> 1. Ray/Slurm launchers only works for more than 1 node (`cluster.n_nodes > 1`). For
+> 1. Ray / Slurm launchers only works for more than 1 node (`cluster.n_nodes > 1`). For
 >    single node scenario, please use `LocalLauncher`.
-> 1. In Ray/Slurm launchers, GPUs are allocated at node granularity, which means #GPUs
+> 1. In Ray / Slurm launchers, GPUs are allocated at node granularity, which means #GPUs
 >    for generation or training must be integer multiples of `cluster.n_gpus_per_node`.
 
 <!--