LongAlign：ChatGLM 团队发布的超长文指令数据及训练评估方案

木尧大兄弟

已于 2024-08-24 18:08:33 修改

阅读量1k

点赞数 25

分类专栏： ML&DL&NLP 文章标签：人工智能机器学习深度学习大模型 LLM

于 2024-08-24 18:07:29 首次发布

本文链接：https://blog.csdn.net/muyao987/article/details/140610399

版权

ML&DL&NLP 专栏收录该内容

31 篇文章 11 订阅

订阅专栏

前言

LongAlign: A Recipe for Long Context Alignment of Large Language Models

这是一个由清华 ChatGLM 团队提出的长上下文指令微调数据、模型训练、评测方案一条龙。主要包括：

长文数据集：从九个来源收集长序列，通过 Self-Instruct 用 Claude 2.1 构建的一个长指令遵循数据集，10000条，长度在8k到64k之间，10%中文数据。
长文训练：通过 packing 和 sorted batching 策略加速不同长度分布数据的 sft 效率。
- packing：将序列打包到最大长度后再将它们分派到GPU。然而因为包含不同数量序列的 pack 在最终 loss 计算中被赋予了相等的权重，所以损失平均计算存在偏差。为了减轻这种偏差，我们提出了一种损失加权策略，以平衡不同序列对 loss 的贡献。
- sorted batching：排序批处理方法，将长度相似的序列组合在一起，以减少批内空闲时间。
- 训练了三个模型 （rope调整、继续与训练等过程见实验部分）：
  - ChatGLM3-6B → LongAlign-6B-64k
  - Llama-2-7B → LongAlign-7B-64k
  - Llama-2-13B → LongAlign-13B-64k
长文评测基准：由博士生们标注搞了个 LongBench-Chat 长文基准测试，评估长度为10k-100k的模型能力，包含40个英文任务+10个中文任务。它涵盖了指令遵循能力的多个方面，如推理、编码、摘要和多语言翻译等长上下文。用 GPT-4 标注参考答案，结合一些少样本评分示例给LLM生成的响应打分。

附录里还整了个活，为了测试模型在长文 OOD（Out-Of-Distribution）场景的能力，作者把这篇论文本身送给了模型，“从专业审稿人的角度，指出本文的优点和缺点” 和 “用 Rap 来讲述这篇论文”，作者是懂整活的😂。

Github：https://github.com/THUDM/LongAlign
数据集：https://huggingface.co/datasets/THUDM/LongAlign-10k
论文：https://arxiv.org/abs/2401.18058

介绍

大型语言模型（LLMs）具有大的上下文窗口，有助于完成诸如长文本摘要、长文本和代码问答等任务（Bai et al., 2023a）。现有的构建长上下文LLMs的工作主要集中在上下文扩展上，即位置编码扩展和对长文本的继续预训练（Chen et al., 2023a; Xiong et al., 2023; Peng et al., 2023）。在这项工作中，我们关注长上下文对齐，即通过指令微调LLMs来处理超长Query。然而缺数据、训练效率低、没有一个强大的benchmark，于是有了这篇work。

除了前面简介中的主要工作之外，还有几个有趣发现：

数据量和多样性的影响：长指令数据量和多样性显著影响对齐模型处理长上下文的能力，影响最终性能高达30%。
长指令数据的好处：长指令数据的数量对长上下文任务的性能有积极影响，同时不会损害模型处理短上下文的能力。
训练策略的有效性：采用的打包和排序批处理策略可以加速训练100%以上，而不会牺牲性能。此外，提出的损失加权技术提高了10%长上下文性能。

LongAlign 详情

一. 数据：长文SFT数据集构建

长指令数据通常涉及长篇的上下文材料，如书籍、百科、论文、长代码，并伴随着需要根据材料进行摘要、推理或计算的任务 Query。在构建过程中，我们首先从9个不同的来源收集长文章和文档，然后我们使用Claude 2.1（Anthropic, 2023）根据给定的长上下文生成任务和答案，如图2所示。
在这里插入图片描述 9个来源包括：

Arxiv（学术论文）：开放获取，任何人都可以自由下载。
Books3（书籍）：来自The Pile，目前没有授权下载。
C4 数据集（各类文章）：公开可用的数据集，带有ODC-BY许可。
CLUECorpus2020（各类中文文章）：由Xu等人（2020）从CommonCrawl语料库中提取的中文实例。
CommonCrawl 语料库（各类文章）：公开可用的数据集，任何人都可以自由下载。
Github（代码库）：开放获取，任何人都可以自由下载。
Stack Exchange（问答网站）：可自由下载，采用CC BY-SA许可。
Wikipedia（百科全书）：提供免费访问，采用CC BY-SA许可。
WuDaoCorpora（各类文章）：开放获取的数据集。

