The author has integrated a diagnostic tool into TraceML, an open-source PyTorch observability project, to help users distinguish between GPU compute-bound and input-bound training performance issues. This new feature allows developers to quickly identify whether low GPU utilization is caused by slow data loading or actual model computation limits.

  • The integration adds a `TraceMLTrainerCallback` to the Hugging Face `Trainer`, requiring minimal setup changes.
  • It separates input wait time from step work to provide a practical verdict on the bottleneck type.
  • A Colab example with ResNet-50 demonstrated that adjusting DataLoader settings changed the diagnosis from input-bound to compute-bound and improved speed by 1.83×.

This tool helps users avoid optimizing the wrong component by clarifying the specific constraint affecting their training runs.