profiler bug fix for agent loop

Author: tangmengcheng

verl/experimental/agent_loop/agent_loop.py

@@ -781,7 +781,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
         Returns:
             DataProto: Output batch.
         """
-
+        self.start_profile(async_start=True)
         if self.config.actor_rollout_ref.rollout.free_cache_engine:
             self.wake_up()
         if self.reward_model_manager and self.config.reward_model.rollout.free_cache_engine:
@@ -803,7 +803,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
         # calculate performance metrics
         metrics = [output.meta_info.pop("metrics") for output in outputs]  # List[List[Dict[str, str]]]
         timing = self._performance_metrics(metrics, output)
-
+        self.stop_profile()
         output.meta_info = {"timing": timing, **outputs[0].meta_info}
         return output
 
@@ -837,6 +837,14 @@ def sleep(self):
         """Sleep all rollout replica instances."""
         self._run_all([replica.sleep() for replica in self.rollout_replicas])
 
+    def start_profile(self, **kwargs):
+        """Start profiling on all replicas."""
+        self._run_all([replica.start_profile(**kwargs) for replica in self.rollout_replicas])
+
+    def stop_profile(self):
+        """Stop profiling on all replicas."""
+        self._run_all([replica.stop_profile() for replica in self.rollout_replicas])
+
     def _run_all(self, tasks: list[asyncio.Task]):
         async def run_all():
             await asyncio.gather(*tasks)

verl/trainer/ppo/ray_trainer.py

@@ -897,7 +897,7 @@ def _load_checkpoint(self):
     def _start_profiling(self, do_profile: bool) -> None:
         """Start profiling for all worker groups if profiling is enabled."""
         if do_profile:
-            self.actor_rollout_wg.start_profile(role="e2e", profile_step=self.global_steps)
+            self.actor_rollout_wg.start_profile(profile_step=self.global_steps)
             if self.use_reference_policy:
                 self.ref_policy_wg.start_profile(profile_step=self.global_steps)
             if self.use_critic:

verl/utils/profiler/mstx_profile.py

@@ -192,24 +192,31 @@ def __init__(self, rank: int, config: ProfilerConfig, tool_config: NPUToolConfig
     def start(self, **kwargs):
         role, profile_step = kwargs.get("role", None), kwargs.get("profile_step", None)
         profile_step = str(profile_step) if profile_step is not None else None
+        self.async_start = kwargs.get("async_start", False)
         if self.enable and self.this_rank:
             self.this_step = True
-            if not self.discrete and NPUProfiler._define_count == 0:
+            if (not self.discrete or self.async_start) and NPUProfiler._define_count == 0:
+                if not self.discrete:
+                    prof_role = "e2e"
+                    prof_step = profile_step
+                else:
+                    prof_role = role
+                    prof_step = None
                 self.profile_npu = get_npu_profiler(
                     contents=self.profile_contents,
                     profile_level=self.profile_level,
                     profile_save_path=self.profile_save_path,
                     analysis=self.analysis,
-                    role=role,
-                    profile_step=profile_step,
+                    role=prof_role,
+                    profile_step=prof_step,
                 )
                 self.profile_npu.start()
                 NPUProfiler._define_count += 1
 
     def stop(self):
         if self.enable and self.this_rank:
             self.this_step = False
-            if not self.discrete and NPUProfiler._define_count == 1:
+            if (not self.discrete or self.async_start) and NPUProfiler._define_count == 1:
                 self.profile_npu.step()
                 self.profile_npu.stop()
                 NPUProfiler._define_count -= 1

verl/workers/rollout/vllm_rollout/vllm_async_server.py

@@ -440,6 +440,16 @@ async def sleep(self):
     async def wait_for_requests_to_drain(self):
         await self.engine.wait_for_requests_to_drain()
 
+    def start_profile(self, **kwargs):
+        """Start profiling on all workers."""
+        if self.workers:
+            ray.get([worker.start_profile.remote(**kwargs) for worker in self.workers])
+
+    def stop_profile(self):
+        """Stop profiling on all workers."""
+        if self.workers:
+            ray.get([worker.stop_profile.remote() for worker in self.workers])
+
 
 @ray.remote(num_cpus=1)
 class vLLMHttpServer(vLLMHttpServerBase):
@@ -559,6 +569,14 @@ async def sleep(self):
         await self.servers[0].wait_for_requests_to_drain.remote()
         await asyncio.gather(*[server.sleep.remote() for server in self.servers])
 
+    async def start_profile(self, **kwargs):
+        """Start profiling on all servers."""
+        await asyncio.gather(*[server.start_profile.remote(**kwargs) for server in self.servers])
+
+    async def stop_profile(self):
+        """Stop profiling on all servers."""
+        await asyncio.gather(*[server.stop_profile.remote() for server in self.servers])
+
 
 def _qwen2_5_vl_dedup_image_tokens(prompt_ids: list[int], processor):
     """Deduplicate consecutive image tokens in prompt_ids for Qwen2.5-VL, since vLLM will replicate the

verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py

@@ -77,7 +77,7 @@
 from verl.utils.distributed import initialize_global_process_group_ray
 from verl.utils.import_utils import deprecated
 from verl.utils.model import get_lora_rank_from_adapter
-from verl.utils.profiler import GPUMemoryLogger
+from verl.utils.profiler import GPUMemoryLogger, mark_start_range, mark_end_range
 from verl.utils.ray_utils import ray_noset_visible_devices
 from verl.utils.torch_functional import get_response_mask, pad_2d_list_to_length
 from verl.utils.vllm import TensorLoRARequest, VLLMHijack, is_version_ge