预训练数据处理--长度分解

发表于 2024-07-23 更新于 2024-07-24 分类于 CS ， NLP ， LLM 本文字数： 4.9k 阅读时长 ≈ 9 分钟

【本文已在同名微信公众号 / 知乎 / 个人博客linsight.cn 上线】

LLM预训练最重要的工作就是数据的准备，可以说90%的时间都在处理数据。

苹果提出Dataset Decomposition，对数据按长度进行分桶，提升预训练的效率。

预训练数据准备

concat-and-chunk

目前一般的做法是准备好文档数据后，获取对应的token，再通过特殊的分隔token把数据拼接起来，最后按训练窗口大小，切分成相应的长度，这就是concat-and-chunk。

这样的做法存在一些问题：
- 不同的文档拼接在一起，如果在训练的时候没有对注意力进行特殊处理，那就会出现模型在对当前文档数据进行next token prediction的时候，关注到前面拼接的不相关内容。
- 不同文档的拼接训练效果上未必有很大帮助，但是却带来了额外的计算成本，因为attention的复杂度和输入长度的平方相关。
- 多个文档拼接后切分，这样就有很多处于窗口边界的文档被切开，导致模型所见到的实际长度小于平均文档长度，影响了模型的长文本能力。

也有一些方法在concat-and-chunk基础上进行了改进，比如 document masking (DM)、 best-fit sequence packing和in-context pretraining等，后面会一起比较效果。

Variable Sequence Length

针对以上问题，一个自然的想法就是做Dataset Decomposition（DD）。

数据集分解的规则有：
- 把原文档数据集重构成多个bucket，每个bucket内的sequence有不同的长度
- 每个bucket内的sequence都是一个完整文档的subsequence
- 所有bucket内的token没有重复

符合这样规则的分解方法并不唯一，苹果提出一个具体的做法：对于长度为l的文档，按 \(l=2^{i_1}+2^{i_2}+\ldots+2^{i_k}\) 进行长度分解，k个subsequence的长度都是2的幂次方。这就是Variable Sequence Length（VSL）。

比如一个文档的长度为200，那么就分解成128 + 64 + 8三个长度的子序列。不同长度的子序列分别集中到对应长度的bucket中（D_i），如下图

最终各个bucket所包含的token数量分布如下图（b）。可以看到长度为512的序列总token数是最多的。此外长度为2^i的bucket主要也都是来源于长度为\(2^i\leq l<2^{i+1}\) 的文档。

这样的切分方法能够保持切分后的序列长度分布基本和原数据一致。此外，上图（c）给出了concat-and-chunk和VSL的context length分布对比。

实验

基于VSL训练的时候，保证了每个batch中的token数b总不变的。训练时会先从多个bucket中进行采样（假设选中的bucket i），选择当前要训练的长度，然后根据当前选择的长度，计算sample数量 = b / 2^i。

这样保持不同长度下都能有相同token数的做法，能够使得训练时不需要改变超参，比如learning rate等。

使用所用模型基于OpenLM，具体参数如下

另外把RoPE的base frequency从10,000增加到100,000。不同base frequency下的下游任务评测效果如下，使用100k的效果更好。

评测任务分成两部分，regular language modeling benchmarks和long context task。

regular包括：
- Commonsense Reasoning (CSR): PIQA-0-shot，COPA-0-shot，OpenBookQA-10-shots
- Language Understanding (LU): Lambada-OpenAI，Hellaswag-0-shot，Winograd-3-shots，WinoGrande-5-shots
- Reading Comprehension (RC): SQuAD-3-shots，BoolQ-0-shot，CoQA-0-shot
- World Knowledge (WK): Jeopardy-3-shots，ArcEasy-3-shots，ArcChallenge-3-shots，WikiDataQA-3-shots

long context包括：
- Multi-Document Question Answering (MDQA)：从NaturalQuestions-Open里抽的一些问题，多个从Wikipedia抽的文档
- TOEFL：多项选择题
- QuALITY：多项选择题

训练效率

保持batch size = 8 * 8192，用OpenLM-1B/3B/7B分别训练 \(2^6\) 到 \(2^13\) 的窗口长度。

各个模型在不同窗口下的单步训练时间如下图

随着窗口长度增大，每步的training overhead增长的速度也在加快。

对于concat-and-chunk方案，每步所需平均时间是固定的，而对于VSL，每步训练所需的平均时间是不同长度下的加权平均。

按前面分桶获得的长度分布，同样训练8k长度的情况下，VSL可以节省20%的训练时间。

Sequence length bias

这个实验是为了了解使用不同长度的预训练数据对模型的效果有什么影响。

首先是在OpenLM-1B上，固定用不同长度的窗口进行训练，各个模型在评测任务上的标下如下图（a）

结果上，reasoning, language understanding和world knowledge的准确性随着预训练长度的增加呈现出倒U型行为，而阅读理解则从随训练长度增加而变好。这样的结果可能和“与任务相关的数据在预训练语料中的长度分布”相关，比如和阅读理解任务相关的预训练数据，长度往往较长，这也就使得在更大窗口长度上训练出来的模型有更好的阅读理解能力。

上图（b）给出了各个下游任务的长度分布，结果也符合上面的分析。

