HyperOffload: 图驱动的分层内存管理用于超节点架构上的大型语言模型
HyperOffload: Graph-Driven Hierarchical Memory Management for Large Language Models on SuperNode Architectures
AI总结:
HyperOffload通过图驱动的编译器方法优化超节点架构上的LLM内存管理,减少峰值内存使用并提升计算效率。
摘要:
大型语言模型(LLMs)的快速发展朝着长上下文推理和稀疏架构方向演进,推动了内存需求远超单个设备HBM的容量。尽管新兴的超节点架构通过高速互连提供了TB级共享内存池,但现有软件栈未能有效利用这一硬件。当前基于运行时的卸载和交换技术采用局部视角,导致反应式调度和暴露的通信延迟,阻碍了计算流水线。在本文中,我们提出了SuperNode内存管理框架(HyperOffload)。它采用编译器辅助方法,利用图驱动的内存管理将远程内存访问视为计算图中的显式操作,专门设计用于分层超节点架构。与反应式运行时系统不同,超节点在编译器的中间表示(IR)中使用缓存操作来表示数据移动。这种设计使得能够对张量生命周期和执行依赖性进行全局、编译时分析。利用这种可见性,我们开发了一个全局执行顺序细化算法,通过静态调度数据传输,将远程内存延迟隐藏在计算密集型区域之后。我们在生产深度学习框架MindSpore中实现了SuperNode,添加了远程内存后端和专门的编译器传递。在代表性LLM工作负载上的评估表明,SuperNode在推理时将峰值设备内存使用量减少了高达26%,同时保持端到端性能。我们的工作证明了将内存增强型硬件整合到编译器优化框架中对于扩展下一代AI工作负载至关重要。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。