Merge pull request #813 from Booyaka101/fix/dataloader-pin-memory

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

toolkit/config_modules.py

@@ -1041,6 +1041,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)
         # threads used to prep (decode/resize) items ahead of the VAE while caching latents
         self.cache_latents_num_workers: int = kwargs.get('cache_latents_num_workers', min(6, os.cpu_count() or 1))
         self.extra_values: List[float] = kwargs.get('extra_values', [])

toolkit/data_loader.py

@@ -742,6 +742,16 @@ def get_dataloader_from_datasets(
         os.environ.setdefault('DIFFUSERS_VERBOSITY', 'error')
         os.environ.setdefault('NO_ALBUMENTATIONS_UPDATE', '1')
 
+    # 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:
         # make sure they all have buckets
         for dataset in datasets: