profile.py — 分布式分析器核心

文件路径: verl/utils/profiler/profile.py

文件概述

定义 DistProfiler 核心调度器类，根据配置委托给具体的分析后端。还包含 TorchMemoryProfiler 和 DistProfilerExtension。

核心类

DistProfiler — 统一调度器

class DistProfiler:
    def __init__(self, rank, config=None, tool_config=None):
        if self._tool == "nsys":
            self._impl = NsightSystemsProfiler(...)
        elif self._tool == "npu":
            self._impl = NPUProfiler(...)
        elif self._tool == "torch":
            self._impl = Profiler(...)
        elif self._tool == "torch_memory":
            self._impl = TorchMemoryProfiler(...)
        else:
            self._impl = _NoOpProfiler()

    def start(self, **kwargs):
        if self.check_enable() and self.check_this_rank():
            self._impl.start(**kwargs)

    def stop(self):
        if self.check_enable() and self.check_this_rank():
            self._impl.stop()

    @classmethod
    def annotate(cls, message=None, ...):
        """装饰器：为函数添加分析标记"""

DistProfilerExtension

class DistProfilerExtension:
    """在单控制器模式下分发分析命令到所有 rank"""
    @register(dispatch_mode=Dispatch.ONE_TO_ALL)
    def start_profile(self, **kwargs):
        self.profiler.start(**kwargs)

与其他模块的关系

• 被 trainer 创建并传递给 worker
• 调度到 nvtx_profile.py、mstx_profile.py 或 torch_profile.py

小结

统一的分析器调度中心，根据配置选择合适的后端。