本文主要探讨了在长序列训练时如何通过精细的激活内存管理来解决长序列训练中的挑战，为此提出了一种名为MEMo的新框架，旨在实现高效的激活内存管理和碎片优化。
1. 问题
由于LLM需要大量的计算资源和存储空间，因此在长序列训练时会遇到很多问题，如GPU内存不足、计算效率低下等。在长上下文训练中，线性增长的激活内存仍然是主要的挑战。例如，在训练一个具有32 layer和4096个隐藏单元的7B GPT模型时，仅使用单个长度为1MB的序列就需要消耗4096GB的内存，这远远超过了GPU显存的容量限制。 目前主要采用recompute和swap方式减少显存占用，其中recompute技术通过选择性存储某些层输入而不是所有中间激活的方法，可以在反向传播过程中重新计算所需的激活， swap方式通过将tensor转移到CPU内存并在需要时将其检索； 但是这两种技术都以时间和空间换取——recompute技术会产生额外的计算开销，而swap技术则需要从CPU内存传输激活到GPU内存，会导致训练速度显著下降。只使用recompute技术会带来显著的额外计算开销，而只使用swap技术可能会导致CPU OOM错误（当序列太长时）或阻塞GPU计算（当交换时间不能完全重叠时）。因此，我们期望精心协调这两种内存节省技术来管理骨架激活，以便在最小化额外开销的同时适应LLM长期训练中的巨大内存需求。为此，本文提出一种基于token的激活重算和交换机制， 通过协调管理这两种技术，以满足在最小化额外开销的同时适应LLM长期训练中的巨大内存需求
2. Memo架构
MEMo是一种针对长序列训练中大规模激活管理的方法，其主要特点在于它采用了精细的激活内存管理策略，通过细粒度激活重算和交换策略来充分利用GPU计算中闲置的PCIe带宽，从而降低激活重构的成本。同时，作者采用一个双层混合整数规划技术来实现对单个Transformer层内内存分配的优化，再利用相同的内存空间来处理每个相同层的数据，避免了内存碎片化的问题；
3. 总结
本文提出一种名为MEMo的LLM训练框架，通过精细的recompute和swap机制来管理激活数据，同时采用一种双层内存规划方法来重复利用瞬态激活的空间。这种方法可以在不增加额外时间和空间成本的情况下减少峰值内存消耗，提高训练效率。