为了促进生成任务的多样性，我们将任务类型描述纳入提示中，例如用于总结、信息提取、推理等任务。通过这种方法，我们为10k长文本创建了任务和答案，总共产生了10k个监督数据实例，其中10%是中文。这些数据采样时长度范围从8k到64k，长度用中文字符的更高压缩率的ChatGLM3-6B tokenizer计算出来的。而且对更长的文章进行了上采样（采的比例高一些），以确保长文数据平衡。

用 Claude 构造数据时用的 Prompt（基本都是根据上文提出5个xxx问题）：
在这里插入图片描述

二. 训练：长文训练策略（打包和排序）

已知 LLM SFT 在训练期间，指令和响应通常被串联形成一个序列 $[x, y]$ ，然后通过自回归语言模型 $π$ 处理，以最大化概率 $P_{\pi}(y|x)$ 。损失类似于语言建模损失，但只计算与y（目标标记）相关的损失：

$-\sum_{i=1}^{|y|} \log P_{\pi}(y_i | [x, y]_{<i})$

为确保模型在监督微调（SFT）后仍能处理长文本和短文本（一般能力），我们将长指令数据与一般指令数据集混合进行训练。大量短指令数据与相对较少的长指令数据的混合导致长尾数据长度分布。如图所示，大多数数据属于0-8k长度范围，而其余数据在8k-64k长度区间。如图，这导致同一个batch有的巨长，有的巨短，就会导致相当长的空闲时间。所以本文探索了两种训练方式（packing 和 sorted batching）。
在这里插入图片描述

策略1：Packing

将不同长度的数据串联在一起，直到达到最大长度。得到的打包数据，其长度通常接近最大长度（比如128k），然后被批量处理并在多GPU上进行处理。如图所示：

在这里插入图片描述
此外，为了防止在自注意力计算期间同一个包内不同序列之间的交叉污染，我们传递一个包含不同序列的起始和结束位置的列表，并使用 FlashAttention 2 中的flash_attn_varlen_func函数，它支持块对角注意力的高效计算（见附录B了解更多细节），在计算过程中，每个序列的 Query 只能关注同一序列内的 Key。与使用2D注意力掩码的传统方法相比，它需要更少的计算和IO时间。

为了实现损失加权策略，我们首先预处理训练数据，为每个序列包生成一个加权的一维掩码，其中权重在当前序列中目标标记的位置设置为1/N（N是当前序列中目标标记的数量），否则为0。

在训练过程中，我们根据批次的配置动态设置M和K，即当前整个 batch 中的全部序列数 M 和 Pack数 K。然后，损失计算为每个标记的交叉熵损失之和，乘以 K/MN 进行缩放。

具体来说，考虑将 M 个序列打包到 K 个 pack 的batch 中，其中第i个pack由索引在 $P_{i-1}, P_i )$ 范围内的序列组成，因此有 $P_0 = 1$ ， $P_K = M + 1$ 。具体看看论文中的几个公式啥意思：

公式(2) 是说正常没 pack 时，M个序列中每个序列的损失 $L_i$ 是平均计算的，每个序列对最终损失 $L$ 的贡献是相等的（权重都是 $\frac{1}{M}$ ）：

$\frac{1}{M} \sum_{i=1}^{M} L_i$
公式(3) 更进一步，使用 pack 训练策略，但不考虑序列中 label token 数量的不同。这里出现了一个问题，即不同 pack 中的序列因为包含的 label token 数量不同，按说对最终损失的贡献得是不同才对，所以这里会导致损失计算出现偏差，导致最终跟正确的 $L$ 不一致：
为了解决这种不一致，通过引入损失加权策略来平衡每个序列对损失的贡献，确保最终损失能够反映所有序列的平均损失。具体实现是公式(4) ，通过乘以权重 $K/(N_iM)$ 来调整每个序列的损失，其中 $N_i$ 是第 $i$ 个序列中需计算loss的 label token 的数量。

策略2：Sorted Batching

为确保每个 batch 中的序列长度相似，我们按长度对数据进行排序，并为每个 batch 选择一个随机连续的数据组，不重复。

在这里插入图片描述
然而，这种策略不可避免地在不同 batch 中引入了数据分布的偏差（也就是当前batch全都很长或很短）。这可能对SGD优化器造成灾难性影响。

