著者は、GPU計算バウンドと入力バウンドのトレーニングパフォーマンスの問題を区別するのを支援するために、オープンソースのPyTorch観測性プロジェクトであるTraceMLに診断ツールを統合しました。この新機能により、開発者は低GPU利用率が低速なデータ読み込みによるものか、実際のモデル計算の限界によるものかを迅速に特定できます。

  • 統合により、Hugging Face `Trainer` に `TraceMLTrainerCallback` が追加され、最小限の設定変更で済みます。
  • 入力待機時間とステップ作業を分離し、ボトルネックの種類に関する実用的な判断を提供します。
  • ResNet-50を使用したColabの例では、DataLoaderの設定を変更することで診断が入力バウンドから計算バウンドに変更され、速度が1.83倍向上しました。

このツールは、トレーニング実行に影響を与える特定の制約を明確にすることで、ユーザーが間違ったコンポーネントを最適化するのを防ぎます。