MINI-SEQUENCE TRANSFORMER: Optimizing Intermediate Memory for Long Sequences Training

我们介绍了 Mini-Sequence Transformer（MsT），这是一种简单而有效的方法，可用于高效准确地训练具有极长序列的语言模型。MsT将输入序列分区，并迭代地处理小序列以减少中间内存使用。与激活重计算相结合，它在前向和后向传递中都能显著节省内存。在使用 Llama3-8B 模型进行实验时，我们通过精心的内存优化，使用 MsT，即使使用比标准实现长12倍的序列，也没有观察到吞吐量或收敛方面的降级。MsT是完全通用的，不依赖于具体实现，且需要最少的代码更改即可与现有的语言模型训练框架集成。

论文旨在解决长序列语言模型（LLM）训练中的内存问题，通过MsT方法，实现高效、准确的LLM训练。

MsT方法将输入序列分割成小序列，并迭代处理小序列以减少中间内存使用。结合激活重计算，它在前向和后向传递中实现了显著的内存节省。

论文中的实验表明，MsT方法能够处理比标准实现长12倍的序列，而不会降低吞吐量或收敛速度。MsT方法通用、实现无关，与现有的LLM训练框架集成所需的代码更少。值得关注的是，论文提出的MsT方法可以在其他领域的序列处理中应用。

最近的相关研究包括《Reformer: The Efficient Transformer》、《Longformer: The Long-Document Transformer》等。