但是还有一个问题，比如在适中长度数据上训练有收益的regular任务，其收益是来自于这些有合适长度的“数据”内容，还是只要窗口长度合适就行，而不论是什么数据？

对此另外训了两个模型：\(\mathcal{D}_{13\to10}\) 和 \(\mathcal{D}_{7\to10}\)，分别把8k数据拆成8个1k长度，和把8个128个数据拼接成1k长度窗口。结果如上图（c）。

首先 \(\mathcal{D}_{13\to10}\) 比 \(\mathcal{D}_{13}\) 有2.6分的提升，这证明了窗口长度的影响。

而 \(\mathcal{D}_{13\to10}\) 又比 \(\mathcal{D}_{10}\) 的得分低一些，这说明长的文档其内容和段文档与特定任务的惯性性确实有所不同。

此外，\(\mathcal{D}_{7\to10}\) 相比 \(\mathcal{D}_{7}\) 没有提升，这说明拼接并不能缓解长度相关性的问题。

Data mixture

训练数据集中，不同长度的比例如何影响结果呢？为了研究这个问题，这里用了7种不同的数据混合比例做实验（保持总token数相同）。

首先，观察到较小上下文长度的mixture在MDQA上的表现较差，因为MDQA需要长上下文理解的能力。而较大的平均上下文长度也与阅读理解任务的性能呈正相关，但是训练的代价更大。此外，1k-only在regular任务的效果不错，但在长上下文任务则表现较差。

最后Natural分布在regular任务和MDQA任务上的效果都接近最佳，这证明了VSL方法在更大数据集上的可扩展性。

Length-based curriculum

经验上，短文本的训练难度应该更低，而长文本则难度更大。

这个实验使用不同的采样比例，探索长度的curriculum learning。各个采样比例和模型效果如下

由于lr会衰减，这样可能导致采样比例较低的数据学习不足，因此使用了多个epoch来缓解这样的差异影响。

结果显示，“Grow-P2” curriculum在不同的指标上接近最优。

curriculum learning另一个好处是训练的稳定性。《The stability-efficiency dilemma: Investigating sequence length warmup for training gpt models》中就提到长序列在训练的开始阶段会导致比较极端的梯度方差，从而造成训练的不稳定。

Scaling

下图（a）显示，VSL方法在不同的总训练量下，在regular任务上的效果都比更高，而训练效率也大概提升了一倍。

（但是这里没有报告长文本的效果？）

而图（b）则显示在不同规模的模型上，VSL也有稳定收益。

和sota的对比

VSL和document masking (DM)、 best-fit sequence packing、in-context pretraining的效果对比如下

DM相比baseline-8k在regular的效果有提升，但是长文本的能力下降了。如《Effective long-context scaling of foundation models》指出的，即使的不想管的文档拼接，在长文本上也有一定的收益，因此DM直接屏蔽不相关文档的做法对长文本效果有损害。

best-fit sequence packing通过更好地组合文档，在regular任务和长上下文任务都有所提升。

in-context pretraining则是拼接相似文档构建长数据序列，因此长文本效果有提升，不过相似性搜索对资源的消耗比较大。

和以上的方法相比，VSL方法在效果和成本上算是有比较好的平衡。

小结

这些实验探索了窗口长度、文档长度和下游各个任务效果的关系，能给我们做特定任务提升提供一些参考。
通过长度对数据集进行分解的方案在效果上看起来略有提升，主要是能在资源消耗模型训练稳定性上有帮助，在大规模模型上，这应该更为重要。
DeepSeek-V2、LLAMA3.1和苹果的DCLM也都提到了关于长度的curriculum learning，这还是一个值得尝试的方法。

读到这了，来一发点赞收藏关注吧~

博客：http://www.linsight.cn/
知乎：Linsight
微信公众号：Linsight

【往期文章】
- MoE：
MoE模型的前世今生
 DeepSeek-V2和MLA
昆仑万维-SkyworkMoE
成本10w刀的JetMoE
MoE的top-p routing
对MoE模型的一些观察
 从dense到MoE -- sparse upcycling
MoE路由--expert choice routing
- 预训练：
Qwen2技术报告
 Yi技术报告-划重点看细节
 MiniCPM
GLM4报告的一些技术点
 Gemma2
苹果的OpenELM
从Yuan2.0到Yuan2.0-M32
bilibili的index-1.9B
从loss视角理解大模型涌现能力
- 长上下文：
LLM长上下文的问题
 解锁大模型长上下文能力
 大模型推理窗口-从有限到无限大
- 推理加速：
大模型推理加速-投机解码
 大模型推理加速-MEDUSA
- 对齐：
大模型偏好对齐-DPO
大模型偏好对齐-ODPO
大模型偏好对齐-simPO
大模型偏好对齐-IPO
- Transformer：
理解Attention:从起源到MHA,MQA和GQA
LLM的重复生成和ICL
transformer中normalization的二三事
 从代码实现看normalization-到底做了什么
 稀疏注意力计算:sliding window attention
理解LLM位置编码:RoPE
RoPE的远距离衰减
- 大模型算法题：
(1) (2) (3) (4) (5) (6) (7) (8)

Reference

【1】Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum https://arxiv.org/abs/2405.13226