Feature Request: NVIDIA GDS Support for PyTorch IterableDataset & Checkpointing

username: karan6181
cooked: <p>Hey PyTorch Team,</p>
<p>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:</p>
<ul>
<li>Lower Data Loading Latency: Faster data transfer to the GPU.</li>
<li>Higher I/O Throughput: Leverages NVMe storage speed for large datasets.</li>
<li>Better Resource Utilization: Frees up CPU resources for other tasks.</li>
</ul>
<p>An example of leveraging GDS with map-style datasets can be seen in the<a href="https://github.com/Project-MONAI/tutorials/blob/main/modules/GDS_dataset.ipynb" rel="noopener nofollow ugc"> MONAI tutorials</a>, 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.</p>
<p>Our own<a href="https://docs.mosaicml.com/projects/streaming/en/stable/index.html" rel="noopener nofollow ugc"> StreamingDataset</a> 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.</p>
<p>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:</p>
<ol>
<li>First-class support for GDS within PyTorch’s IterableDataset.</li>
<li>First-class support for GDS in checkpoint saving and loading, both distributed and non-distributed.</li>
</ol>
<p>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.</p>
<p>Thanks for considering this request!</p>

username: karan6181
cooked: <p>NVIDIA GDS: <a href="https://developer.nvidia.com/blog/gpudirect-storage/" class="inline-onebox" rel="noopener nofollow ugc">GPUDirect Storage: A Direct Path Between Storage and GPU Memory | NVIDIA Technical Blog</a><br>
cu-CIM’s GDS integrations:<a href="https://github.com/rapidsai/cucim" class="inline-onebox" rel="noopener nofollow ugc">GitHub - rapidsai/cucim: cuCIM - RAPIDS GPU-accelerated image processing library</a></p>

username: andrewkho
cooked: <p>Hi Karan, thanks for reaching out! I think this was implied but let me know if I’ve misunderstood:<br>
the problem arises when you want num_workers &gt; 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?</p>
<p>for 2) are you referring to checkpointing with <a href="https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader" rel="noopener nofollow ugc">StatefulDataLoader</a>? What are the gaps here?</p>
<p>Also not sure if you’re on the PyTorch slack but it may be a better place for communication</p>