libai 设计文档之数据加载篇

[dangkai4u]

由于预训练的数据量比较大（数据大小可能大于内存），所以libai在数据加载上要考虑更多高效性。

经翻阅megatron历史版本，发现GLM中的数据处理方法来自于megatron 0.1版本中，从1.0版本到现在，megatron都采用了现在的方法。估计是考虑了高效性，希望减少pad，减少重复的预处理步骤，从而缩短训练时长。因此，libai应该基于megatron设计数据处理。

megatron中的数据处理：

1. 将原始文本预处理为numpy张量，并保存为二进制文件，之后模型读取二进制文件（以bin和idx为后缀文件）。
2. 预处理是单独步骤，在训练开始前完成。调用tools/preprocess_data.py脚本实现，保存为bin和idx文件。该脚本支持多进程处理，且对于大的数据集，如果有多个文件，可以分别处理，分别保存，之后指明文件即可训练。
3. megatron的底层dataset有三种，IndexedDataset，IndexedCacheDataset，MMapIndexedDataset。分别对应lazy，带缓存的lazy，和直接内存映射的。mmap是速度最快的，lazy会增加io通信。但mmap要求所有数据都放进内存，lazy可以只放一部分数据进内存。因此，当训练语料极其大时，lazy是可用的，其他可能无法使用。当语料较大，例如几百万个样本，mmap是比较好的。使用哪种底层，在预处理时通过--dataset-impl来切换。
4. indexed_dataset（上面那三种）存储numpy数组，可以理解为每一条数据对应一个句子经过tokenize和convert tokens to ids后结果。
5. 上层的dataset，就是与具体任务相关，例如bert_dataset，gpt_dataset，t5_dataset。底层的dataset都相同，无需改动，用户只可能改上层的dataset。
6. bert dataset：对于每个文档中的句子，在不超过max_seq_length的前提下尽可能放入多个句子。也就是说，bert dataset的一个样本通常由多个indexed_datset样本组成。新合成的句子由多个句子组成，然后选一个位置（句子的idx），将新合成的句子分成前后两部分，并置换顺序。这其实不是bert的NSP任务，其实是albert提出的SOP任务，不过SOP任务好像比NSP任务更有效。里面的加噪方法也稍微复杂一点，支持wwm mask和span mask，这也是对bert的一种改进。
7. gpt dataset：将所有的文档拼接成一个，按照顺序截取max_seq_length，如果前一个文档结束了，且句子不足max_seq_length，就从下一个文档进行读取。
8. 在6、7中介绍了megatron对于bert和gpt在数据处理上的做法，优点在于，尽可能减少了padding，训练中的无效计算大大降低。同时，这些都是操作numpy数组，没有考虑原始文本，因此执行效率上也较高。缺点在于，上层dataset中的sample idx和底层indexed dataset中的sample idx不一致，因此使用build_sample_idx、build_index_mappings等函数做了id的重映射。这些操作破坏了一些易用性，简洁性，同时，用户自定义数据集的难度也增加了。
9. megatron重写了data sampler，根据数据并行和rank，每个数据并行stage为1的机器获取数据。这样也能减少保存重复的数据。
10. 预训练可能包含多个数据集，megatron提供了blendable dataset，用于混合各数据集，同样由于8中提到的缺点，blendable dataset无法像普通的concat dataset那样简单，需要调用build_blending_indices函数重新映射。另外，也支持对数据集加权，实现降采样或重采样。
11. 对于大语料，直接切分训练集、验证集、测试集也不方便，因此没有split datset这样的类。它通过改变indexed dataset中的doc idx，让训练集、验证集、测试集保存不同的文档，实现这样的切分。
12. 上面提到的idx重映射，也会保存为文件，如果不存在，则重新生成，如果存在，则直接读取。另外，构造idx映射时，有一个多机通信操作，暂时看不懂，可能检查机器间是否同步。
13. megatron没有num_epoch概念，在构建数据集时获取或计算num_epoch和num_samples，然后构造对应的索引范围和对应的数据item。

