Dataloading is not working when used in litgpt's debug pretraining example

[iloshchilov]

The pretraining example of litgpt with

litgpt pretrain \
   --model_name pythia-14m \
   --config https://raw.githubusercontent.com/Lightning-AI/litgpt/main/config_hub/pretrain/debug.yaml

is doing some data preprocessing which slows-down from >100 it/sec to about 14 it/sec (ca. 40 it/s shown below because computed over few last iterations). Overall, it takes about 1 hour which seems >10x longer than it should be for this small dataset. At the end, it does not complete because some workers still have something to do that they don't do:

Worker 18 is terminating.
Worker 18 is done.████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉| 99996/100000 [57:22<00:00, 41.02it/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉| 99986/100000 [57:23<00:00, 40.94it/s]
Progress: 92%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎ | 45/49 [57:38<05:07, 76.85s/it]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉| 99996/100000 [57:23<00:00, 41.58it/s]

99%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌ | 99024/100000 [57:12<00:24, 40.35it/s]

99%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏| 99429/100000 [57:22<00:13, 42.13it/s]
99%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏| 99434/100000 [57:22<00:13, 42.12it/s]

99%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏| 99449/100000 [57:22<00:13, 42.32it/s]

99%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏| 99459/100000 [57:22<00:12, 42.21it/s]

99%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏| 99479/100000 [57:23<00:12, 40.69it/s]

100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉| 99999/100000 [57:35<00:00, 42.37it/s]

When I relaunch that code, it restarts the whole data preprocessing from scratch. Could you please have a look at it?

[github-actions]

Hi! thanks for your contribution!, great first issue!

[javaarchive]

Keep an eye on your memory usage, and check sudo dmesg for mentions of python. I spent a few hours debugging this just to realize that the OOM killer was killing the process leading to the progress bar stalling after some but not all workers complete.

[tchaton]

Interesting, I had a mechanism to detect whether one of the process died: https://github.com/Lightning-AI/litdata/blob/main/src/litdata/processing/data_processor.py#L1002C1-L1004C1. Maybe, something else is happening there.

I will have a look this week.

[tchaton]

cc @awaelchli