实验中，我们观察到排序批处理显著加速训练，而对性能没有明显负面影响。这可能归因于我们使用大的梯度累积Step和优化器的强大适应性。

三. 评测：评测基准 LongBench-Chat

包含50个长文本真实世界 Query，长度从10k到100k不等，共40个英文任务+10个中文任务，涵盖各种关键用户密集型场景，如文档问答、摘要和编码。

为确保评估真正反映模型遵循长文本指令的能力，我们避免使用可能在预训练期间被模型看到并记忆的流行长文本。我们还避免使用模型在不阅读长文本的情况下就能直接回答的问题。

在评估中，我们采用GPT-4（OpenAI, 2023b）根据给定的人工标注的参考答案和每个问题的少样本评分示例来评分模型的响应。为确保评估器能够仅根据真值和少样本评分示例做出明智的判断，我们避免过于开放式的问题，例如“根据前面的文本写一首诗”。

怎么证明GPT-4打分靠谱？表1展示了GPT-4使用零样本提示直接打分、使用额外的少样本打分、与人类打分相关性，而且发现，使用少样本提示，GPT-4与人工注释的相关性不仅一致，而且超过了注释者之间的一致性水平，证明了这种度量在LongBench-Chat上的可靠性
在这里插入图片描述

图1 可以看到，还得是闭源模型更胜一筹（GLM-4吊打 Claude? 哈哈），而且上下文长度为32k或更少的模型在 LongBench-Chat 上表现不佳，表明更长的上下文窗口对于完成这些长任务是必要的。
在这里插入图片描述

实验

实验部分主要研究以下几个有趣的问题：

Q1：长指令数量和多样性如何影响模型性能？
Q2：长指令是否会影响通用能力和短指令的遵循和对话能力？
Q3：packing 和 Sorted Batching 对模型训练效率和性能有啥影响？
其他问题：
- 这波策略能扩展到更大参数（13B）模型吗？更扩展到更长上下文（128k）吗？答案是能，见上面图1的结果。
- 长指令和短指令学习趋势有啥区别？答案是趋势非常一致，如图都显示出在0-500步之间快速提升，随后缓慢上升，并在1000步后趋于稳定。所以两者相辅相成。

实验设置如下（关键）：

数据：短指令用的 ShareGPT，把 assistant content 为空的过滤掉。长指令用的 LongAlign-0k（没有长指令的baseline）、LongAlign5k 和 LongAlign-10k，以及 LongAlpaca-12k（9k长QA+3k短QA，其中长QA数据仅基于学术论文和书籍生成多样性较少）
模型：ChatGLM3-6B、Llama-2-7B 和 Llama2-13B。鉴于它们8k和4k的上下文窗口，我们首先把它们的上下文窗口扩展到64k，分别得到ChatGLM3-6B-64k、Llama-2-7B-64k 和 Llama2-13B-64k。这涉及将RoPE位置编码的基频率 base b（Su et al., 2024）扩大200倍（从 10,000 到 2,000,000） 并在长度不超过64k的预训练数据上继续预训练，总共100亿个token（按每条64k，大约15.6w条）。
训练：8xA800 80G GPU + DeepSpeed+ZeRO3+CPU offload。模型可以在不超过64k token的最大长度下训练，而不会OOM。因此，我们将训练数据的最大长度设置为64k，超出此长度的数据将从右侧截断。对于packing训练，每个pack平均包含12条指令数据，我们将总 batch size 设置为 8，从而得到 96 的 global batch size。为了公平比较，我们将非pack训练方法的其他batch size 设置为8，梯度累积步数为12。我们在训练数据上训练2个epcoh（大约1500-2000步）。
评估：我们在评估中包括了长文本任务和短文本任务（指令遵循、对话、通用能力）。
- 对于长文本任务，我们使用我们提出的 LongBench-Chat 来评估模型的长文本对齐能力，并使用 LongBench（也是之前清华自家提出的）来测试模型的一般长文本理解能力。LongBench是一个双语、多任务长文本基准。我们进行三种类型任务的评估：单文档QA、多文档QA和摘要。由于对齐模型通常产生更长的响应，我们没有使用原始指标（ROUGE、F1）来评估模型的回复，而是使用GPT-4根据LongBench上的地面真实答案来评分模型的输出。
- 对于短文本任务，我们使用MT-Bench（Zheng et al., 2023），一个多轮聊天基准，来衡量模型遵循短指令的能力。我们还评估了Open LLM Leaderboard（Beeching et al., 2023）上的一般任务，包括ARC（Clark et al., 2018）、HellaSwag（Zellers et al., 2019）、Truthful QA（Lin et al., 2022）和MMLU（Hendrycks et al., 2021）。我们遵循Open LLM Leaderboard中的评估设置，并使用lmevaluation-harness框架（Gao et al., 2023）对这些任务进行评估。为确保评估结果最稳定，我们使用GPT-4在LongBench-Chat和MT-Bench上两次评分，并取平均值作为最终得分。

