论文页面 - 压平长上下文混合专家模型训练中的所有内存峰值

来源：https://huggingface.co/papers/2609.14306

摘要

在长上下文或大批量训练混合专家模型时，只要任意一个组件的峰值分配超过设备内存容量，训练就会失败。因此，目标是同时压平每一个峰值，而非仅降低平均内存占用。常见的并行方案中有四项内存开销未被严格限制，且各自增长方式不同：涉及路由矩阵的专家分发、与词表大小和令牌数相关的词表投影、与深度和序列长度相关的梯度检查点边界，以及与参数量相关的优化器状态。随着模型、上下文长度和设备数量的变化，首先耗尽内存的组件也会不同，仅降低某一峰值只会暴露下一个瓶颈。

我们通过以下调度方案将所有四项内存占用限定在固定值，使 GPU 工作集在启动时即确定：Pipelined LLEP 扩展了最少加载专家并行策略，限制了每个源在分发块中贡献的令牌数；Ring-DTP 让激活值或权重分片在环中流转，处理词表投影并将每个 logits 块折算为在线 log-sum-exp 运算；选择性检查点卸载将每个检查点边界中唯一长期存在的张量保留在 CPU 内存中；OffloadStreamAdamW 则将优化器卸载的串行 CPU Adam 更新转为桶式流水线。这四种方案仅改变计算与数据移动的顺序和粒度，因此损失值与梯度保持精确。在组件对照测试中，它们将 MoE 分发峰值降低最多 59.3% 而不损失吞吐量，将词表投影峰值降低 86.6%，并将卸载的优化器步骤加速达 2.05 倍。组合应用于参数量从 120B 到 667B 的 MoE 模型时，可实现 1M 上下文长度的训练，达到调优后 FSDP2 基线方案 8–32 倍的覆盖范围，吞吐量最高可达其 10.4 倍。

查看 arXiv 页面 (https://arxiv.org/abs/2609.14306)查看 PDF (https://arxiv.org/pdf/2609.14306)添加到收藏 (https://huggingface.co/login?next=%2Fpapers%2F2609.14306)

通过代理获取此论文：

hf papers read 2609\.14306

尚未安装最新 CLI？curl \-LsSf https://hf\.co/cli/install\.sh \| bash

引用此论文的模型0

无模型链接此论文

在模型 README.md 中引用 arxiv.org/abs/2609.14306 以从本页链接。

引用此论文的数据集0

无数据集链接此论文

在数据集 README.md 中引用 arxiv.org/abs/2609.14306 以从本页链接。

引用此论文的空间0

无空间链接此论文

在空间 README.md 中引用 arxiv.org/abs/2609.14306 以从本页链接。

包含此论文的合集0

无合集包含此论文

将此论文添加到合集 (https://huggingface.co/new-collection) 以从本页链接。