已开启
[Feature]: 重构 Trainer Profiling 配置与采集流程 #389
townwish4gitcode创建于  25 天前
townwish4gitcode
25 天前 创建

🚀 功能描述

当前 Trainer profiling 实现配置能力有限,仅支持单 rank 和基础 PyTorch
Profiler 参数,与MindFormers等生产级训练框架的profiler 生命周期、输出结构等使用方式
存在较大差异。

希望重构 Trainer profiling,提供统一的配置入口,支持按训练步、rank、PP stage、
采集等级、memory、stack 和 MSTX 等维度控制性能采集,并确保 profiler 在采集窗口
结束后及时停止和释放。

现有替代方案

现有 ProfilingCallback 和 trainer.runtime.profiling 可以完成基础采集,但扩展能力
有限,继续在旧实现上叠加功能会保留两套配置语义并增加维护成本,因此计划直接替换
旧实现。

与DFX相关性DF

提议的新API

将配置入口统一为 TrainerConfig.profiler,主要支持:

  • enabled
  • start_step / stop_step
  • start_on_init
  • memory
  • rank_ids
  • pipeline_stage_leaders
  • output_path
  • level
  • with_stack
  • data_simplification
  • mstx

是否影响现有API

有影响。本次改动将删除旧 profiling 配置、ProfilingCallback 和
trainer.runtime.profiling,不提供旧接口兼容层。

示例 YAML、CLI dotted override 和 Trainer callback 初始化逻辑需要同步更新。

补充信息

Thanks for contributing 🎉!

likedislike