Q1：长指令数量和多样性如何影响模型性能？+ Q2：长指令是否会影响通用能力和短指令的遵循和对话能力？

在这里插入图片描述
表2是 ChatGLM3-6B-64k 使用 ShareGPT 短指令配上不同长度的长指令的SFT实验结果（除LongAlign-0k外，所有模型都使用更高效的打包策略和损失加权）。

LongBench-Chat和MTBench结果是GPT-4在所有测试实例中的平均评分（1-10），而其他数据集的结果在0-100之间进行了归一化。

在这里插入图片描述
图4是这4种不同配比训出来的模型大海捞针实验。

结论：

更多的长指令数据增强了长任务中的性能，并且没有影响短任务的性能。
- LongAlign-0k在长任务中的较差性能（特别是在LongBench-Chat上 3.73）表明仅对基础模型进行上下文扩展是不足以确保在下游长任务中表现良好的；
- 捞针结果还表明，更多的长数据增强了模型利用长文本中不同位置信息的能力，导致模型检索错误减少；
长指令数据的多样性有助于模型的指令遵循能力。
- LongAlign-10k 在长指令和短指令遵循任务（LongBench-Chat和MTBench）上的结果显示出比LongAlpaca-12k明显更好的结果。
- LongAlpaca-12k 在LongBench上3个任务略微优于LongAlign-10k。这主要是因为它在 2WikiMQA 和 NarrativeQA 数据集上表现更优，这些数据集基于Wikipedia和小说，与LongAlpaca-12k中的指令数据来源更为相似。

Q3：packing 和 Sorted Batching 对模型训练效率和性能有啥影响？

图5是每种方法所需的训练时间比较。
在这里插入图片描述
表3是在各种训练方式模型在下游任务上的性能。

结论：

打包和排序批处理提高了训练效率，同时保持良好的性能。 从图5可以看出，打包和排序批处理的训练效率相当，都不到朴素批处理所需时间的一半。此外，根据表3，使用这两种高效方法训练的模型在长任务和短任务上的表现与使用朴素批处理训练的模型相当。我们还发现，这两种训练方法的有效性因不同模型而异。
损失加权显著提高了打包训练在长指令任务上的性能。 通过比较在打包训练期间有无损失加权策略的模型性能，很明显地看出，在LongBench-Chat上，采用损失加权策略显著提高了能力（大约提高了5%到10%），而对其他任务的性能影响较小且变化不定。我们认为，这主要是因为在没有损失加权的情况下，不同的长指令数据对损失的贡献不同——更长的数据倾向于对损失贡献更多（参见公式3）。这种不自然的权重偏差通常对模型训练不利，可能导致训练不稳定，使其偏离最优学习轨迹。

总结

本文旨在为长上下文对齐找到最佳实践，涵盖数据、训练方法和评估等方面。我们提出的解决方案，即LongAlign，使用Self-Instruct构建多样化的长指令数据集，并结合打包和损失加权或排序批处理策略高效地微调模型。此外，我们引入了LongBench-Chat，以促进对大型语言模型在实际长上下文交互中指令遵循能力的可靠评估。通过对照实验，我们发现数据的数量、多样性以及正确的训练方法对最终性能至关重要。

我们通过LongAlign方法的实验结果表明，该方法能够有效地对齐模型以处理长达64k令牌的上下文，同时在一般任务上保持性能不下降。此外，我们还发现长指令数据的数量和多样性对模型处理长上下文的能力有显著影响，能够提升最终性能达30%。训练策略，包括打包和排序批处理，能够提高训练效率超过100%，而提出的损失加权技术能够进一步提升长上下文性能10%。

尽管LongAlign在长上下文对齐方面取得了显著进展，但仍存在一些局限性。例如，我们尚未探索LongAlign在更大模型规模或更长上下文窗口上的可扩展性，也未能进行深入的人类评估研究来验证GPT-4评估器在LongBench-Chat上的性能。未来的工作可以包括对LongAlign进行更广泛的评估，探索其在更大规模数据和模型上的应用，并深入研究训练过程中的学习动态。