search

jasperzhong / read-papers-and-code

My paper/code reading notes in Chinese
43 stars 3 forks source link

code reading | Language Datasets and Data Loaders #357

Closed jasperzhong closed 11 months ago

jasperzhong commented 11 months ago

https://github.com/NVIDIA/LDDL/tree/main

jasperzhong commented 11 months ago

非常有用的repo. 像language model pretraining需要大量数据，这些数据基本是保存成一个一个parquet文件. 这么多文件无法直接load进memory，但如果一个一个文件在训练的时候ad-hoc读取，势必造成很大的latency.

LDDL做的事情是提供了一个PyTorch Dataloader，会"streams samples from disk into memory"，使用方法和普通的dataloader一样.

其主要实现是ParquetDataset类，继承PyTorch的IterableDataset.

__init__()需要给定file_paths和buffer size等信息. 首先遍历所有file，得到total num samples. 这里会尝试读取file path下面的".num_samples.json"的文件，里面保存了每个file的num samples信息，因为直接读取parquet文件非常expensive.

这里有一个问题，就是不同file可能samples数量不一样，但data parallel需要每个worker有相同数量的samples. 他们这里这个问题不严重，不同file最多差一个sample，所以直接统一选择minimum num samples作为每个file的samples数量. 这里的file的包含的samples数量算是非常均匀了.

IterableDataset需要实现__iter__()方法. 里面主要是取对应rank的files，以及dataloader还有多个worker，再给每个rank的不同worker分files.

rank_files = files[self._rank::self._world_size]
worker_files = rank_files[worker_rank::num_workers_per_rank]

然后创建一个ShuffleBuffer，输入是worker_files，buffer size和warmup factor. ShuffleBuffer实现了__iter__()的方法，里面有一个for loop读取worker_files然后读取samples，有一个in memory的buffer来保存读取的samples. 首先会填满buffer，等buffer满了就从buffer里面yield一个sample ，然后buffer对应位置replace成新的sample，最后再yield完buffer里面剩下的samples. 这里有一个warmup的做法，就是一开始buffer是空的，等buffer全部完全满了才开始yield sample有一个很长的等待时间，所以就有每填充warmup factor个samples，就yield一个sample，降低等待时间.

jasperzhong commented 11 months ago

总体来讲，我觉得这个解决方案算比较general，单卡多卡都适用.

但LDDL没有解决的问题是：如果不同file的num samples差别数量很大怎么办？那就应该在分配files到不同rank和每个rank的不同dataloader worker的时候，考虑这个问题.