discrete profiler support agent loop

Co-authored-by: tardis-key <[email protected]>

Author: mengchengTang

recipe/fully_async_policy/agent_loop/agent_loop.py

@@ -325,3 +325,9 @@ async def sleep(self):
 
     async def clear_kv_cache(self):
         await asyncio.gather(*[replica.clear_kv_cache() for replica in self.rollout_replicas])
+
+    async def start_profile(self, **kwargs):
+        await asyncio.gather(*[replica.start_profile(**kwargs) for replica in self.rollout_replicas])
+
+    async def stop_profile(self):
+        await asyncio.gather(*[replica.stop_profile() for replica in self.rollout_replicas])

recipe/one_step_off_policy/agent_loop/agent_loop.py

@@ -62,3 +62,9 @@ async def sleep(self):
 
     async def clear_kv_cache(self):
         await asyncio.gather(*[replica.clear_kv_cache() for replica in self.rollout_replicas])
+
+    async def start_profile(self, **kwargs):
+        await asyncio.gather(*[replica.start_profile(**kwargs) for replica in self.rollout_replicas])
+
+    async def stop_profile(self):
+        await asyncio.gather(*[replica.stop_profile() for replica in self.rollout_replicas])

verl/experimental/agent_loop/agent_loop.py

@@ -807,6 +807,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
 
         # Fix for Issue #4147: Always call wake_up() to ensure weight sync
         # The wake_up()/sleep() methods internally check free_cache_engine
+        self.start_profile(role="agent_loop_rollout_generate")
         self.wake_up()
         if self.reward_model_manager:
             self.reward_model_manager.wake_up()
@@ -829,6 +830,7 @@ def generate_sequences(self, prompts: DataProto) -> DataProto:
         timing = self._performance_metrics(metrics, output)
 
         output.meta_info = {"timing": timing, **outputs[0].meta_info}
+        self.stop_profile()
         return output
 
     def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: DataProto) -> dict[str, float]:
@@ -865,6 +867,14 @@ def clear_kv_cache(self):
         """Clear all rollout kv cache, but don`t sleep."""
         self._run_all([replica.clear_kv_cache() for replica in self.rollout_replicas])
 
+    def start_profile(self, **kwargs):
+        """Start profiling on all rollout replicas."""
+        self._run_all([replica.start_profile(**kwargs) for replica in self.rollout_replicas])
+
+    def stop_profile(self):
+        """Stop profiling on all rollout 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

@@ -990,24 +990,24 @@ 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_e2e_profiler(role="e2e", profile_step=self.global_steps)
             if self.use_reference_policy:
-                self.ref_policy_wg.start_profile(profile_step=self.global_steps)
+                self.ref_policy_wg.start_e2e_profiler(profile_step=self.global_steps)
             if self.use_critic:
-                self.critic_wg.start_profile(profile_step=self.global_steps)
+                self.critic_wg.start_e2e_profiler(profile_step=self.global_steps)
             if self.use_rm and not self.use_reward_loop:
-                self.rm_wg.start_profile(profile_step=self.global_steps)
+                self.rm_wg.start_e2e_profiler(profile_step=self.global_steps)
 
     def _stop_profiling(self, do_profile: bool) -> None:
         """Stop profiling for all worker groups if profiling is enabled."""
         if do_profile:
-            self.actor_rollout_wg.stop_profile()
+            self.actor_rollout_wg.stop_e2e_profiler()
             if self.use_reference_policy:
-                self.ref_policy_wg.stop_profile()
+                self.ref_policy_wg.stop_e2e_profiler()
             if self.use_critic:
-                self.critic_wg.stop_profile()
+                self.critic_wg.stop_e2e_profiler()
             if self.use_rm and not self.use_reward_loop:
-                self.rm_wg.stop_profile()
+                self.rm_wg.stop_e2e_profiler()
 
     def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqlen", keep_minibatch=False):
         """Reorder the data on single controller such that each dp rank gets similar total tokens"""

verl/utils/profiler/mstx_profile.py

@@ -173,46 +173,59 @@ def __init__(self, rank: int, config: ProfilerConfig, tool_config: NPUToolConfig
             config = ProfilerConfig(ranks=[], enable=False)
         if not tool_config:
             assert not config.enable, "tool_config must be set when profiler is enabled"
-        self.enable: bool = config.enable
-        if not config.enable:
-            return
-        self.this_step: bool = False
         self.discrete: bool = tool_config.discrete
-        self.this_rank: bool = False
-        self.profile_npu = None
+        self.e2e_profile_npu = None
+        self.capture_profiler_npu = None
         self.profile_contents = tool_config.contents
         self.profile_level = tool_config.level
         self.profile_save_path = config.save_path
         self.analysis = tool_config.analysis
-        if config.all_ranks:
-            self.this_rank = True
-        elif config.ranks:
-            self.this_rank = rank in config.ranks
-
-    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
-        if self.enable and self.this_rank:
-            self.this_step = True
-            if not self.discrete and NPUProfiler._define_count == 0:
-                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,
-                )
-                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:
-                self.profile_npu.step()
-                self.profile_npu.stop()
-                NPUProfiler._define_count -= 1
+
+    def start_e2e_profiler(self, **kwargs):
+        role = kwargs.get("role", None)
+        if not self.discrete and NPUProfiler._define_count == 0:
+            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,
+            )
+            self.profile_npu.start()
+            NPUProfiler._define_count += 1
+
+    def stop_e2e_profiler(self):
+        if not self.discrete and NPUProfiler._define_count == 1:
+            self.profile_npu.step()
+            self.profile_npu.stop()
+            NPUProfiler._define_count -= 1
+
+    def start_capture_profiler(self, **kwargs):
+        """Start an on-demand profiling segment."""
+        role = kwargs.get("role", "")
+
+        if self.discrete:
+            self.capture_profiler_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,
+            )
+            self.capture_profiler_npu.start()
+
+        self._capture_range_id = mark_start_range(message=role)
+
+    def stop_capture_profiler(self):
+        """Stop the on-demand profiling segment."""
+        if hasattr(self, "_capture_range_id"):
+            mark_end_range(self._capture_range_id)
+            del self._capture_range_id
+
+        if self.discrete and getattr(self, "capture_profiler_npu", None):
+            self.capture_profiler_npu.step()
+            self.capture_profiler_npu.stop()
+            del self.capture_profiler_npu
 
     def annotate(self, message: Optional[str] = None, role: Optional[str] = None, **kwargs_outer) -> Callable:
         """Decorate a Worker member function to profile the current rank in the current training step.
