[P] TraceML: wrap your PyTorch training step in single context manager and see what’s slowing training live

End-summary

Building TraceML, an open-source tool for PyTorch training runtime visibility.

You add a single context manager:

with trace_step(model): ... 

and get a live view of training while it runs:

• dataloader fetch time
• forward / backward / optimizer timing
• GPU memory
• median vs worst rank in single-node DDP
• skew to surface imbalance
• compact end-of-run summary with straggler rank and step breakdown

The goal is simple: quickly show answer
why is this training run slower than it should be?

Current support:

• single GPU
• single-node multi-GPU DDP
• Hugging Face Trainer
• PyTorch Lightning callback

Useful for catching:

• slow dataloaders
• rank imbalance / stragglers
• memory issues
• unstable step behavior

Repo: https://github.com/traceopt-ai/traceml/

Please share your runtime summary in issue or here and tell me whether it was actually helpful or what signal is still missing.

If this looks useful, a star would also really help.