perf: enable pin_memory on CUDA DataLoaders for faster CPU->GPU transfer

torch DataLoader's pin_memory was never set, defaulting to False. Setting
pin_memory=True allocates batch tensors in page-locked memory, which lets
PyTorch issue async cuda host->device copies via DMA - typically 10-30%
faster transfer for image/video training tensors and can reduce per-step
overhead noticeably on Windows where num_workers is forced to 0 (see
6d31c6db / #770) and the main process pays the full transfer cost.

pin_memory is orthogonal to num_workers and is safe with num_workers=0,
so this benefits Windows/macOS users (who currently can't use multi-worker
loading due to pickling constraints) without changing any other behavior.
Gated on torch.cuda.is_available() so it's a no-op on CPU/MPS/ROCm-without
-CUDA-API setups.

Refs #758.
This commit is contained in:
Booyaka101 committed 2026-05-01 23:07:43 +08:00
1 parent 0d91fcee9e
commit 67e223b502
1 file changed
+7
+7
View File
@@ -694,6 +694,13 @@ def get_dataloader_from_datasets(
dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers
dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor
# Enable pinned memory for faster CPU->GPU transfer when CUDA is available.
# pin_memory is independent of num_workers and works even with num_workers=0
# (the forced setting on native Windows / macOS due to dataset pickling
# constraints). No-op on non-CUDA backends.
if torch.cuda.is_available():
dataloader_kwargs['pin_memory'] = True
if has_buckets:
# make sure they all have buckets
for dataset in datasets: