[功能] 在训练脚本中添加 DeepSpeed/ZeRO 支持

作者: Tianyi-Franklin-Wang创建于 2025年12月5日更新于 2025年12月5日
标签enhancement

** 您的特性请求与问题有关吗 ? 请说明。 ** DeepSpeed/ZeRO为大规模训练提供了有意义的加速和内存减少. 然而,Timm’s'培训脚本目前不包括官方或内置方式,以启用DeepSpeed/ZeRO. 我注意到早先曾尝试过第490期, 我执行DeepSpeed/ZeRO的支持(有一些妥协), 我很乐意准备一份公关文件,

说明你想要的解决方案 目前的实施将DeepSpeed作为一个可选的依赖性加以整合,并遵循`优化'现有培训结构。 具体而言:

  • 引入一套与当前参数解析样式相匹配的新的 CLI 旗帜,同时允许 DepSpeed 和 pass 配置选项,例如:
  • " -- -- 速度 "
  • '-ds-0级 {0、1、2、3}
  • '--ds-offload-优化器{没有,cpu,nvme}{
  • '-ds-offload-param {没有,cpu,nvme}`}
  • 使用一个小的帮助函数Building ds config',直接从现有的tim'参数(批量大小、梯度累积、AMP dtype、剪接等)构造一个DeepSpeed config dict/JSON。
  • 只在“-----------------------------------------------------------------------------------------------------------------------------------------`------------------------------------------------------------------------------------------------------------------ 这确保了与现有脚本的完全后向兼容,同时为用户提供ZeRO加速的选入路径.

** 其他情况** 目前对"深通"路径中模型EMA,伐木和检查站的处理尚不优雅. 我可改进这些组成部分,如果接受这些组成部分,则使之与现有的公用事业更加一致。 实际上,在DeepSpeed/ZeRO的加速上,我的训练经验相当体面.

内容来源: huggingface/pytorch-image-models