🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

原作者： Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

发布于 2026-05-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教一个巨大且超级聪明的机器人（大型语言模型）如何写作、聊天和推理。为此，你向它展示了一个庞大的图书馆（训练数据）。机器人通过犯错、检查自己的工作，并调整其内部“大脑设置”（参数）来变得更好。

进行这种调整的标准方法就像一位非常谨慎的会计师，名叫Adam。Adam 为每一个大脑设置都保留一本详细的账本，不仅追踪当前的错误，还追踪过去错误的历史，以此决定调整幅度。

问题：会计师过于沉重
问题在于，这位“会计师”（优化器）极其沉重。对于一个拥有数十亿个大脑设置的机器人来说，会计师的账本所占用的内存是机器人“大脑”本身的两倍。

类比：想象你要搬一座房子。家具（模型）很大，但搬家卡车（优化器状态）甚至更大。如果你只有一辆小卡车（有限的计算机内存），无论房子多好，你都根本搬不动它。这就是阻碍研究人员训练更大、更聪明机器人的“内存瓶颈”。

旧方案：走捷径
以往试图缩小搬家卡车的尝试都存在缺陷：

冻结房屋的部分区域：你停止移动某些家具，只调整几个小箱子。这节省了空间，但使房屋变得不够灵活（性能降低）。
拍摄模糊照片：与其移动实际家具，不如拍摄低分辨率的照片来节省空间。但计算这些照片需要大量的时间和能量（计算开销）。
共享账本：你让不同的房间共用同一本笔记本。这节省了空间，但使调整变得不够精确。

新方案：FOAM（智能折叠法）
这篇论文介绍了FOAM（带有近似动量的折叠优化器）。将 FOAM 想象为一位精通打包的大师，它使用巧妙的“折叠”技术，将巨大的搬家卡车塞进一辆微型货车中，而不会丢失任何重要细节。

以下是 FOAM 的工作原理，使用简单的比喻：

1. “块折叠”（压缩）

FOAM 不再单独追踪每一个大脑设置，而是将它们分组为小块（例如 8 个或 16 个设置的网格）。

比喻：想象你有一排 100 个人，你需要记住他们的平均身高。与其写下 100 个独立的数字，不如将他们分成 10 个 10 人的小组。你只需写下一个数字：每个小组的平均身高。
结果：这将“账本”所需的内存大幅缩减（最多减少 90%）。

2. “残差校正”（安全网）

如果只使用平均值，你就会丢失每个人的独特细节。也许小组里有人很高，有人很矮。平均值会掩盖这些差异。

比喻：FOAM 很聪明。在写下“小组平均”后，它会迅速计算出真实人员与平均值之间的差异（误差）。它称之为“残差”。
技巧：它不会永久保留这个差异。它计算出差异，用它来修正那一刻的更新，然后将其丢弃。这就像厨师尝了一口汤，调整了盐量，然后忘记了那一勺的具体味道，而不是记录每一次品尝的精确味道。
重要性：这确保了机器人仍能学习每一个大脑设置的独特细节，即使内存已被压缩。

3. “展开”（恢复）

当需要实际更新机器人的大脑时，FOAM 将这些压缩的“小组平均”重新展开为完整尺寸，并在其上叠加“残差”校正。

结果：机器人获得了精确的更新，就像使用沉重的会计师一样，但计算机只需携带微小的折叠版本。

论文声称（结果）
作者在训练各种机器人模型（从 6000 万参数的小模型到 70 亿参数的大模型）上测试了 FOAM。他们发现：

巨大的内存节省：FOAM 将优化器所需的内存减少了高达90%。总体而言，它将训练的内存占用减少了约50%。这意味着你可以在相同的硬件上训练更大的模型，或者更快地训练相同的模型。
无性能损失：尽管进行了重度压缩，使用 FOAM 训练的模型表现与使用标准沉重会计师训练的模型一样好（有时甚至更好）。它们学习得更快，并达到更高的准确率。
速度：由于无需携带沉重负载，训练过程更快。
通用性：它适用于不同类型的机器人架构（如 LLaMA、Qwen 等），甚至可以与其他节省内存的技巧结合使用。

简而言之
FOAM 就像一个魔法手提箱。它通过将大量信息（优化器的内存）整齐地折叠起来，并添加一个快速的“校正备注”以确保不丢失任何内容，从而将其装入极小的空间。这使得研究人员能够在不需要超级昂贵的大型计算机的情况下，构建更智能、更大的 AI 模型。

技术摘要：FOAM——用于内存高效 LLM 训练的阻塞状态折叠

1. 问题陈述

大型语言模型（LLM）通过增加参数数量和训练数据取得了卓越的性能，但其规模在训练过程中引入了显著的内存瓶颈。事实上的优化器 Adam（及 AdamW）需要为每个参数存储一阶和二阶矩估计，导致优化器状态的内存开销达到模型大小的两倍。这使得大模型（例如 7B+ 参数）的预训练和微调受限于内存，即使使用较小的批量大小，通常也需要数百 GB 的显存。

