作者将诊断工具集成到了 TraceML（一个开源 PyTorch 可观测性项目）中，以帮助用户区分由 GPU 计算限制和输入限制导致的训练性能问题。这个新功能允许开发人员快速确定低 GPU 利用率是由数据加载缓慢还是模型计算能力限制引起的。

  • 该集成向 Hugging Face `Trainer` 添加了 `TraceMLTrainerCallback`，只需进行最少的设置更改。
  • 它将输入等待时间与步骤工作分开，以提供关于瓶颈类型的实用判断。
  • 使用 ResNet-50 的 Colab 示例表明，调整 DataLoader 设置将诊断从输入限制改为计算限制，并将速度提高了 1.83 倍。

该工具通过澄清影响训练运行的具体约束，帮助用户避免优化错误的组件。