HyperOffload: 图驱动的分层内存管理让大模型突破显存限制

原文链接: arXiv:2602.00748 PDF

摘要

随着大语言模型(LLM)向长上下文推理和稀疏架构演进,内存需求已远超单个设备 HBM 的容量。新兴的 SuperNode 架构通过高带宽互连提供 TB 级共享内存池,但现有软件栈无法有效利用这一硬件能力。本文介绍的 HyperOffload 框架采用编译器辅助方法,通过图驱动的内存管理将远程内存访问作为计算图中的显式操作,在 MindSpore 框架中的实现表明:该方法可将推理峰值设备内存使用降低高达 26%,同时保持端到端性能。

1. 问题定义

大语言模型正面临严重的”内存墙”问题。在推理场景中,多级 KV 缓存主导内存消耗;在训练场景中,峰值激活内存经常触发 OOM 错误。

“内存容量已成为模型可扩展性和系统性能的首要约束”

现有技术方案如 ZeRO、激活重计算等主要在运行时层操作,存在两个根本性限制:

  1. 运行时缺乏全局视野:调度决策基于瞬时内存压力被动做出,无法主动预取数据
  2. 执行顺序控制有限:数据传输常与计算争夺资源,通信延迟无法被有效隐藏

即使拥有 TB 级共享内存和高带宽链路,当前系统仍遭受流水线停顿,无法接近硬件暗示的性能极限。

2. 方法框架

HyperOffload 的核心洞察是:远程内存访问必须从运行时机制提升为计算图中的头等概念

方法框架对比 图 1:运行时驱动 vs 编译器驱动的内存调度对比(来源:原文 Figure 1)

2.1 核心创新

算子化的远程缓存管理

HyperOffload 在 MindIR 计算图中引入了一等缓存算子:

  • Prefetch 算子:异步启动从远程内存到设备内存的数据传输
  • Store 算子:将数据安全释放回远程存储
  • Detach 算子:标记可释放设备驻留的张量

“通过将内存移动从隐式副作用转换为可调度单元,编译器能够联合推理计算和通信”

图驱动的执行顺序优化

计算图中无数据依赖的算子可以任意顺序执行,但这种非确定性可能导致低效模式。HyperOffload 提出了一种基于成本模型的执行顺序优化算法:

  1. 分析每个缓存算子的可行执行位置
  2. 评估每个位置的传输完成时间和与计算的重叠程度
  3. 选择最优位置以最小化延迟和内存占用

执行顺序优化 图 4:不同执行顺序对通信重叠和峰值内存的影响(来源:原文 Figure 4)

2.2 技术优势

全局内存生命周期可见性

算子化使编译器能够全面追踪数据何时产生、消耗、卸载和重新加载,这是纯运行时方法根本无法实现的。

可预测的内存管理

在静态计算图中,内存分配和释放点在编译时定义,消除了运行时不确定性驱动的保守保留。

系统性 JIT 优化

编译器自动流水线化通信算子和计算算子,用户无需手动插入低级同步原语。

2.3 SuperNode 架构支持

SuperNode 架构 图 2:SuperNode 硬件架构对比传统 GPU(来源:原文 Figure 2)

SuperNode 架构的关键特征是大规模共享内存池,通过超高带宽互连直接连接多个 NPU/GPU 设备。在华为 CloudMatrix384 架构中,一个 SuperNode 由 384 个 Ascend 910 NPU 和 192 个 CPU 通过统一总线互连,使整个 SuperNode 作为单一逻辑实体运行。

3. 实验结果

3.1 实验设置

平台:Ascend 910C NPU 平台,单节点 8 卡配置

模型

  • 训练:LLaMA-8B、DeepSeek-V3
  • 推理:DeepSeek-V3 + NSA(稀疏注意力)

基线:原生 MindSpore 框架,未启用分层内存或执行顺序优化

3.2 训练场景结果

LLaMA-8B 训练性能

配置 DP/TP/PP Batch GBS 重计算 端到端成本
No.1 8/1/1 2 16 开启 8000ms+
No.2 2/2/2 1 16 关闭 5200ms

