高效内存管理策略

高效内存管理是任何机器学习 (machine learning)运行时系统性能的根本。机器学习模型经常处理大型张量，这需要大量的内存分配。此外，某些模型的动态特性以及推理 (inference)过程中产生的中间激活会导致频繁的分配和释放请求。在性能敏感的循环中，天真地依赖 malloc 或 cudaMalloc 等通用分配器会带来很大的开销和潜在的内存碎片，严重影响执行速度。因此，专门的内存管理策略是必不可少的。

机器学习 (machine learning)内存分配的挑战

机器学习工作负载带来了独特的内存管理挑战：

1. 大额分配： 表示模型权重 (weight)、激活以及输入/输出的张量可能占用数千兆字节的内存。
2. 频繁分配/释放： 计算图中的中间张量通常生命周期较短，导致频繁的内存变动。
3. 性能敏感： 分配/释放延迟会直接增加模型的端到端执行时间。与标准分配器相关的系统调用开销通常无法接受。
4. 异构内存： 系统通常包含不同的内存空间（例如，CPU DRAM、GPU 高带宽内存 (HBM)），它们具有不同的容量、带宽和访问特性。这些空间之间的数据移动代价很高。
5. 动态形状： 张量的确切大小在运行时可能未知，这使得静态内存规划变得复杂。

竞技场分配器和内存池化

在机器学习 (machine learning)运行时中，减轻分配开销最常用的技术是使用竞技场分配器，也称为内存池。其基本思想很简单：

1. 预分配： 运行时在初始化期间或执行特定子图之前，从底层系统预先分配一个或多个大的连续内存块（“竞技场”）（例如，使用 cudaMalloc 用于 GPU 内存，或 mmap/VirtualAlloc 用于 CPU 内存）。
2. 子分配： 后续对张量内存的请求通过从该竞技场中划出块来满足。这涉及到运行时内部的轻量级记录，避免了代价高昂的系统调用。
3. 释放： 当不再需要张量的内存时，竞技场中对应的内存块被标记 (token)为可用，以便未来请求使用。实际内存仍分配给竞技场，直到竞技场本身被销毁或重置。

竞技场分配器通过从预分配的内存块中进行子分配并管理空闲空间来处理请求的视图。

优点：

• 减少开销： 消除了大多数针对单个分配/释放的系统调用开销。
• 提高局部性： 竞技场内的分配通常是物理连续的或在虚拟地址空间中更接近，可能提升缓存性能（尽管大张量的大小通常会主导局部性效应）。
• 碎片控制： 尽管可能发生内部碎片（已分配块中未使用的空间），但外部碎片（已分配块之间不可用的空闲空间）在竞技场边界内进行管理。

实现策略：

• 首次适应/下次适应： 查找下一个可用块的简单策略。长期可能导致碎片。
• 最佳适应： 找到适合请求的最小空闲块。可以减少碎片，但需要搜索空闲列表。
• 分离空闲列表（分离适应）： 为不同大小类别维护单独的空闲列表。特定大小的分配请求首先检查相应的列表。这对于具有常见张量大小的工作负载非常有效，可以显著加快分配速度并减少碎片。
• 伙伴分配器： 以2的幂次大小的块管理内存，简化了已释放块的合并，但可能导致更高的内部碎片。

策略的选择取决于预期的分配模式、内存限制和性能目标。对于动态形状，竞技场可能需要调整大小，或者可能采用具有不同增长策略的多个竞技场。

重新审视静态与动态分配

如第三章（图级优化）所述，静态内存规划会提前分析计算图，以确定张量生命周期并找到缓冲区共享和重用的机会。这能最大限度地减少内存峰值占用。然而，静态规划依赖于预先知道张量形状。

当存在动态形状时，运行时内存管理器必须处理在执行期间确定大小的分配。即使对图中已知部分进行了静态规划，动态部分也仍非常依赖高效的运行时分配。通常，会采用混合方法：静态规划尽可能地优化，而动态竞技场分配器则处理其余部分，包括根据启发式方法或运行时反馈进行的潜在超额分配，以适应动态大小。

显式内存重用与活跃度跟踪

除了竞技场分配器返回已释放块提供的重用之外，运行时还可以实现更积极的显式内存重用。这需要跟踪每个张量缓冲区的活跃度：精确知道缓冲区中的数据何时不再被任何后续操作需要。

一旦缓冲区被识别为“非活跃”，运行时可以立即将其别名为新的分配请求，甚至在生成它的相应操作完全完成之前（前提是同步确保正确性）。这需要与运行时执行调度器（稍后讨论）仔细集成，以正确管理依赖关系。活跃度信息通常由编译器计算，并传递给运行时以指导这些决策。

优化主机-设备传输：锁定内存

在 CPU（主机）和 GPU（设备）内存之间传输数据是一个常见的瓶颈。通过 malloc 分配的标准主机内存通常是可分页的，这意味着操作系统可以移动其物理位置。对于 GPU 用于实现高带宽传输的直接内存访问 (DMA) 引擎，物理地址必须是固定的。

因此，从可分页内存发起传输通常涉及一个中间步骤：GPU 驱动程序将数据从可分页源缓冲区复制到主机 RAM 中一个临时的锁定（或页锁定）缓冲区，该缓冲区的物理地址是固定的。然后 DMA 引擎将数据从这个锁定缓冲区传输到 GPU。这个额外的复制增加了延迟并消耗了带宽。

比较使用可分页主机内存与锁定主机内存的数据传输路径。锁定内存允许直接 DMA，消除了暂存复制。

机器学习 (machine learning)运行时通过将参与 GPU 传输的主机端缓冲区直接分配为锁定内存（例如，使用 cudaMallocHost 或 cudaHostAlloc）来优化此过程。

权衡：

• 性能： 由于直接 DMA 访问，主机到设备和设备到主机的传输速度显著加快。这对于计算和通信的重叠非常重要。
• 资源消耗： 锁定内存是有限的系统资源。过度分配锁定内存会减少操作系统和其他应用程序可用的可分页内存量，从而对整体系统性能产生负面影响。
• 分配开销： 分配锁定内存有时会比分配标准可分页内存略慢。

运行时必须仔细管理锁定内存分配，通常为锁定缓冲区使用专用竞技场，并且仅在传输性能重要时才审慎分配。

统一内存系统

统一内存 (UM) 旨在通过提供一个 CPU 和 GPU 都可以访问的单一、一致的虚拟地址空间，来简化异构系统的编程。程序员只需分配一次内存（例如，使用 cudaMallocManaged），指针即可从任一处理器解引用。

底层系统（GPU 驱动程序、操作系统和硬件）会自动管理物理 CPU DRAM 和 GPU HBM 之间的数据迁移，通常是根据页面错误按需进行。

优点：

• 简化编程： 消除了在不同空间中显式内存分配和手动数据传输 (cudaMemcpy) 的需要。
• 潜在的超额订阅： 允许应用程序分配比 GPU 物理可用内存更多的内存，系统会将数据分页进出。

缺点：

• 迁移开销： 由页面错误触发的自动迁移会引入延迟。如果访问模式导致频繁来回迁移（“抖动”），性能可能无法预测。
• 控制粒度： 与手动管理相比，开发者对数据放置和移动的显式控制较少。性能调优通常涉及使用提示（例如 cudaMemAdvise）来指导驱动程序的迁移决策或预取数据（cudaMemPrefetchAsync）。
• 硬件/驱动依赖： 性能特性在不同 GPU 代次和驱动版本之间差异很大。

尽管 UM 简化了开发，但高性能机器学习 (machine learning)运行时通常仍偏爱显式内存管理（为 cudaMalloc 和 cudaMallocHost 使用竞技场），并结合与计算内核一同调度的异步内存复制（cudaMemcpyAsync）。这提供了对数据放置和移动的最大控制，通常是实现峰值性能所必需的，尽管在优先考虑开发简单性或特定访问模式下自动迁移表现良好的场景中，UM 也可以是一个可行的选择。

高级分配器考量

为机器学习 (machine learning)运行时构建高性能内存管理器需要进一步的考量：

• 线程安全： 运行时通常使用多线程来调度操作或管理数据传输。内存分配器必须是线程安全的，这通常通过锁机制或线程局部竞技场实现。锁竞争可能成为瓶颈，从而促使采用更细粒度的锁或无锁方法的设计。
• NUMA 感知： 在多插槽 CPU 系统上，内存访问延迟取决于非均匀内存访问 (NUMA) 域。NUMA 感知的运行时会在最接近执行线程或连接设备（例如 GPU）的节点上分配 CPU 内存，以最小化访问延迟。
• 设备亲和性： 分配应放置在进行计算的特定设备上（例如，GPU 0 对比 GPU 1），以最大程度地减少昂贵的跨设备通信。
• 与异步执行的交互： 在异步运行时中释放内存时，管理器必须确保所有使用该内存的操作都已完成。这通常涉及事件或栅栏等同步原语，延迟内存块的实际重用直到安全为止。
• 调试和性能分析： 分配器应包含检测内存泄漏、缓冲区溢出以及跟踪内存使用模式的机制，以帮助调试和性能分析。

总之，高效内存管理是高性能机器学习运行时系统的基础。竞技场分配、内存锁定、基于活跃度的仔细重用以及可能使用的统一内存等技术都是重要的工具。最优策略通常是这些技术的组合，根据具体的机器学习模型、硬件平台和性能要求进行仔细调整。