From 0c8793ff47ecb6ae7d090f5879b7f41cadd101bf Mon Sep 17 00:00:00 2001 From: Booyaka101 Date: Sat, 2 May 2026 11:02:41 +0800 Subject: [PATCH] fix: make pin_memory opt-in via dataset config (default off) @zappazack found a real regression: pin_memory page-locks RAM that NVIDIA's Windows driver wants for shared-memory VRAM-overflow fallback when VRAM saturates. Page-locked pages can't be relocated, so the driver falls back to PCIe thrashing - step times went from 1.5s to 20+s when VRAM saturated on RTX 5090 + Qwen LoRA + 1024x1024. The original "strict-improvement" claim was wrong. Convert to opt-in: - Add `pin_memory: bool` field on DatasetConfig (default False) - Only pin when both CUDA is available AND user opts in via config Users with stable VRAM headroom can opt in for the transfer speedup. Users at the VRAM ceiling stay safe with the default. Refs #758. --- toolkit/config_modules.py | 6 ++++++ toolkit/data_loader.py | 13 ++++++++----- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 5ed490e..4c95909 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -977,6 +977,12 @@ class DatasetConfig: self.num_workers: int = kwargs.get('num_workers', 2) self.prefetch_factor: int = kwargs.get('prefetch_factor', 2) + # Pin DataLoader output tensors in page-locked RAM for faster CPU->GPU + # transfer. Off by default because page-locked RAM cannot be relocated + # by NVIDIA's Windows driver shared-memory VRAM-overflow fallback, + # which can cause severe PCIe thrashing for users at the VRAM ceiling. + # Opt in if you have stable VRAM headroom and want the transfer speedup. + self.pin_memory: bool = kwargs.get('pin_memory', False) self.extra_values: List[float] = kwargs.get('extra_values', []) self.square_crop: bool = kwargs.get('square_crop', False) # apply same augmentations to control images. Usually want this true unless special case diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 8f6fb81..6b39e26 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -694,11 +694,14 @@ 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(): + # Enable pinned memory only when explicitly opted in via dataset config and + # CUDA is available. pin_memory speeds up CPU->GPU transfer (helpful even at + # num_workers=0), but page-locked RAM cannot be relocated by NVIDIA's + # Windows driver shared-memory VRAM-overflow fallback, which can cause + # severe PCIe thrashing for users at the VRAM ceiling. Off by default; opt + # in via 'pin_memory: true' on the dataset config when VRAM headroom is + # stable. + if torch.cuda.is_available() and dataset_config_list[0].pin_memory: dataloader_kwargs['pin_memory'] = True if has_buckets: