torch_profile.py — PyTorch Profiler

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

文件概述

基于 PyTorch 内置 torch.profiler 的性能分析实现，生成 Chrome Trace 格式的性能报告。

核心类

Profiler

class Profiler(DistProfiler):
    def start(self, **kwargs):
        self.prof = get_torch_profiler(contents=self.contents, save_path=self.save_path, ...)
        self.prof.start()

    def stop(self):
        self.prof.step()
        self.prof.stop()  # 触发 trace_handler 保存文件

    def annotate(self, message=None, ...):
        def wrapper(*args, **kwargs):
            with torch.profiler.record_function(profile_name):
                return func(*args, **kwargs)

输出格式

生成 .json.gz 文件，可以在 Chrome 浏览器的 chrome://tracing 中打开查看。

与其他模块的关系

• 被 profile.py 的 DistProfiler 在 tool="torch" 时使用
• 不需要额外安装（PyTorch 内置）

小结

跨平台的 PyTorch Profiler 封装，生成可视化的性能报告。