Potential memory leak when num_workers > 0 and gc_handler enabled

[thuliu-yt16]

I am running a task with high I/O demand. So I modified the ParallelAwareDataloader to enable num_workers > 0. As the training went, I observed increasing memory used and finally the task crashed after some dataloader workers are killed due to OOM.

I tried to delete the gc_handler and just let gc do as it was. The memory leak disappeared but the training slows down. From #148

[thuliu-yt16]

I guess it is because the worker processes are forked from the training process which disables gc. But during training, the gc_handler's collect operation does not affect the worker processes and left the latter never collected. Just FYI, I have found two ways to somewhat fix the dataloader worker memory leak.

Either

1. manually call gc.collect() at the end of dataloader's collate_fn.
   or
2. call gc.enable() in worker_init_fn.

[rakkit]

You mean CPU memory OOM? If so, setting num_workers to a small number is all you need.

We saw this before (num_workers>0 and parquet/pyarrow data). This is kind of expected and not the fault of torchtitan.
Each of your workers will open a minimal "block" to speed up data access. The code logic from HF and PyArrow is that it does not release the memory (of opened blocks) until you have finished reading data from the entire parquet shard. And in the LLM training case, with multiple workers, there is none of the shard that will run out until you use all tokens in the dataset, meaning, e.g. >15T tokens if you are using finetune.

Instead, if you set e.g. num_workers=0/1 (which does not slow down training too much), you can finish one shard of data very quickly, and release this memory for next shard.

You can cross-check by using lots of num_workers + jsonl metadata

[thuliu-yt16]

Hi rakkit, yes I mean cpu memory OOM. I agree with you on the LLM training context but in my case, the training process requires reading lots of data. I have to set many workers to catch up the training speed. I hope this issue can help those who face similar problems.

[rakkit]

The problem comes from multi-worker + parquet/pyarrow data. You could either try multi-worker + JSON or stay parquet, but do pre-process to reshard the parquet and make each shard contains less data that can be used quickly.