Efficiently Training 7B LLM with 1 Million Sequence Length on 8 GPUs

现如今，使用更长的上下文长度训练大型语言模型（LLMs）以促进更具创意的应用已经成为了趋势。然而，长上下文训练考虑到GPU内存的限制面临着巨大的挑战。这不仅会导致训练期间大量的激活内存消耗，还会造成相当大的内存碎片。为了促进长上下文的训练，现有的框架采用了重新计算和各种形式的并行策略等策略。然而，这些技术依赖于冗余计算或大量通信，导致模型FLOPS利用率低。在本文中，我们提出了MEMO，一种新的LLM训练框架，旨在进行精细的激活内存管理。考虑到使用FlashAttention时计算的二次扩展和内存的线性扩展，我们在每个层的前向传递后将消耗内存的激活转移到CPU内存，并在反向传递期间获取它们。为了最大化激活的交换而不妨碍计算，并避免耗尽有限的CPU内存，我们实现了一种基于令牌的激活重新计算和交换机制。此外，我们通过采用双层混合整数规划（MIP）方法来解决内存碎片问题，优化跨Transformer层的内存重用。实证结果表明，MEMO相对于Megatron-LM和DeepSpeed平均实现了2.42倍和2.26倍的模型FLOPS利用率。这种改进归因于MEMO最小化内存碎片、减少重新计算和密集通信，并绕过由于碎片化而导致的内存重组过程的延迟的能力。通过利用精细的激活内存管理，MEMO在仅使用8个A800 GPU的情况下，以1百万序列长度高效地训练了7B LLM，实现了52.30%的MFU。

本论文旨在解决长序列训练大语言模型时GPU内存限制的问题，提出了一种fine-grained activation memory management的框架MEMO。

MEMO框架通过在每个层的前向传递后将占用内存较多的激活值转移到CPU内存中，并在反向传递时重新获取，以实现长序列训练。同时，利用token-wise激活值重计算和交换机制最大化交换激活值，避免CPU内存不足，并通过双层混合整数规划(MIP)方法解决内存碎片化问题。

MEMO框架相比于Megatron-LM和DeepSpeed等框架，平均MFU提高了2.42倍和2.26倍，实现了7B LLM在仅使用8个A800 GPU上训练1百万序列长度，MFU达到52.30%。此外，论文还探讨了其他一些值得关注的问题，如内存碎片化问题和交换机制的影响。

近期的相关研究包括通过增加GPU内存容量、采用多卡并行和重新计算等方法解决长序列训练的问题。相关论文有Scaling Laws for Neural Language Models、Megatron-LM和DeepSpeed。