diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index fd6f630..98a97ed 100644 --- a/toolkit/config_modules.py +++ b/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', []) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 38b7ad4..5147870 100644 --- a/toolkit/data_loader.py +++ b/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: