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.
This commit is contained in:
Booyaka101 committed 2026-05-02 11:02:41 +08:00
1 parent 67e223b502
commit 0c8793ff47
2 files changed
+14 -5

No files matched your search

+6
View File
@@ -977,6 +977,12 @@ class DatasetConfig:
self.num_workers: int = kwargs.get('num_workers', 2) self.num_workers: int = kwargs.get('num_workers', 2)
self.prefetch_factor: int = kwargs.get('prefetch_factor', 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.extra_values: List[float] = kwargs.get('extra_values', [])
self.square_crop: bool = kwargs.get('square_crop', False) self.square_crop: bool = kwargs.get('square_crop', False)
# apply same augmentations to control images. Usually want this true unless special case # apply same augmentations to control images. Usually want this true unless special case
+8 -5
View File
@@ -694,11 +694,14 @@ def get_dataloader_from_datasets(
dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers
dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor
# Enable pinned memory for faster CPU->GPU transfer when CUDA is available. # Enable pinned memory only when explicitly opted in via dataset config and
# pin_memory is independent of num_workers and works even with num_workers=0 # CUDA is available. pin_memory speeds up CPU->GPU transfer (helpful even at
# (the forced setting on native Windows / macOS due to dataset pickling # num_workers=0), but page-locked RAM cannot be relocated by NVIDIA's
# constraints). No-op on non-CUDA backends. # Windows driver shared-memory VRAM-overflow fallback, which can cause
if torch.cuda.is_available(): # 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 dataloader_kwargs['pin_memory'] = True
if has_buckets: if has_buckets: