通过编译器驱动的图计算优化,实现超节点内存高效管理。
HyperOffload: Graph-Driven Hierarchical Memory Management for Large Language Models on SuperNode Architectures
- 用编译器中间表示显式表达远程内存访问操作
- 静态调度数据传输,减少26%峰值显存占用
- 适合大规模语言模型推理场景使用
大型语言模型(LLM)向长上下文推理和稀疏架构演进,其内存需求已远超单个设备的HBM容量。尽管新兴的超节点架构通过高带宽互连提供了太字节级共享内存池,但现有软件栈未能有效利用该硬件。当前基于运行时的卸载与交换技术仅具备局部视角,导致反应式调度和暴露的通信延迟,阻塞计算流水线。本文提出超节点内存管理框架(HyperOffload),采用编译器辅助的图驱动内存管理,将远程内存访问视为计算图中的显式操作,专为分层超节点架构设计。与反应式运行时系统不同,该框架在编译器中间表示(IR)中使用缓存操作符表达数据移动,实现对张量生命周期和执行依赖的全局编译时分析。基于此可见性,我们开发了全局执行顺序优化算法,静态调度数据传输以隐藏远程内存延迟。我们在生产级深度学习框架MindSpore中实现了HyperOffload,添加了远程内存后端和专用编译器阶段。在代表性LLM工作负载上的评估显示,该方法在保持端到端性能的前提下,将峰值设备内存使用降低至多26%。结果表明,将内存增强型硬件集成到编译器优化框架中,是扩展下一代AI工作负载的关键。
原文摘要 · Abstract (English)
The rapid evolution of Large Language Models (LLMs) towards long-context reasoning and sparse architectures has pushed memory requirements far beyond the capacity of individual device HBM. While emerging supernode architectures offer terabyte-scale shared memory pools via high-bandwidth interconnects, existing software stacks fail to exploit this hardware effectively. Current runtime-based offloading and swapping techniques operate with a local view, leading to reactive scheduling and exposed communication latency that stall the computation pipeline. In this paper, we propose the SuperNode Memory Management Framework (\textbf{HyperOffload}). It employs a compiler-assisted approach that leverages graph-driven memory management to treat remote memory access as explicit operations in the computation graph, specifically designed for hierarchical SuperNode architectures. Unlike reactive runtime systems, SuperNode represents data movement using cache operators within the compiler's Intermediate Representation (IR). This design enables a global, compile-time analysis of tensor lifetimes and execution dependencies. Leveraging this visibility, we develop a global execution-order refinement algorithm that statically schedules data transfers to hide remote memory latency behind compute-intensive regions. We implement SuperNode within the production deep learning framework MindSpore, adding a remote memory backend and specialized compiler passes. Evaluation on representative LLM workloads shows that SuperNode reduces peak device memory usage by up to 26\% for inference while maintaining end-to-end performance. Our work demonstrates that integrating memory-augmented hardware into the compiler's optimization framework is essential for scaling next-generation AI workloads.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。