来自英伟达和UIUC的华人团队提出了一种高效训练方法，该方法将LLM上下文长度从128K扩展至惊人的400万token，创造了SOTA纪录！基于Llama3.1-Instruct打造的UltraLong-8B模型，在长上下文基准测试中表现卓越，在标准任务中也保持着顶尖竞争力。

大语言模型即LLM，在文本任务方面，展现出了惊艳的表现，在多模态任务方面，同样展现出了惊艳的表现。

像是最新的Gemini 2.5 Pro在文本方面有顶尖性能，在代码方面也有顶尖性能，以及GPT-4o具备原生生图能力，这些都很好地证明了这点。

然而，很多实际应用场景需要模型能够处理超长的token序列，比如长文档和视频理解，比如上下文学习，比如推理时扩展。

在这些场景里，模型的上下文窗口受到限制，这常常成为一个大瓶颈，原因是分布于长文档中的关键信息有可能被忽视。

为解决这些问题，英伟达的研究者提出了一种训练方法，UIUC的研究者也提出了一种训练方法，且这两种方法是高效的 。

这种方法能够从现有的指令微调模型着手，构建超长上下文的语言模型，它最高能够把上下文长度推进到400万token的极限！

论文地址：

研究人员运用上面方法训练出了UltraLong-8B模型，该模型在长上下文任务方面达到了顶尖水准，并且在标准任务上也具备竞争力 。

主要贡献：

实验方法

如图1所示，本文方法主要分成两个阶段，一个阶段是持续预训练，另一个阶段是指令微调。

以Llama 3.1 - 8B - Instruct为基础，持续预训练阶段会逐步扩展模型的上下文窗口至目标长度，比如100万token，200万token，400万token 。之后，指令微调阶段会优化模型的指令遵循能力以及推理能力。

这两个阶段相互结合，使得模型能够高效地处理超长输入，并且在长短上下文任务中有着出色的表现。

第一阶段，通过持续预训练来扩展模型的上下文窗口，采用特殊文档分隔符以及基于YaRN的缩放技术处理超长序列。第二阶段，使用精心挑选的数据集进行指令微调，以此提升模型的指令遵循和推理能力。

持续预训练：扩展上下文长度

在第一阶段，研究者进行持续预训练，借此将 Llama - 3.1 - 8B - Instruct 的上下文窗口扩展到目标长度。

研究者对短文档进行下采样，这些短文档少于4000 token，对长文档进行上采样，这些长文档超过8000 token，最终形成一个包含10亿token的语料库。

这些文档被拼接成更长的序列，该序列对应目标上下文的长度，例如100万、200万、400万token 。在拼接时，他们用特殊字符分隔不同文档，而非使用保留的开始和结束标记。

此外，在持续预训练过程中，研究人员并未使用跨文档注意力掩码，如此一来，便允许模型关注整个输入序列。

研究人员采用了基于YaRN的缩放方法来支持超长上下文，而非之前工作常用的NTK感知缩放策略，他们固定超参数α为1、β为4，还根据目标上下文长度计算缩放因子s 。

当输入长度快要接近最大限制的时候，Llama - 3.1模型的性能会出现下降的情况。为了处理这个问题，他们给RoPE嵌入采用了更大的缩放因子，借此能够更好地适应超长序列。

研究者构建了长上下文模型，该模型针对三种上下文长度，分别是100万token、200万token和400万token，并且将RoPE缩放因子分别设置为128、256和512 。

每个模型都要在10亿token的语料上进行训练，训练时长为一个epoch，其学习率为3×10⁻⁵ 。

他们使用Megatron-LM框架，目的是提升训练的可扩展性。采用了张量并行和上下文并行，用于处理超长输入序列。

训练是在256个NVIDIA H100 GPU上开展的，1M模型的训练时间大约是5小时，2M模型的训练时间大约是6小时，4M模型的训练时间大约是13小时。

指令微调

在第二阶段，研究者提升长上下文模型的指令遵循和推理能力，借助的是监督微调（SFT），他们使用的是精心挑选的一些数据集。

他们对多个开源SFT数据集进行了整合，还进行了优化，这些数据集覆盖三个关键领域，分别是通用领域，数学和代码。

他们想要进一步提升SFT数据集的质量，于是利用GPT - 4o和4o - mini对这些数据集的回答内容进行了优化。

需要注意的是，研究者的SFT数据集仅仅包含上述短上下文数据，也就是少于8000 token的样本，并未加入合成长上下文指令数据。

他们发现，仅依靠短上下文数据，就能够取得优异效果，这和之前研究的观察是一致的。

最终，研究者构建了一个SFT数据集，该数据集包含10万个样本。对于每种目标上下文长度的模型，他们使用的批大小为128，学习率为5×10⁻⁶ 。

训练基于Megatron-LM框架，在256个NVIDIA H100 GPU上开展，张量并行度设置为tp=8，每次训练大概需要30分钟。

基线模型与评估基准

研究者把他们的模型，与基于Llama家族的最先进长上下文模型作对比，目的是确保对训练方法进行公平且可控的评估。

本文研究者专注于Llama家族的模型，如此一来能更清晰地展示其扩展上下文长度训练方法的有效性，并且能确保在标准任务上的性能依旧具备竞争力。

他们通过以下基准测试来评估模型的长上下文能力：

实验结果

