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:
1 parent
67e223b502
commit
0c8793ff47
2 files changed
+14
-5
No files matched your search
@@ -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
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in new issue
Block a user