[fully_async] feat: standalone log prob server (Model Engine Server) support

tests/special_sanity/check_device_api_usage.py

@@ -53,6 +53,7 @@
     "verl/checkpoint_engine",  # checkpoint engine backend are device specific
     "verl/utils/modelopt/megatron_qat_patch.py",  # appear in torch.cuda.empty_cache()
     "verl/models/mcore/patch.py",  # checkpoint patch only on cuda
+    "verl/verl/workers/rollout/model_engine_server/model_engine_server.py",  # appear in default device_name
 ]
 
 # directory or file path must contain keyword "nccl"

verl/experimental/agent_loop/agent_loop.py

@@ -531,6 +531,7 @@ async def generate_sequences(self, batch: DataProto) -> DataProto:
             sampling_params["top_p"] = config.val_kwargs.top_p
             sampling_params["top_k"] = config.val_kwargs.top_k
             sampling_params["temperature"] = config.val_kwargs.temperature
+            sampling_params["validate"] = True
 
         # by default, we assume it's a single turn agent
         if "agent_name" not in batch.non_tensor_batch:
@@ -675,6 +676,14 @@ async def _agent_loop_postprocess(self, output, validate, **kwargs) -> _Internal
         if output.response_logprobs is not None:
             pad_size = self.rollout_config.response_length - len(output.response_logprobs)
             response_logprobs = torch.tensor(output.response_logprobs + [0.0] * pad_size).unsqueeze(0)
+        _es_logprobs = output.extra_fields.pop("engine_server_logprobs", None)
+        if _es_logprobs is not None:
+            pad_size = self.rollout_config.response_length - len(_es_logprobs)
+            output.extra_fields["engine_server_logprobs"] = torch.tensor(_es_logprobs + [0.0] * pad_size).unsqueeze(0)
+        _es_entropys = output.extra_fields.pop("engine_server_entropys", None)
+        if _es_entropys is not None:
+            pad_size = self.rollout_config.response_length - len(_es_entropys)
+            output.extra_fields["engine_server_entropys"] = torch.tensor(_es_entropys + [0.0] * pad_size).unsqueeze(0)
 
         response_mask = response_mask_output["input_ids"] * response_output["attention_mask"]
         attention_mask = torch.cat([prompt_output["attention_mask"], response_output["attention_mask"]], dim=1)
@@ -884,6 +893,14 @@ def _postprocess(
         optional_outputs = {}
         if inputs[0].response_logprobs is not None:
             optional_outputs["rollout_log_probs"] = torch.cat([input.response_logprobs for input in inputs], dim=0)
+        if inputs[0].extra_fields.get("engine_server_logprobs") is not None:
+            optional_outputs["engine_server_logprobs"] = torch.cat(
+                [input.extra_fields.pop("engine_server_logprobs") for input in inputs], dim=0
+            )
+        if inputs[0].extra_fields.get("engine_server_entropys") is not None:
+            optional_outputs["engine_server_entropys"] = torch.cat(
+                [input.extra_fields.pop("engine_server_entropys") for input in inputs], dim=0
+            )
         if inputs[0].routed_experts is not None:
             optional_outputs["routed_experts"] = torch.cat([input.routed_experts for input in inputs], dim=0)
         if inputs[0].teacher_logprobs is not None and inputs[0].teacher_ids is not None:
@@ -1132,6 +1149,7 @@ async def _init_agent_loop_workers(self):
                     teacher_servers,
                     teacher_load_balancer_handle,
                     self.reward_loop_worker_handles,
+                    self.model_engine_server_handle,
                 )
             )
 

verl/experimental/agent_loop/single_turn_agent_loop.py

@@ -68,6 +68,12 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutpu
             metrics["num_preempted"] = output.num_preempted if output.num_preempted is not None else -1
         response_mask = [1] * len(output.token_ids)
 
+        extra_fields = output.extra_fields
+        if extra_fields.get("engine_server_logprobs"):
+            extra_fields["engine_server_logprobs"] = extra_fields["engine_server_logprobs"][: self.response_length]
+        if extra_fields.get("engine_server_entropys"):
+            extra_fields["engine_server_entropys"] = extra_fields["engine_server_entropys"][: self.response_length]
+
         output: AgentLoopOutput = AgentLoopOutput(
             prompt_ids=prompt_ids,
             response_ids=output.token_ids[: self.response_length],
@@ -81,7 +87,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutpu
             multi_modal_data=multi_modal_data,
             num_turns=2,
             metrics=metrics,
-            extra_fields=output.extra_fields,
+            extra_fields=extra_fields,
         )
 
         # keeping the schema consistent with tool_agent_loop

verl/experimental/agent_loop/tool_agent_loop.py

@@ -184,6 +184,12 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutpu
         if agent_data.video_data is not None:
             multi_modal_data["videos"] = agent_data.video_data
 
+        extra_fields = agent_data.extra_fields
+        if extra_fields.get("engine_server_logprobs"):
+            extra_fields["engine_server_logprobs"] = extra_fields["engine_server_logprobs"][: self.response_length]
+        if extra_fields.get("engine_server_entropys"):
+            extra_fields["engine_server_entropys"] = extra_fields["engine_server_entropys"][: self.response_length]
+
         output: AgentLoopOutput = AgentLoopOutput(
             prompt_ids=prompt_ids,
             response_ids=response_ids[: self.response_length],
@@ -195,7 +201,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutpu
             num_turns=agent_data.user_turns + agent_data.assistant_turns + 1,
             metrics=agent_data.metrics,
             routed_experts=agent_data.routed_experts,
-            extra_fields=agent_data.extra_fields,
+            extra_fields=extra_fields,
         )
         output.extra_fields.update({"turn_scores": agent_data.turn_scores, "tool_rewards": agent_data.tool_rewards})
         return output
@@ -232,8 +238,9 @@ async def _handle_generating_state(
         else:
             agent_data.metrics["num_preempted"] += output.num_preempted if output.num_preempted is not None else 0
 
+        _ACCUMULATED_KEYS = {"engine_server_logprobs", "engine_server_entropys"}
         if not agent_data.extra_fields:
-            agent_data.extra_fields.update(output.extra_fields)
+            agent_data.extra_fields.update({k: v for k, v in output.extra_fields.items() if k not in _ACCUMULATED_KEYS})
         else:
             # Multi-round calls, only update the maximum max_global_steps.
             max_global_steps = output.extra_fields.get("max_global_steps", None)
@@ -246,7 +253,12 @@ async def _handle_generating_state(
         agent_data.response_mask += [1] * len(agent_data.response_ids)
         if output.log_probs:
             agent_data.response_logprobs += output.log_probs
-
+        if output.extra_fields.get("engine_server_logprobs"):
+            agent_data.extra_fields.setdefault("engine_server_logprobs", [])
+            agent_data.extra_fields["engine_server_logprobs"] += output.extra_fields["engine_server_logprobs"]
+        if output.extra_fields.get("engine_server_entropys"):
+            agent_data.extra_fields.setdefault("engine_server_entropys", [])
+            agent_data.extra_fields["engine_server_entropys"] += output.extra_fields["engine_server_entropys"]
         if output.routed_experts is not None:
             agent_data.routed_experts = output.routed_experts
 
@@ -378,6 +390,10 @@ async def _handle_processing_tools_state(self, agent_data: AgentData) -> AgentSt
         agent_data.response_mask += [0] * len(response_ids)
         if agent_data.response_logprobs:
             agent_data.response_logprobs += [0.0] * len(response_ids)
+        if agent_data.extra_fields.get("engine_server_logprobs"):
+            agent_data.extra_fields["engine_server_logprobs"] += [0.0] * len(response_ids)
+        if agent_data.extra_fields.get("engine_server_entropys"):
+            agent_data.extra_fields["engine_server_entropys"] += [0.0] * len(response_ids)
         agent_data.user_turns += 1
         return AgentState.GENERATING
 
@@ -410,6 +426,10 @@ async def _handle_interacting_state(self, agent_data: AgentData) -> AgentState:
         agent_data.response_mask += [0] * len(response_ids)
         if agent_data.response_logprobs:
             agent_data.response_logprobs += [0.0] * len(response_ids)
+        if agent_data.extra_fields.get("engine_server_logprobs"):
+            agent_data.extra_fields["engine_server_logprobs"] += [0.0] * len(response_ids)
+        if agent_data.extra_fields.get("engine_server_entropys"):
+            agent_data.extra_fields["engine_server_entropys"] += [0.0] * len(response_ids)
 
         # double check prompt
         # Check termination condition

verl/experimental/fully_async_policy/agent_loop/agent_loop.py

@@ -44,6 +44,16 @@ class FullyAsyncLLMServerManager(AsyncLLMServerManager):
     invisible to the AgentLoop.
     """
 
+    def __init__(
+        self,
+        config: DictConfig,
+        servers: list[tuple[str, ray.actor.ActorHandle]],
+        load_balancer_handle: ray.actor.ActorHandle,
+        model_engine_server_handle: ray.actor.ActorHandle = None,
+    ):
+        super().__init__(config=config, servers=servers, load_balancer_handle=load_balancer_handle)
+        self.model_engine_server_handle = model_engine_server_handle
+
     @rollout_trace_op
     async def generate(
         self,
@@ -74,6 +84,7 @@ async def generate(
         elif "max_new_tokens" in sampling_params:
             limit_key = "max_new_tokens"
         original_max_tokens = sampling_params.get(limit_key) if limit_key else None
+        validate = sampling_params.pop("validate", None)
 
         final_output = TokenOutput(
             token_ids=[],
@@ -91,11 +102,25 @@ async def generate(
                 image_data=image_data,
                 video_data=video_data,
             )
-
+            current_prompt_ids = prompt_ids + final_output.token_ids
+            current_temperature = sampling_params.get("temperature", 1.0)
+            # Skip old log prob computation during validation
+            if not validate:
+                output = await self._compute_old_log_prob(output, current_prompt_ids, current_temperature)
             # 2. merge output into final_output
             final_output.token_ids.extend(output.token_ids)
             if output.log_probs is not None:
                 final_output.log_probs.extend(output.log_probs)
+            if output.extra_fields.get("engine_server_logprobs") is not None:
+                final_output.extra_fields.setdefault("engine_server_logprobs", [])
+                final_output.extra_fields["engine_server_logprobs"].extend(
+                    output.extra_fields["engine_server_logprobs"]
+                )
+            if output.extra_fields.get("engine_server_entropys") is not None:
+                final_output.extra_fields.setdefault("engine_server_entropys", [])
+                final_output.extra_fields["engine_server_entropys"].extend(
+                    output.extra_fields["engine_server_entropys"]
+                )
             if output.routed_experts is not None:
                 if final_output.routed_experts is None:
                     final_output.routed_experts = output.routed_experts
@@ -121,11 +146,30 @@ async def generate(
             # 4. check stop reason
             if output.stop_reason not in ("aborted", "abort") or not self.config.async_training.partial_rollout:
                 break
+        if validate:  # restore for multi-turn use
+            sampling_params["validate"] = validate
         final_output.extra_fields["global_steps"] = global_steps
         final_output.extra_fields["min_global_steps"] = min_global_steps
         final_output.extra_fields["max_global_steps"] = max_global_steps
         return final_output
 
+    async def _compute_old_log_prob(self, output: TokenOutput, context_prompt_ids, temperature: float):
+        if self.model_engine_server_handle is None:
+            return output
+
+        # Only recompute old_log_probs for newly generated tokens in this turn.
+        if len(output.token_ids) == 0:
+            output.extra_fields["engine_server_logprobs"] = []
+            output.extra_fields["engine_server_entropys"] = []
+            return output
+
+        result = await self.model_engine_server_handle.compute_log_prob.remote(
+            context_prompt_ids, output.token_ids, temperature
+        )
+        output.extra_fields["engine_server_logprobs"] = result["log_probs"]
+        output.extra_fields["engine_server_entropys"] = result["entropy"]
+        return output
+
 
 @ray.remote
 class FullyAsyncAgentLoopWorker(AgentLoopWorker):
@@ -137,8 +181,14 @@ def __init__(
         teacher_servers: list[tuple[str, ray.actor.ActorHandle]] = None,
         teacher_load_balancer_handle: ray.actor.ActorHandle = None,
         reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,
+        model_engine_server_handle: ray.actor.ActorHandle = None,
     ):
-        self.server_manager = FullyAsyncLLMServerManager(config, servers, load_balancer_handle)
+        self.server_manager = FullyAsyncLLMServerManager(
+            config,
+            servers,
+            load_balancer_handle,
+            model_engine_server_handle,
+        )
         super().__init__(
             config,
             servers,
@@ -163,6 +213,42 @@ def __init__(
         if self.distillation_enabled:
             raise NotImplementedError("Distillation is not implemented in FullyAsyncAgentLoopManager yet.")
 
+    async def _initialize_llm_servers(self):
+        """Extend base class to also create ModelEngineReplica when configured."""
+        await super()._initialize_llm_servers()
+        if self.config.model_engine_server.enable_standalone:
+            await self._init_model_engine_replica()
+
+    async def _init_model_engine_replica(self):
+        """Create ModelEngineReplica, call init_standalone, and append to rollout_replicas.
+
+        ModelEngineReplica.init_standalone() self-allocates a Ray resource pool,
+        spawns ModelEngineWorker actors, calls init_model(), and creates the
+        OldLogProbServer — all in one call, exactly like vLLM/SGLang replicas.
+
+        After this, self.model_engine_server_handle is set so that
+        _init_agent_loop_workers() passes it to every FullyAsyncAgentLoopWorker.
+        """
+        from verl.workers.rollout.model_engine_server import ModelEngineReplica, ModelEngineWorker
+
+        replica = ModelEngineReplica(
+            replica_rank=len(self.rollout_replicas),
+            full_config=self.config,
+            worker_cls=ModelEngineWorker,
+            rollout_config=self.rollout_config,
+        )
+        await replica.init_standalone()
+        self.rollout_replicas.append(replica)
+
+        # Expose the server handle so _init_agent_loop_workers passes it to workers.
+        self.model_engine_server_handle = replica.servers[0]
+
+    async def shutdown(self):
+        """Shut down OldLogProbServer if one was created."""
+        for replica in self.rollout_replicas:
+            if hasattr(replica, "shutdown"):
+                await replica.shutdown()
+
     @auto_await
     async def generate_sequences_single(self, prompts: DataProto) -> DataProto:
         """Split input batch and dispatch to agent loop workers.

verl/experimental/fully_async_policy/config/fully_async_ppo_megatron_trainer.yaml

@@ -4,6 +4,9 @@ hydra:
 
 defaults:
   - ppo_megatron_trainer
+
+  - model_engine_server@model_engine_server: megatron_model_engine_server
+
   - _self_
 
 trainer:

verl/experimental/fully_async_policy/fully_async_rollouter.py

@@ -609,6 +609,10 @@ async def fit(self):
             # Wait for the task to complete
             await asyncio.gather(generation_task, monitor_task, return_exceptions=True)
 
+            # Shut down OldLogProbServer (if any) — it lives on the Rollouter side.
+            if self.async_rollout_manager is not None:
+                await self.async_rollout_manager.shutdown()
+
         print("[FullyAsyncRollouter] Rollouter fit completed")
 
     async def _async_monitor_loop(self):

verl/experimental/fully_async_policy/fully_async_trainer.py

@@ -481,6 +481,14 @@ def _compute_old_log_prob(self, batch: DataProto):
         If local_trigger_step == 2, 3, ..., restore the parameters of version 1 to calculate the old_log_prob,
         then restore the parameters of the current version.
         """
+        if "engine_server_logprobs" in batch.batch:
+            batch_dict = {
+                "old_log_probs": batch.batch.pop("engine_server_logprobs"),
+                "entropys": batch.batch.pop("engine_server_entropys"),
+            }
+            old_log_prob = DataProto.from_dict(batch_dict)
+            return old_log_prob, 0.0
+
         if self.local_trigger_step == 1:
             self.actor_rollout_wg.save_model_to_cpu(1)
             old_log_prob, old_log_prob_mfu = super()._compute_old_log_prob(batch)