From 7c844a8a85ff8bc9b14d791ef3e70032da897f4c Mon Sep 17 00:00:00 2001 From: David Lemming Date: Tue, 26 May 2026 17:17:55 +0200 Subject: [PATCH 1/5] Fix misleading data_wait on cuda --- .../_callbacks/tqdm_progress_bar.py | 60 +++++++++++++++--- src/lightly_train/_methods/method.py | 63 ++++++++++++------- 2 files changed, 93 insertions(+), 30 deletions(-) diff --git a/src/lightly_train/_callbacks/tqdm_progress_bar.py b/src/lightly_train/_callbacks/tqdm_progress_bar.py index be27519e0..61f5a93ed 100644 --- a/src/lightly_train/_callbacks/tqdm_progress_bar.py +++ b/src/lightly_train/_callbacks/tqdm_progress_bar.py @@ -10,6 +10,7 @@ import time from typing import Any +import torch from pytorch_lightning import LightningModule, Trainer from pytorch_lightning.callbacks import TQDMProgressBar @@ -17,21 +18,59 @@ class DataWaitTQDMProgressBar(TQDMProgressBar): """ Customizes the progress bar to include compute efficiency. + + On CUDA, ``on_train_batch_end`` fires before the GPU finishes the backward + pass. Recording ``batch_end_time`` there with ``time.perf_counter()`` + causes the remaining async GPU work to bleed into the next step's + ``data_time``, inflating ``data_wait`` (e.g. ~0.5 % true → ~60 % reported). + + Fix for CUDA: + - Bracket each step with CUDA events to get true GPU duration. + - Measure wall-clock time from step-start to step-start (= GPU + data). + - Subtract GPU duration to isolate data-loading wait. + - Events are queried at the next step start, by which point the stream + has advanced past both events — no CPU stall. + + CPU training falls back to simple wall-clock bookkeeping (synchronous + compute means ``on_train_batch_end`` timing is accurate). """ def __init__(self) -> None: super().__init__(refresh_rate=5) self.batch_start_time: float | None = None - self.batch_end_time: float | None = None + self.batch_end_time: float | None = None # CPU-only self.data_time: float | None = None self.batch_time: float | None = None + # CUDA-only: events bracketing the GPU step. + self._step_start_event: torch.cuda.Event | None = None + self._step_end_event: torch.cuda.Event | None = None def on_train_batch_start( self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int ) -> None: - self.batch_start_time = time.perf_counter() - if self.batch_end_time is not None: - self.data_time = self.batch_start_time - self.batch_end_time + now = time.perf_counter() + + if trainer.strategy.root_device.type == "cuda": + # 1. Resolve previous step's GPU duration from events + if self._step_start_event is not None and self._step_end_event is not None: + self.batch_time = ( + self._step_start_event.elapsed_time(self._step_end_event) / 1e3 # type: ignore[no-untyped-call] + ) + + # 2. data_time = wall gap (start -> start) minus GPU duration + if self.batch_start_time is not None and self.batch_time is not None: + wall_gap = now - self.batch_start_time + self.data_time = max(0.0, wall_gap - self.batch_time) + + # 3. Begin new step: record start event + self._step_start_event = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] + self._step_start_event.record() # type: ignore[no-untyped-call] + else: + # CPU: data_time is simply the gap between previous step end and now. + if self.batch_end_time is not None: + self.data_time = now - self.batch_end_time + + self.batch_start_time = now return super().on_train_batch_start(trainer, pl_module, batch, batch_idx) def on_train_batch_end( @@ -42,9 +81,16 @@ def on_train_batch_end( batch: Any, batch_idx: int, ) -> None: - self.batch_end_time = time.perf_counter() - if self.batch_start_time is not None: - self.batch_time = self.batch_end_time - self.batch_start_time + if trainer.strategy.root_device.type == "cuda": + # Record end event — queried at next step start + self._step_end_event = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] + self._step_end_event.record() # type: ignore[no-untyped-call] + elif self.batch_start_time is not None: + # CPU: batch_time is just wall-clock step duration + now = time.perf_counter() + if self.batch_start_time is not None: + self.batch_time = now - self.batch_start_time + self.batch_end_time = now return super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) def get_metrics( diff --git a/src/lightly_train/_methods/method.py b/src/lightly_train/_methods/method.py index 950681549..ae2280d1e 100644 --- a/src/lightly_train/_methods/method.py +++ b/src/lightly_train/_methods/method.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from typing import Any, Literal, Mapping +import torch from lightly.utils.scheduler import CosineWarmupScheduler from pytorch_lightning import LightningModule from pytorch_lightning.loggers import WandbLogger as LightningWandbLogger @@ -44,7 +45,10 @@ class TrainingStepResult: @dataclass class BatchStartEndTime: batch_start_s: float | None = None - batch_end_s: float | None = None + batch_end_s: float | None = None # CPU-only + # CUDA-only: events bracketing the GPU step. + _start_event: torch.cuda.Event | None = None + _end_event: torch.cuda.Event | None = None class Method(LightningModule): @@ -191,27 +195,40 @@ def _log_example_views(self, train_batch: Batch) -> None: ) def _log_time_batch_start(self) -> None: - self.batch_start_end_time.batch_start_s = time.perf_counter() - if self.batch_start_end_time.batch_end_s is not None: - assert ( - self.batch_start_end_time.batch_start_s - > self.batch_start_end_time.batch_end_s - ) - self.log( - "profiling/data_time", - self.batch_start_end_time.batch_start_s - - self.batch_start_end_time.batch_end_s, - ) + now = time.perf_counter() + t = self.batch_start_end_time + + if self.device.type == "cuda": + # 1. Resolve previous step's GPU duration from events. + if t._start_event is not None and t._end_event is not None: + batch_time = t._start_event.elapsed_time(t._end_event) / 1e3 # type: ignore[no-untyped-call] + self.log("profiling/batch_time", batch_time) + + # 2. data_time = wall gap (start → start) minus GPU duration. + if t.batch_start_s is not None: + wall_gap = now - t.batch_start_s + self.log("profiling/data_time", max(0.0, wall_gap - batch_time)) + + # 3. Begin new step: record start event. + t._start_event = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] + t._start_event.record() # type: ignore[no-untyped-call] + else: + # CPU: data_time is simply the gap between previous step end and now. + if t.batch_end_s is not None: + self.log("profiling/data_time", now - t.batch_end_s) + + t.batch_start_s = now def _log_time_batch_end(self) -> None: - self.batch_start_end_time.batch_end_s = time.perf_counter() - if self.batch_start_end_time.batch_start_s is not None: - assert ( - self.batch_start_end_time.batch_end_s - > self.batch_start_end_time.batch_start_s - ) - self.log( - "profiling/batch_time", - self.batch_start_end_time.batch_end_s - - self.batch_start_end_time.batch_start_s, - ) + t = self.batch_start_end_time + + if self.device.type == "cuda": + # Record end event — queried at next step start. + t._end_event = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] + t._end_event.record() # type: ignore[no-untyped-call] + else: + # CPU: batch_time is just wall-clock step duration. + now = time.perf_counter() + if t.batch_start_s is not None: + self.log("profiling/batch_time", now - t.batch_start_s) + t.batch_end_s = now From 936f64f1344ba3b4e6a7bfc939ca2cd667cf17c9 Mon Sep 17 00:00:00 2001 From: David Lemming Date: Tue, 26 May 2026 17:28:49 +0200 Subject: [PATCH 2/5] Update CHANGELOG.md --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 025010e1b..455a1c352 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm ### Removed ### Fixed +- Fix `data_time`, `batch_time` and `data_wait` to report correct metrics on cuda. ### Security From e6eba50ddee2e850887d0cd3fc0df124582f359c Mon Sep 17 00:00:00 2001 From: David Lemming Date: Wed, 27 May 2026 14:31:29 +0200 Subject: [PATCH 3/5] Format CHANGELOG.md --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 455a1c352..ce569a2da 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm ### Removed ### Fixed + - Fix `data_time`, `batch_time` and `data_wait` to report correct metrics on cuda. ### Security From e281e1468ffda8cf9f3055fadf19dc217376ca58 Mon Sep 17 00:00:00 2001 From: David Lemming Date: Wed, 3 Jun 2026 15:58:11 +0200 Subject: [PATCH 4/5] Deduplicate event-bracketing timing logic by reading from method's logs --- .../_callbacks/tqdm_progress_bar.py | 86 ++----------------- 1 file changed, 6 insertions(+), 80 deletions(-) diff --git a/src/lightly_train/_callbacks/tqdm_progress_bar.py b/src/lightly_train/_callbacks/tqdm_progress_bar.py index 61f5a93ed..13139fb81 100644 --- a/src/lightly_train/_callbacks/tqdm_progress_bar.py +++ b/src/lightly_train/_callbacks/tqdm_progress_bar.py @@ -7,98 +7,24 @@ # from __future__ import annotations -import time -from typing import Any - -import torch from pytorch_lightning import LightningModule, Trainer from pytorch_lightning.callbacks import TQDMProgressBar class DataWaitTQDMProgressBar(TQDMProgressBar): - """ - Customizes the progress bar to include compute efficiency. - - On CUDA, ``on_train_batch_end`` fires before the GPU finishes the backward - pass. Recording ``batch_end_time`` there with ``time.perf_counter()`` - causes the remaining async GPU work to bleed into the next step's - ``data_time``, inflating ``data_wait`` (e.g. ~0.5 % true → ~60 % reported). - - Fix for CUDA: - - Bracket each step with CUDA events to get true GPU duration. - - Measure wall-clock time from step-start to step-start (= GPU + data). - - Subtract GPU duration to isolate data-loading wait. - - Events are queried at the next step start, by which point the stream - has advanced past both events — no CPU stall. - - CPU training falls back to simple wall-clock bookkeeping (synchronous - compute means ``on_train_batch_end`` timing is accurate). - """ + """Customizes the progress bar to include compute efficiency.""" def __init__(self) -> None: super().__init__(refresh_rate=5) - self.batch_start_time: float | None = None - self.batch_end_time: float | None = None # CPU-only - self.data_time: float | None = None - self.batch_time: float | None = None - # CUDA-only: events bracketing the GPU step. - self._step_start_event: torch.cuda.Event | None = None - self._step_end_event: torch.cuda.Event | None = None - - def on_train_batch_start( - self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int - ) -> None: - now = time.perf_counter() - - if trainer.strategy.root_device.type == "cuda": - # 1. Resolve previous step's GPU duration from events - if self._step_start_event is not None and self._step_end_event is not None: - self.batch_time = ( - self._step_start_event.elapsed_time(self._step_end_event) / 1e3 # type: ignore[no-untyped-call] - ) - - # 2. data_time = wall gap (start -> start) minus GPU duration - if self.batch_start_time is not None and self.batch_time is not None: - wall_gap = now - self.batch_start_time - self.data_time = max(0.0, wall_gap - self.batch_time) - - # 3. Begin new step: record start event - self._step_start_event = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] - self._step_start_event.record() # type: ignore[no-untyped-call] - else: - # CPU: data_time is simply the gap between previous step end and now. - if self.batch_end_time is not None: - self.data_time = now - self.batch_end_time - - self.batch_start_time = now - return super().on_train_batch_start(trainer, pl_module, batch, batch_idx) - - def on_train_batch_end( - self, - trainer: Trainer, - pl_module: LightningModule, - outputs: Any, - batch: Any, - batch_idx: int, - ) -> None: - if trainer.strategy.root_device.type == "cuda": - # Record end event — queried at next step start - self._step_end_event = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] - self._step_end_event.record() # type: ignore[no-untyped-call] - elif self.batch_start_time is not None: - # CPU: batch_time is just wall-clock step duration - now = time.perf_counter() - if self.batch_start_time is not None: - self.batch_time = now - self.batch_start_time - self.batch_end_time = now - return super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) def get_metrics( self, trainer: Trainer, pl_module: LightningModule ) -> dict[str, int | str | float | dict[str, float]]: metrics = super().get_metrics(trainer, pl_module) - if self.batch_time is not None and self.data_time is not None: - if self.batch_time + self.data_time > 0: - data_wait = self.data_time / (self.batch_time + self.data_time) + batch_time = trainer.callback_metrics.get("profiling/batch_time") + data_time = trainer.callback_metrics.get("profiling/data_time") + if batch_time is not None and data_time is not None: + if batch_time + data_time > 0: + data_wait = data_time / (batch_time + data_time) metrics["data_wait"] = f"{data_wait:.1%}" return metrics From 088302784e900d55cb681d8e3027f91cf140ae7a Mon Sep 17 00:00:00 2001 From: David Lemming Date: Wed, 3 Jun 2026 15:58:39 +0200 Subject: [PATCH 5/5] Ensure _end_event has been recorded before elapsed_time() --- src/lightly_train/_methods/method.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/lightly_train/_methods/method.py b/src/lightly_train/_methods/method.py index ae2280d1e..b60754654 100644 --- a/src/lightly_train/_methods/method.py +++ b/src/lightly_train/_methods/method.py @@ -200,7 +200,11 @@ def _log_time_batch_start(self) -> None: if self.device.type == "cuda": # 1. Resolve previous step's GPU duration from events. - if t._start_event is not None and t._end_event is not None: + if ( + t._start_event is not None + and t._end_event is not None + and t._end_event.query() # type: ignore[no-untyped-call] + ): batch_time = t._start_event.elapsed_time(t._end_event) / 1e3 # type: ignore[no-untyped-call] self.log("profiling/batch_time", batch_time)