现有的内存高效解决方案面临权衡：

参数冻结（例如 LoRA）： 将更新限制在低秩子空间，可能限制模型的表达能力和性能。
低秩投影（例如 GaLore、APOLLO）： 使用奇异值分解（SVD）或随机投影压缩优化器状态。虽然有效，但 SVD 引入了巨大的计算开销，而随机投影在严格内存约束下可能会降低性能。
块共享（例如 Adam-Mini）： 在块之间共享二阶矩估计以减少存储，但通常一阶矩开销仍然很高，或将内存减少限制在约 50%。

2. 方法论：FOAM

作者提出了FOAM（带近似矩的折叠优化器），这是一种结构化状态折叠框架，旨在压缩优化器状态而不牺牲优化保真度。FOAM 基于阻塞平均结合残差校正的原则运行。

核心机制

给定时间步 $t$ 的梯度矩阵 $G_t \in \mathbb{R}^{m \times n}$ 和折叠级别 $l$ ：

阻塞平均（折叠）： 将梯度划分为包含 $2^l$ 个连续元素的块。折叠算子 $A^{(l)}$ 计算每个块的均值，将梯度维度从 $n$ 压缩至 $n/2^l$ 。一阶和二阶矩（ $\tilde{M}_t, \tilde{V}_t$ ）基于这些压缩后的梯度进行更新。
展开： 为了将更新应用到完整的参数空间，折叠算子 $E^{(l)}$ 将压缩后的条目复制回原始维度。
残差校正： 直接复制更新会导致块内所有元素共享相同的值，仅依赖初始化进行区分。为了恢复丢失的信息，FOAM 计算残差项 $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ ，代表原始梯度与折叠 - 展开版本之间的误差。
- 一阶矩重构为： $M_t = \tilde{M}_t E^{(l)} + R_t$ 。
- 二阶矩重构为： $V_t = \tilde{V}_t E^{(l)} + R_t^2$ 。
- 关键在于， $R_t$ 是即时计算并立即释放的；它不存储在内存中。

混合应用

FOAM 专门应用于 MLP 和 Attention 模块，而 Embedding 层使用标准 Adam。缩放系数 $\alpha$ 调整 FOAM 与 Adam 模块之间的学习率比率，这是内存高效优化器中的常见做法。

3. 主要贡献

创新的优化器设计： 与 Full-Adam 相比，FOAM 实现了高达90% 的优化器状态内存开销减少以及约50% 的端到端总内存占用减少。与基于投影的方法不同，它不需要投影矩阵，并支持全参数优化。
理论保证： 在标准非凸优化假设（L-平滑、有界方差、有界梯度）下，作者证明了 FOAM 达到了与原始 Adam 相同的 $O(\frac{\log T}{\sqrt{T}})$ 收敛速率。残差误差比率 $\delta_l$ 被限制在 1 以内，且不会对主要收敛速率产生负面影响。
广泛评估： 该方法在模型规模（60M 至 7B）、架构（LLaMA、Qwen、GPT、DeBERTa）和任务（C4 上的预训练、GLUE 和 MMLU 上的微调）方面进行了验证。

4. 实验结果

论文报告了以下实证发现：

预训练性能： 在验证困惑度（PPL）方面，FOAM 始终优于内存高效基线（GaLore、APOLLO、Adam-Mini、Muon），同时保持显著更低的内存使用量。例如，在 LLaMA-1.3B 上，FOAM-3 将总内存消耗减少了约 50%，同时实现了更低的 PPL。
收敛速度： FOAM 表现出更快的收敛速度，通常比其他内存高效方法需要更少的训练迭代次数即可达到目标 PPL。
可扩展性： FOAM 在更大模型（LLaMA-3B、7B）和长序列训练设置（高达 2048 个 token）下保持优越性能。
微调： 在 GLUE 和 MMLU 基准测试上，FOAM 实现了与 Full-Adam 及其他基线相当或更好的性能，证明了其在预训练之外的实用性。
鲁棒性： 消融研究证实，残差项对性能至关重要；移除它会导致收敛变慢和最终 PPL 升高。FOAM 在结合 8 位量化以及与其他优化器（如 Adam-Mini 和 Muon）集成时，也保持鲁棒性。

5. 意义与主张

论文声称，FOAM 为大规模语言模型训练中的内存瓶颈提供了一种实用、高效且与优化器无关的解决方案。通过消除对权重冻结或昂贵投影计算的需求，FOAM 允许在大幅降低内存开销的情况下进行全参数优化。作者将 FOAM 定位为低秩投影技术的补充替代方案，提供了一种在理论上与 Adam 收敛性相匹配，同时在经验上提供卓越内存效率且通常具有更好训练性能的方法。

该研究表明，FOAM 可以在现有硬件上训练更大的模型，或降低训练当前最先进模型的计算成本，从而使高性能 LLM 训练更加易于获取。