@@ -230,42 +243,33 @@ def annotate(self, message: Optional[str] = None, role: Optional[str] = None, **
         def decorator(func):
             @functools.wraps(func)
             def wrapper(*args, **kwargs_inner):
-                if not self.enable:
-                    return func(*args, **kwargs_inner)
-
                 profile_name = message or func.__name__
                 discrete_mode = self.discrete
-                profile_enable = self.this_step and self.enable
-
-                if not profile_enable:
-                    return func(*args, **kwargs_inner)
-
-                if profile_enable:
-                    if not discrete_mode:
-                        mark_range = mark_start_range(message=profile_name)
-                    else:
-                        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_npu.start()
-                        mark_range = mark_start_range(message=profile_name)
+
+                if not discrete_mode:
+                    mark_range = mark_start_range(message=profile_name)
+                else:
+                    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_npu.start()
+                    mark_range = mark_start_range(message=profile_name)
 
                 result = func(*args, **kwargs_inner)
 
-                if profile_enable:
-                    if not discrete_mode:
-                        mark_end_range(mark_range)
-                    else:
-                        mark_end_range(mark_range)
-                        profile_npu.step()
-                        profile_npu.stop()
+                if not discrete_mode:
+                    mark_end_range(mark_range)
+                else:
+                    mark_end_range(mark_range)
+                    profile_npu.step()
+                    profile_npu.stop()
 
                 return result
 
             return wrapper
 
-        return decorator
+        return decorator

verl/utils/profiler/nvtx_profile.py

@@ -129,25 +129,25 @@ def __init__(self, rank: int, config: Optional[ProfilerConfig], tool_config: Opt
         self.enable = config.enable
         if not config.enable:
             return
-        self.this_step: bool = False
         self.discrete: bool = tool_config.discrete
-        self.this_rank: bool = False
-        if config.all_ranks:
-            self.this_rank = True
-        elif config.ranks:
-            self.this_rank = rank in config.ranks
-
-    def start(self, **kwargs):
-        if self.enable and self.this_rank:
-            self.this_step = True
-            if not self.discrete:
-                torch.cuda.profiler.start()
-
-    def stop(self):
-        if self.enable and self.this_rank:
-            self.this_step = False
-            if not self.discrete:
-                torch.cuda.profiler.stop()
+
+    def start_e2e_profiler(self, **kwargs):
+        if not self.discrete:
+            torch.cuda.profiler.start()
+
+    def stop_e2e_profiler(self):
+        if not self.discrete:
+            torch.cuda.profiler.stop()
+
+    def start_capture_profiler(self, **kwargs):
+        """Start an on-demand profiling segment."""
+        if self.discrete:
+            torch.cuda.profiler.start()
+
+    def stop_capture_profiler(self):
+        """Stop the on-demand profiling segment."""
+        if self.discrete:
+            torch.cuda.profiler.stop()
 
     def annotate(
         self,
@@ -176,22 +176,17 @@ def annotate(
         def decorator(func):
             @functools.wraps(func)
             def wrapper(*args, **kwargs_inner):
-                if not self.enable:
-                    return func(*args, **kwargs_inner)
-
                 profile_name = message or func.__name__
 
-                if self.this_step:
-                    if self.discrete:
-                        torch.cuda.profiler.start()
-                    mark_range = mark_start_range(message=profile_name, color=color, domain=domain, category=category)
+                if self.discrete:
+                    torch.cuda.profiler.start()
+                mark_range = mark_start_range(message=profile_name, color=color, domain=domain, category=category)
 
                 result = func(*args, **kwargs_inner)
 
-                if self.this_step:
-                    mark_end_range(mark_range)
-                    if self.discrete:
-                        torch.cuda.profiler.stop()
+                mark_end_range(mark_range)
+                if self.discrete:
+                    torch.cuda.profiler.stop()
 
                 return result