Slow DataLoader Performance in Virtual Cell Challenge Training

[yxwucq]

Hi, I follow the notebook from https://colab.research.google.com/drive/1QKOtYP7bMpdgDJEipDxaJqOchv7oQ-_l#scrollTo=F0vEOd_nmpTR and find that training with the st model is 100x slower than expected. DataLoader takes up to 70 seconds to load a single batch instead of the expected sub-second performance shown in tutorials.

Epoch 0: 100%|██████████| 200/200 [56:16<00:00,  0.06it/s, v_num=0]

Has anyone experienced similar performance issues?

[Griffan]

it seems, for each epoch, the dataset is re-preprocessed. Need confirmation from the developers though.

[Griffan]

There is a preprint https://arxiv.org/abs/2506.01883, "scDataset: Scalable Data Loading for Deep Learning on Large-Scale Single-Cell Omics" claimed to have better performance and support "streaming" mode.

Maybe the developer of this repo can consider adopting it as backend. Specifically, this change can happen within the "data_modules":

cell-load/src/cell_load/data_modules/perturbation_dataloader.py

Line 364 in 1702978

def _create_dataloader(

[zboldyga]

my initial impression walking through this with a debugger is there's a lot of room for optimization

since the dataset for training is fairly small, on linux, you can get around a lot of the performance pains if you have enough system memory that the OS page will keep the h5 file data cached (you can observe with watch -n 1 free -h). Then at least you're just dealing with shuffling that's happening in-memory and not with slow disk reads.

the dataloader is still the bottleneck in training in all hardware I've tested though. @abhinadduri if I happen to create a patch for optimization, would the ARC team accept? (is someone already working on this? any things I need to take into consideration that might not be apparent in the code?)

Not sure on timing. Currently it's not the limiting factor for me in the competition work. But it could be relevant if I expand the dataset size significantly, and would be more economical scale-wise than jumping straight to clusters of GPUs

[zboldyga]

fwiw, for anyone experiencing pain with this in a situation where you have more h5 data than os page cache memory to work with, you can also refactor the h5 files with chunking that is aligned with your batch size. this will probably reduce disk reads

[zboldyga]

I see there are some similar contributions helping with dataloader speed in the discord already, e.g. scplode. will leverage those before doing anything here, will probably circle back eventually if I still feel optimization is worthwhile

[abhinadduri]

hi folks - thanks for engaging and for the discussion. the data loader should not be that slow. I have found that on colab, it's significantly slower than colab+, where it's closer to 3 minutes / epoch.

more generally, @zboldyga is correct in that data loading is by far the bottle neck, and also leads to spiky gpu util. we would appreciate help and would absolutely accept patches after review!

however, I do expect it would be a significant code change. my suspicion is that the core issue lies with reading from h5py with multiple workers.