adjust len=120

verl/utils/profiler/__init__.py

@@ -23,4 +23,15 @@
     from .performance import marked_timer
     from .profile import DistProfiler, mark_annotate, mark_end_range, mark_start_range
 
-__all__ = ["GPUMemoryLogger", "log_gpu_memory_usage", "mark_start_range", "mark_end_range", "mark_annotate", "DistProfiler", "DistProfilerExtension", "ProfilerConfig", "simple_timer", "marked_timer"]
+__all__ = [
+    "GPUMemoryLogger",
+    "log_gpu_memory_usage",
+    "mark_start_range",
+    "mark_end_range",
+    "mark_annotate",
+    "DistProfiler",
+    "DistProfilerExtension",
+    "ProfilerConfig",
+    "simple_timer",
+    "marked_timer",
+]

verl/utils/profiler/empty_annotations.py

@@ -15,15 +15,25 @@
 from typing import Callable, Optional
 
 
-def mark_start_range(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> None:
+def mark_start_range(
+    message: Optional[str] = None,
+    color: Optional[str] = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+) -> None:
     pass
 
 
 def mark_end_range(range_id: str) -> None:
     pass
 
 
-def mark_annotate(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> Callable:
+def mark_annotate(
+    message: Optional[str] = None,
+    color: Optional[str] = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+) -> Callable:
     def decorator(func):
         return func
 

verl/utils/profiler/nvtx_profile.py

@@ -23,7 +23,12 @@
 from .profile import DistProfiler, ProfilerConfig
 
 
-def mark_start_range(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> None:
+def mark_start_range(
+    message: Optional[str] = None,
+    color: Optional[str] = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+) -> None:
     """Start a mark range in the profiler.
 
     Args:
@@ -49,7 +54,12 @@ def mark_end_range(range_id: str) -> None:
     return nvtx.end_range(range_id)
 
 
-def mark_annotate(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> Callable:
+def mark_annotate(
+    message: Optional[str] = None,
+    color: Optional[str] = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+) -> Callable:
     """Decorate a function to annotate a mark range along with the function life cycle.
 
     Args:
@@ -71,7 +81,13 @@ def decorator(func):
 
 
 @contextmanager
-def marked_timer(name: str, timing_raw: Dict[str, float], color: str = None, domain: Optional[str] = None, category: Optional[str] = None):
+def marked_timer(
+    name: str,
+    timing_raw: Dict[str, float],
+    color: str = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+):
     """Context manager for timing with NVTX markers.
 
     This utility function measures the execution time of code within its context,
@@ -128,7 +144,12 @@ def stop(self):
                 torch.cuda.profiler.stop()
 
     @staticmethod
-    def annotate(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> Callable:
+    def annotate(
+        message: Optional[str] = None,
+        color: Optional[str] = None,
+        domain: Optional[str] = None,
+        category: Optional[str] = None,
+    ) -> Callable:
         """Decorate a Worker member function to profile the current rank in the current training step.
 
         Requires the target function to be a member function of a Worker, which has a member field `profiler` with NightSystemsProfiler type.

verl/utils/profiler/performance.py

@@ -144,7 +144,13 @@ def simple_timer(name: str, timing_raw: Dict[str, float]):
 
 
 @contextmanager
-def marked_timer(name: str, timing_raw: Dict[str, float], color: str = None, domain: Optional[str] = None, category: Optional[str] = None):
+def marked_timer(
+    name: str,
+    timing_raw: Dict[str, float],
+    color: str = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+):
     """Context manager for timing with platform markers.
 
     This utility function measures the execution time of code within its context,

verl/utils/profiler/profile.py

@@ -72,7 +72,9 @@ def _validate(self):
                 self.config.profile_ranks = [0]
             assert self.config.step_start >= 0, "[ERROR] Profile step start must be greater than 0"
             assert self.config.step_end >= 0, "[ERROR] Profile step end must be greater than 0"
-            assert self.config.step_start < self.config.step_end, "[ERROR] Profile step start must be less than step end"
+            assert self.config.step_start < self.config.step_end, (
+                "[ERROR] Profile step start must be less than step end"
+            )
 
     def check(self):
         return self.prof is not None and not self.skip_prof
@@ -112,7 +114,12 @@ def stop_trace(self):
             self.skip_prof = True
 
 
-def mark_start_range(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> None:
+def mark_start_range(
+    message: Optional[str] = None,
+    color: Optional[str] = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+) -> None:
     """Start a profiling range marker (no-op implementation).
 
     Args:
@@ -133,7 +140,12 @@ def mark_end_range(range_id: str) -> None:
     pass
 
 
-def mark_annotate(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> Callable:
+def mark_annotate(
+    message: Optional[str] = None,
+    color: Optional[str] = None,
+    domain: Optional[str] = None,
+    category: Optional[str] = None,
+) -> Callable:
     """Decorator to annotate a function with profiling markers (no-op implementation).
 
     Args:
@@ -174,7 +186,12 @@ def stop(self):
         pass
 
     @staticmethod
-    def annotate(message: Optional[str] = None, color: Optional[str] = None, domain: Optional[str] = None, category: Optional[str] = None) -> Callable:
+    def annotate(
+        message: Optional[str] = None,
+        color: Optional[str] = None,
+        domain: Optional[str] = None,
+        category: Optional[str] = None,
+    ) -> Callable:
         def decorator(func):
             return func