I kept running into a frustrating training-performance question: GPU utilization would hover around 50%, loss looked normal, and nothing was obviously broken. Was the model actually the limit, or was the GPU waiting for the next batch?
Those are opposite problems with opposite fixes. But nvidia-smi
and a healthy loss curve do not distinguish them, and I do not want to open a full PyTorch Profiler trace for every ordinary run.
So I added a Hugging Face Trainer
integration to TraceML, an open-source PyTorch diagnostics project I have been building.
from traceml.integrations import huggingface as traceml_hf
from traceml.integrations.huggingface import TraceMLTrainerCallback
traceml_hf.init()
trainer = Trainer(
...,
callbacks=[TraceMLTrainerCallback()],
)
It keeps the rest of the Trainer
setup unchanged, separates input wait from step work, and gives a practical input-bound verdict.
I made a runnable Colab with ResNet-50 and real images to test this properly. It runs twice, changing only Data settings. On my included run, it was 1.83× faster and the diagnosis changed from input-bound to compute-bound.
That number will vary with the CPU/GPU ratio. The useful part is finding out which side of the line your own run is on before optimizing the wrong thing.
I would really value feedback from people running real Trainer workloads, especially cases where this diagnosis is surprising or wrong.