libai的做法：

1. 底层采用indexed dataset，具有效率优势。
2. 上层希望各dataset尽可能独立，考虑那些build_sample_idx操作，可以分为哪几类，哪些是共有的，哪些是特有的。
3. sampler采用相同的方法，在调用dataloader后，使用to_consistent转化为sbp机制。
4. 混合数据集和切分数据集，应该也是有必要的。由于build_sample_idx操作，貌似很难写出split dataset类，之后再思考思考。
5. 感觉dataset部分很难使用注册机制，实在不行，就在build_dataset函数中对于每一类数据集分别判定并构造。

[dangkai4u]

数据处理

数据集相关类

底层数据集：

• 按照megatron-lm的做法，数据底层以indexed_dataset方式存储。
• 在数据预处理阶段，将原始文本进行分句、分词，并将token转化为id，以二进制文件保存（分别是.bin和.idx文件），其中bin文件存储向量，idx文件存储位置偏移。
• 在训练阶段读取上述二进制文件，每一条数据对应一个句子的张量表示。
• 也支持从一个给定位置，读取指定长度的数据。
• 可以细分为三种类型，lazy，cache，mmap。mmap会将全部数据读入内存中并做映射，lazy会发生较多的io通信，从文件中读取。对于较小的数据量，选择mmap来减少io通信带来的时间开销。对于预训练，如果数据量特别大，lazy可以支持大数据量的读取。
• 使用megatron-lm中的代码即可，无须改动。

中间层数据集(reindexed_dataset)

• 理论上，indexed_dataset中每一条数据可以代表一个句子，直接加mask噪声(bert)，或者移位(gpt)即可。但这样有个缺点。如果设定最大句子长度为512，大多数句子是没有达到这个长度的。一般来说有两种方法，一种做法是直接将句子padding到512，另一种做法是按照桶排序，将一个batch内的所有数据padding到这个batch里最长数据的长度。第一种做法包含了大量的padding，造成很多无效计算。第二种虽然减少了一定的无效计算，但句子长度可能往往达不到512，模型对于长文本的处理能力不足。
• 我们采用一种新的方法，对indexed_dataset的sample idx进行重映射。有两种重映射规则：
  • 在不超过最大句子长度的前提下，将同一个文档内连续的句子合并为一个句子。
  • 从文档读取最大句子长度的数据，如果这个文档读取完毕，就从下一个文档读取。这种情况下，不能保证所有句子都是完整的，可能包含很多断句。
  • megatron中，gpt采用第二种方法，其他采用第一种方法。
• 由于megatron写的比较乱，对于bert_dataset、gpt_dataset、t5_dataset等，采用分别实现或调用build_sample_idx函数，导致代码难以看懂，所以这里增加了一层抽象。在上层数据集，仅考虑任务相关的处理即可。

上层数据集（任务性，如bert_dataset、gpt_dataset、t5_dataset）

• 上层数据集是对中间层数据集的封装，这里仅考虑任务相关的，不再考虑是否如何减少pad。
• 可能的操作：加噪声，移位，pad，截断，list_to_numpy

功能性数据集

• split dataset：将原始数据集划分为训练集、验证集、测试集
• blendable dataset：将各子数据集融合为一个数据集。因为语料来源可能不同，或者预处理时分批进行。
• blendable dataset和concat dataset功能类似，但实现不同。前者通过构建sample idx映射实现，后者通过bisect_right计算求sample idx，我怀疑是不是这个计算可能会花费较多的时间，所以nvidia没有采用bisect_right方式。

采样器

• 根据data_parallel_rank和data_parallel_size设置batch的sample idx，用于采样。

dataloader

• 需要重新构建一个dataloader，这个dataloader可以支持循环取数据，且不drop_last。

2. 数据集处理流程

1. 用户输入：
   • 输入data-path: data-prefix 或 weight1 data-prefix1 weight2 data-prefix2 ...
   • 训练总样本train_samples，训练迭代轮数train_iters，验证间隔eval_interval，验证轮数eval_iters，测试轮数test_iters
   • 训练批次大小global_batch_size，训练，验证、测试划分比例splits
