Fix misleading data_time / batch_time / data_wait metrics on CUDA (async bleed-through)

[DLemming]

What has changed and why?

On CUDA, on_train_batch_end fires before the GPU finishes the backward pass and optimizer step. The previous implementation recorded batch_end_time at that point, so the remaining async GPU work bled into the next step's data_time, rendering all data/batch-time realted metrics misleading / wildly inaccurate.

The result: data_wait reported ~60% when the true value was ~10%. This is actively misleading — users would chase data-loading bottlenecks that don't exist while the GPU was fully utilized the entire time.

_callbacks/tqdm_progress_bar.py:

• Read profiling/data_time and profiling/batch_time from trainer.callback_metrics

_methods/method.py

• Separate branches for CPU and CUDA
• For CUDA, register non-blocking torch.cuda.Event at on_train_batch_start / on_train_batch_end, measuring real gpu time
• data_time = wall_gap (start → start) − GPU duration (from events)
• CPU-only training falls back to the previous time.perf_counter() approach, which remains accurate when compute is synchronous.

How has it been tested?

Changes have been tested on cuda running a SimCLR-Pretraining (bs=1232, res=224px) and a DINOv2-Pretraining (bs=128, global_crop_res=196px). data_wait dropped from ~60% to ~10% for SimCLR, and increased from ~0.8% to ~2.5% for DINOv2.

Did you update CHANGELOG.md?

• Yes
• Not needed (internal change)

Did you update the documentation?

• Yes
• Not needed (internal change without effects for user)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[CLAassistant]

All committers have signed the CLA.

[mrpositron]

/review

[mrpositron]

Thanks for the contribution! I left a few comments.

[DLemming]

Thanks for the feedback.

Implemented your suggestions. The progress bar now reads profiling/data_time / profiling/batch_time from trainer.callback_metrics instead of re-timing the GPU, so the event-bracketing logic lives only in Method. I also guarded the elapsed_time crash possibility you mentioned with a non-blocking _end_event.query().

One honest limitation worth noting: When the GPU hasn't reached the end event yet we skip that step's update rather than block, resulting in batch_time/data_time going stale, and metrics being slightly biased toward less compute-bound steps. I think that's an acceptable trade-off vs. the old behavior, which reported flat-out wrong numbers (~60% data_wait instad of ~10%).

In practice I haven't been able to trigger the stale path once even on clearly compute-bound runs (e.g. DINOv2), so something in the loop seems to synchronize often enough that the end event is usually ready by the next step start anyway.

Happy to look more into that edge case later, but for now this seems like a solid improvement over the currently available version.