研究人员先从名为“大海捞针”（Needle in a Haystack，NIAH）的测试着手，之后再去探讨长上下文的评估，以及标准基准的评估。

研究人员借助NIAH密码检索测试，来评估模型在长上下文检索方面的能力，在这个任务里，模型要在一大段无意义文本中，找出一个简单密码，例如一个随机的六位数字。

他们进行测试，目的是量化检索的准确性，测试了40种不同的输入序列长度，对于每种长度，密码会被随机插入到10个均匀分布的文档深度中。

结果如图2所示 。针对本文模型 ，测试了输入长度高达100万 、200万以及400万个token的情况 ；而对于基准模型 ，仅测试了最高100万个token的情况 。

如图2a到2c所示，在基准模型里，只有Llama-3-8B-Instruct和Gradient-1048k通过了NIAH测试，Llama-3.1-8B-Instruct出现了错误，Llama-3-8B-ProLong-512k-Instruct也出现了错误，即便在它们声称的上下文长度范围内。

相比之下，如图2d到2f所示，研究者的超长模型在所有输入长度上都达到了100%的准确率，在所有深度上也都达到了100%的准确率，展现了强大的长上下文检索能力。

研究者在RULER上的评估结果如表1所示，研究者在LV-Eval上的评估结果如表1所示，研究者在InfiniteBench上的评估结果如表1所示。加粗的数字表示性能超过了所有基准模型。

总体来说，他们的三个模型在大多数情况下都取得了最高分。

在RULER基准测试里，UltraLong模型在输入长度为512K时表现最佳，在输入长度为100万个token时也表现最佳。在LV-Eval中，他们的模型在token长度为128K时平均F1分数最高，在token长度为256K时平均F1分数同样最高。

此外，他们在InfiniteBench上也取得了最佳表现。

这些结果表明，研究者的训练方法是有效的，它扩展了语言模型的上下文窗口，使其能处理超长输入，并且在这个过程中，保持了原有输入长度的性能。

相比之下，在基准模型里，Llama - 3.1是针对128K输入长度设计的，当输入超过128K token时，其性能会显著下降。ProLong是为512K上下文设计的，不过即便它训练了更多token（410亿对比10亿），在512K长度上的表现也比不上他们的模型。

Gradient是基准模型中支持最长上下文的，其最长上下文为100万个token，不过它在LV-Eval和InfiniteBench上的表现较差，这表明它的设计或许过于偏向人工任务，进而牺牲了现实任务的效果。

本文的模型在人工（RULER）基准测试中始终保持更高的分数，在混合（LV-Eval和InfiniteBench）基准测试中也始终保持更高的分数，凸显了方法的高效性和可扩展性。

研究者对模型进行了评估，评估借助通用、数学和代码领域的标准基准测试，这么做是为了确保扩展上下文长度不会对短上下文任务的性能造成影响。

如表2所示，他们的模型性能与基础模型Llama - 3.1 - 8B - Instruct相当，甚至更高，其平均分数分别为62.47、61.06和60.95，而Llama - 3.1 - 8B - Instruct的平均分数为61.45 。

特别值得一提的是，他们的模型在MMLU基准上表现出明显提升，在MATH基准上也表现出明显提升，同时在GSM8K基准上的表现极具竞争力，在HumanEval基准上的表现也极具竞争力。

相比之下，基准长上下文模型Gradient在这些标准任务上的性能大幅下降，其平均分数仅为37.36，基准长上下文模型ProLong在这些标准任务上的性能也大幅下降，其平均分数仅为40.81。

这些结果显示出，研究者的方法起到了有效扩展上下文窗口的作用，并且该方法保持了模型的通用任务能力，甚至还提升了这种能力。

Llama-3-8B-Instruct-Gradient-1048k性能显著下降，Llama-3-8B-ProLong-512k-Instruct性能也显著下降，这表明它们的超长上下文方法或许存在局限性。

结论

在这项工作里，研究人员提出了一种训练方法，该方法高效且系统化，是用于超长上下文语言模型的，它能把上下文窗口扩展到100万、200万和400万个token，并且在标准基准测试中保持了竞争力。

这种结合了高效的持续预训练与指令微调，它提升了模型对长上下文的理解能力，同时还增强了模型遵循指令的能力。

这一框架树立了新标杆，是为可扩展的长上下文建模，它也铺平了道路，是为未来在实际应用中提升长上下文性能的研究。

作者介绍

Chejian Xu

是伊利诺伊大学厄巴纳 - 香槟分校（UIUC）的计算机科学博士研究生，其导师为Bo Li教授。拥有浙江大学计算机科学与技术专业学士学位，就读于CKC荣誉学院，导师是Shouling Ji教授以及Siliang Tang教授。

专注于提升基础模型的安全性，专注于提升基础模型的可靠性，专注于提升基础模型的一致性，这些基础模型包括LLMs，包括多模态模型，还包括基于LLM的智能体。

Wei Ping

是NVIDIA应用深度学习研究团队的资深研究科学家，其专注于大型语言模型的研究，也专注于生成模型的研究。

拥有加州大学欧文分校机器学习博士学位，热衷于构建用于文本的尖端生成模型，热衷于构建用于音频的尖端生成模型，热衷于构建用于多模态数据的尖端生成模型。此前，曾在由吴恩达创立的百度硅谷人工智能实验室担任文本到语音团队负责人。

参考资料：