在不同 D2H 带宽下的性能表现:

D2H 带宽 性能提升
33.6 GB/s 与基线相当
40-70 GB/s +5.7% ~ +21.5%

DeepSeek-V3 训练性能

DP/TP/PP/EP Batch Size GBS 重计算 端到端成本
2/2/2/4 1 16 禁用 2500ms

启用分层内存后(8/1/1/4 配置):

  • 端到端训练延迟降低 2% ~ 12.3%
  • 由于计算密度更高,通信开销更容易被计算隐藏

3.3 推理场景结果

KV 缓存卸载效果

配置 基线 分层内存 相对变化
峰值设备内存 (GB) 61.2 45.0 ~-26%
最大序列长度 (tokens) 71k 123k ~1.73×

长序列推理性能

指标 基线 分层内存 变化
碎片整理事件 57 0 消除
Prefill 延迟 129.33 s 99.41 s -23.13%
端到端延迟 187.21 s 161.41 s -13.78%

短序列推理性能

阶段 基线 分层内存 相对变化
Prefill 延迟 (s) 62.19 62.49 -0.48%
Decode 延迟 (s) 0.117 0.146 -25.47%
端到端延迟 177.373 177.109 +0.15%

3.4 稀疏块粒度敏感性

指标 基线 分层内存 相对变化
峰值内存使用 58428M 45828M 21.57%
Prefill 预测时间 (s) 120.098 115.186 4.09%
Decode 预测时间 (s) 0.117 0.146 -25.47%
总时间 (s) 177.373 177.109 0.15%

当稀疏块大小显著增加时,Decode 阶段的 CPU 计算和内存拷贝开销明显上升,导致性能进一步下降。

4. 优点与局限

优点

  1. 编译器级全局优化:将远程内存操作提升为计算图算子,实现编译时全局分析和调度
  2. 确定性延迟隐藏:通过执行顺序优化,静态调度数据传输以隐藏在计算密集型区域之后
  3. 非侵入式集成:用户无需修改模型代码,通过 JIT 图重写自动插入缓存算子
  4. 显著内存节省:推理峰值设备内存降低 26%,最大序列长度提升 1.73 倍
  5. 性能稳定:在不同带宽条件下保持稳定的性能表现

局限

  1. Decode 阶段开销:稀疏块粒度较大时,Decode 延迟增加约 25%(但端到端影响<0.2%)
  2. 短生命周期张量:生命周期极短或细粒度访问模式的张量不适合远程缓存
  3. 带宽依赖:性能提升幅度依赖于 D2H 带宽,低带宽条件下提升有限
  4. 框架特定:当前实现在 MindSpore 框架内,需要适配其他框架

5. 总结

HyperOffload 展示了一种实用的编译器辅助分层内存管理方法,通过将远程内存操作作为计算图中的一等算子,实现了确定性内存规划和有效的通信 - 计算重叠。

“将内存增强硬件集成到编译器优化框架中,对于扩展下一代 AI 工作负载至关重要”

实验结果表明,该方法在不修改用户模型的前提下,显著降低了设备内存消耗,提升了长上下文推理能力,并在训练和推理工作负载中提供了稳定的性能表现。这证明了图级、编译器感知的内存管理是充分利用新兴 SuperNode 架构的实用且有效的基础。

参考文献

  1. Fangxin Liu, et al. HyperOffload: Graph-Driven Hierarchical Memory Management for Large Language Models on SuperNode Architectures. arXiv:2602.00748, 2026.
  2. Samyam Rajbhandari, et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20, 2020.
  3. Aaron Grattafiori, et al. The Llama 3 Herd of Models. arXiv:2407.21783, 2024.
  4. Aixin Liu, et al. DeepSeek-V3 Technical Report. arXiv:2412.19437, 2024.
  5. Pengfei Zuo, et al. Serving Large Language Models on Huawei CloudMatrix384. arXiv:2506.12708, 2025.

本文基于论文 HyperOffload: Graph-Driven Hierarchical Memory Management for Large Language Models on SuperNode Architectures 生成,论文作者来自上海交通大学和华为技术有限公司。