AutoSP自动优化大模型长文本训练,提升上下文长度。
AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism

- 通过编译器自动应用序列并行与激活检查点技术
- 在NVIDIA/AMD上分别提升2.7倍和2.5倍上下文长度
- 无需手动调优,适合需要长文本训练的研究者
大语言模型在处理包含数万至数十万标记的长上下文任务中展现出巨大潜力。然而,现有训练框架主要聚焦于参数量大的模型优化(如ZeRO-3/FSDP、张量与流水线并行),缺乏对长上下文训练的易用抽象。这迫使用户手动重构训练代码以集成复杂的长上下文优化(如序列并行),需深厚专业知识,严重降低开发效率。为此,我们提出AutoSP:首个自动优化长上下文训练的解决方案。AutoSP通过编译模型,自动应用序列并行与长上下文感知激活检查点技术,在几乎不损失吞吐的前提下,显著提升大模型的可训练性。评估表明,AutoSP在NVIDIA和AMD硬件上分别比手写基线提升2.7倍和2.5倍的训练上下文长度,且运行性能开销极小。
原文摘要 · Abstract (English)
Large-language-models (LLMs) demonstrate enormous utility in long-context tasks which require processing prompts that consist of tens to hundreds of thousands of tokens. However, existing LLM training libraries do not provide easy to use abstractions to optimize for long-context training, instead focusing on optimizations for models with large parameter counts through ZeRO-3/FSDP, Tensor and Pipeline parallelism. This forces users to rewrite LLM training libraries to incorporate compositions of various complex long-context optimizations, such as sequence-parallelism, to training pipelines; a process that requires in-depth expertise, reducing developer productivity. To tackle these challenges, we introduce AutoSP: the first automated solution to automatically optimize LLM training for longer-contexts. AutoSP compiles models and applies a targeted set of optimizations: automated sequence parallelism, and long-context aware activation-checkpointing, to drastically enhance LLM trainability at negligible cost to throughput. Our evaluation demonstrates AutoSP's capability on both NVIDIA and AMD hardware, increasing training contexts by upto 2.7$\times$ and 2.5$\times$ respectively over competitive hand-written baseline at negligible cost to runtime performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。