Feature Request: NVIDIA GDS Support for PyTorch IterableDataset & Checkpointing

data
1

Hey PyTorch Team,

I’d like to suggest adding first-class support for NVIDIA GPU Direct Storage (GDS) in the IterableDataset class for dataloading, as well as for saving and loading checkpoints. GDS enables direct data transfers between storage and GPU, bypassing the CPU, which offers significant benefits like:

  • Lower Data Loading Latency: Faster data transfer to the GPU.
  • Higher I/O Throughput: Leverages NVMe storage speed for large datasets.
  • Better Resource Utilization: Frees up CPU resources for other tasks.

An example of leveraging GDS with map-style datasets can be seen in the MONAI tutorials, which demonstrate substantial performance improvements. I believe this performance gain could be even more impactful for vision and multimodal workloads, where preprocessing can occur directly on the GPU.

Our own StreamingDataset supports all data types (images, text, video, audio, etc) and is designed to be a drop-in replacement for PyTorch’s IterableDataset. If PyTorch were to offer first-class support for GDS, it could significantly boost data pre-processing and loading efficiency.

With this enhancement, we aim to achieve a ~2x to 7x speed-up in data loading/pre-processing and checkpoint operations. I’m specifically requesting:

  1. First-class support for GDS within PyTorch’s IterableDataset.
  2. First-class support for GDS in checkpoint saving and loading, both distributed and non-distributed.

Additionally, there may be further speed improvements possible with cu-CIM integrations for image processing, depending on the format. For more on this, check out cu-CIM’s rapidsai Github GDS integrations.

Thanks for considering this request!

2

NVIDIA GDS: GPUDirect Storage: A Direct Path Between Storage and GPU Memory | NVIDIA Technical Blog
cu-CIM’s GDS integrations:GitHub - rapidsai/cucim: cuCIM - RAPIDS GPU-accelerated image processing library

3

Hi Karan, thanks for reaching out! I think this was implied but let me know if I’ve misunderstood:
the problem arises when you want num_workers > 0, which leads the IterableDataset to be shoved into a subprocess, and from the subprocess it’s hard to get direct access to the GPU because of how CUDA context is set up, is that correct?

for 2) are you referring to checkpointing with StatefulDataLoader? What are the gaps here?

Also not sure if you’re on the PyTorch slack but it may be a better place for communication