2. 解析：从字符串解析splits，从data-path解析得到weight，data prefix数组
3. build_indexed_dataset
4. build_reindexed_dataset
5. build_bert_dataset
6. 划分训练集、验证集、测试集
7. 对每个数据集循环调用3-6步，之后根据weight进行拼接。
8. 这里的数据集都是一个epoch的数据，对于reindexed_dataset中的两种子情形：
   1. 第一种情形，不会受到影响。构造时，将num_epochs设为1即可。
   2. 第二种情形，会受到影响。原始的实现计算num_samples（来自于输入）个数据，如果一个epoch采样完了，最后一个样本没有达到512长度，就从下一个样本继续读取，直到长度达到512。
   3. 根据num_tokens和seqlen计算得到num_samples，仅构造num_samples个样本。最后一个样本不足512，就丢掉。这里之会丢掉一个样本，所以基本上没影响。

3. idx映射规则

blendable dataset ----> split dataset ----> bert dataset ----> reindexed dataset ----> indexed dataset

sampler产生idx，查blendable dataset中的数组，得到blendable dataset内的索引（哪个数据集的哪条数据）

blendable dataset内的索引，通过split_inds数组，转化为bert dataset的索引

当前不确定的点

• 能否构造出上面描述的dataloader，会自动循环执行，不停地获取数据。

和megatron的区别

• 新增了中间层dataset，split dataset。
• megatron在构造数据集前输入或计算得到num_samples，用于构造数据集。我们仅构造一个epoch对应的dataset，通过train_steps控制停止。
• dataloader需要定制为循环的。
• megatron先划分训练集、验证集、测试集，后构建上层数据集。我们先构建上层数据集，然后通过split dataset划分数据集。

[L1aoXingyu]

我觉得可以先快速从头到尾实现一种 dataset 看看，可以先不用实现那么多功能，先把模块分清楚，然后当用户需要使用自己的数据集的时候方不方便使用

看上去最下面两层都是用户不需要关心的，只有最上层和任务相关的需要用户自己改写

数据处理的部分尽量把大部分工作埋藏在底下，用户只需要写一个简单的处理逻辑就可以加上自己的数据进行训练应该就可以了

[L1aoXingyu]

将 token 转成向量的时候需要一个 vocab->int 的映射吧，这个映射被存在哪里了呀

[dangkai4u]

将 token 转成向量的时候需要一个 vocab->int 的映射吧，这个映射被存在哪里了呀

映射存在了tokenizer，在最开始的预处理阶段完成了转化

[dangkai4u]

我觉得可以先快速从头到尾实现一种 dataset 看看，可以先不用实现那么多功能，先把模块分清楚，然后当用户需要使用自己的数据集的时候方不方便使用

看上去最下面两层都是用户不需要关心的，只有最上层和任务相关的需要用户自己改写

数据处理的部分尽量把大部分工作埋藏在底下，用户只需要写一个简单的处理逻辑就可以加上自己的数据进行训练应该就可以了

我已经尽可能把重复的底层的部分写好了，尽可能让用户只更改上层和任务相关的dataset，数据集部分的完整流水已经实现了，现在sampler和dataloader还需要修改。

[L1aoXingyu]

可以提单独的 pr 上来

[dangkai4u]

分词器是nlp中最基础的操作，是预处理阶段的核心步骤。看起来实现简单，实际上很麻烦，需要考虑很多细节。之前把megatron里的tokenizer搬运过来，里面包含bert和gpt的tokenizer。但这个实现质量较差，甚至接口都没有统一。例如bert里解码是decode函数，gpt里解码是detokenize函数。GLM的实现和megatron相差不多，感觉也不太好。huggingface的tokenizer库实现了很多种tokenizer，考虑了很多细节。我建议以它为基准，在其基础上进行封装，添加预训练相关的特性。你们觉得怎么样？