TraceML integrations
Diagnose training in the stack you already use.
Framework integrations add TraceML's step-level diagnosis without replacing your model, dataset, or training workflow.
Keep the framework your team already knows. Add the measurements needed to explain slow steps and idle GPUs.
A callback or wrapper that leaves your loop in charge.
Choose the adapter that matches your training loop.
Hugging Face Trainer
Initialize TraceML once, then swap in TraceMLTrainer where you build your Trainer. Try the real ResNet-50 input-pipeline experiment in Colab.
Read walkthrough
PyTorch Lightning
Add TraceMLCallback to your Trainer. Keep your LightningModule, callbacks, accelerator settings, and training workflow.
Read integration docs
Ray Train
Wrap your TorchTrainer with TraceMLTorchTrainer. Keep Ray scheduling, scaling, worker startup, and rank setup.
Read integration docs
TraceML also runs on a plain PyTorch loop.
No integration is required. Run your script under traceml run and you get the same step-level breakdown of input, transfer, compute, and optimizer time.