北极长序列训练：面向数百万token序列的可扩展高效训练
机器学习
2025-06-18 v1
摘要
长序列对于RAG、长文档摘要、多模态等应用至关重要，现代LLM（如Llama 4 Scout）支持最长可达1000万token的序列长度。然而，在企业实验室之外，长序列训练对AI社区来说颇具挑战，开源领域的系统支持有限。开箱即用时，即使在现代NVIDIA H100 80GB GPU集群上，使用基本Hugging Face（HF）模型训练序列长度超过32K的Llama 8B模型也会因两个原因而内存不足：i) LLM训练工作负载未针对充分利用单个GPU内存进行优化，ii) 现有的利用多GPU内存的解决方案不易用于HF模型，使得长序列训练难以实现。我们通过北极长序列训练（ALST）解决了这一问题。它结合了与注意力无关的单GPU和多GPU内存优化，使其能够支持多种HF模型的数百万序列长度的开箱即用训练。ALST支持在单个H100 GPU上训练Meta的Llama 8B模型，序列长度为500K；在单个8xH100 GPU节点上训练3.7M；在4节点集群上训练超过15M，与后者的32K基线相比，增加了400倍以上。ALST与HF模型完全兼容，并通过Deepspeed https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-pallellism/ 和Arctic Training https://github.com/snowflakedb/ArcticTraining/blob/main/projects/sequence-parallelism/README.md 开源。
引用
@article{arxiv.2506.13996,
title = {Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences},
author = {Stas Bekman and Samyam Rajbhandari and Michael Wyatt and Jeff Rasley and Tunji Ruwase and Zhewei Yao and Aurick Qiao and Yuxiong He},
journal= {arXiv preprint arXiv:2506.13996},
year = {2025}
}
备注
19 pages, 13 figures