【AI系统】内存分配算法

【AI系统】内存分配算法

在AI系统中，内存分配是一个至关重要的环节。随着神经网络模型的日益复杂，如何高效地管理和优化内存使用成为了AI开发者面临的重要挑战。本文将深入探讨AI编译器中的内存分配算法，从模型和硬件的内存演进、内存的划分与复用，到具体的节省内存算法，为你揭示AI系统内存管理的奥秘。

在AI编译器的前端优化中，内存分配是指基于计算图进行分析和内存的管理，而实际上内存分配的实际执行是在AI编译器的后端部分完成的。本文将包括三部分内容，分别介绍模型和硬件的内存演进，内存的划分与复用好处，节省内存的算法。

模型/硬件内存演进

随着AI技术的迅速发展，神经网络模型在图像识别、自然语言处理等领域的应用日益广泛，模型的参数规模也逐渐增大，通常具有数百万甚至数十亿个参数，这使得神经网络模型在训练过程中需要消耗大量的NPU内存资源，以至于部分神经网络模型难以使用单张NPU进行训练而需要使用多张NPU共同完成训练任务。

例如，以ImageNet数据集为例，它是一个包含数百万张图片的数据集，常用于图像识别任务的训练和测试。在使用ImageNet数据集训练Inception v4模型时，如果将BatchSize设置为32，则需要大约40GB的GPU内存空间，而使用ImageNet数据集训练Wide ResNet-152模型时，若BatchSize设置为64，则训练过程中需要的内存空间高达180GB，这充分展示了神经网络模型在训练过程中对GPU内存的高需求。

在过去多年间，NVIDIA公司推出了若干具有代表性的NVIDIA GPU产品，其基本的硬件性能指标如下表所示：

由上表可以看出，随着NVIDIA公司的GPU产品的不断升级和更新，GPU所具有的CUDA核心数量迅速增长，同时其单精度浮点数运算能力也得到了极大的提升，然而GPU的内存容量却增长相对较慢，与此同时神经网络模型对于GPU内存的需求逐渐增加，这种发展趋势在某种程度上导致了GPU的计算性能与内存容量之间的不匹配，GPU内存资源的利用效率变得越来越重要。

回顾神经网络模型的训练流程，如下图所示。在图示的训练过程中，需要将训练数据加载到GPU内存中，同时也需要将正向传播和反向传播的神经网络计算图也加载到GPU内存中，这些数据将占用大量的GPU内存空间，在GPU内存增长相对较慢的硬件发展趋势下，对有限的GPU内存进行高效地使用具有重要的意义。而在AI编译器前端优化部分，此部分工作通常称为内存分配。

围绕模型/硬件的内存演进，本文最后提出1个基本的问题：

一般认为，估计神经网络模型的NPU内存占用情况有助于提高NPU利用率，为准确估计神经网络模型的NPU内存占用情况，构建一个N层的神经网络需要消耗多大的NPU内存空间，应该如何估算？

内存的划分

本文将介绍内存分配中内存划分的相关内容，内存分配主要包括两类情况——静态内存和动态内存。

静态内存

静态内存主要包含三部分内容。首先是模型的权重参数（Parameter），这些权重参数占用的NPU内存空间在训练过程中基本保持不变，它们在训练过程中会原地更新数值。

其次是值节点（Value Node），这部分内存用于存储模型中在整个训练过程中会持续被使用的那些不能被折叠的常量，从而作为静态内存的一部分并占用相应的NPU内存空间。

最后是输出（Output），模型的输出同样需要占用NPU内存的一部分空间。这三部分共同构成了神经网络模型在训练时所需的静态内存，它们对于维持模型的稳定运行至关重要。

一般情况下，对于一些在整个计算图中都会使用的固定的算子，在模型初始化时将一次性申请完所需要的NPU内存空间，在实际训练和推理过程中不再需要频繁地进行NPU内存申请操作，从而提高系统的性能。下图展示了某一计算图所对应的具体的NPU内存占用情况，图中浅绿色部分所标注的为使用静态内存的部分。

动态内存

动态内存主要包含两部分内容。

第一部分是输出张量（Output Tensor），指的是神经网络中每一层算子的输出的Tensor张量所占用的NPU内存，这些输出张量在计算过程中存在许多可以复用的部分，从而提高了内存的利用效率。

第二部分是工作区张量（Workspace Tensor），部分算子需要额外的内存空间来存储中间结果或进行其他计算操作，以便完成所需的计算任务，这部分内存为网络中某些算子在计算过程中作为临时缓冲区动态使用。这两部分动态内存共同支持着神经网络在训练过程中的灵活性和高效性。

一般情况下，对于临时的NPU内存需求，可以进行临时的申请和释放，以节省NPU内存的使用，提高模型的并发能力。此外，动态内存在神经网络中占据了大部分的NPU内存，例如在图中浅绿色部分所标注的为动态内存的部分，与上文中静态内存的部分相比，可见动态内存所消耗的NPU内存更多。

在上述内容中提到，神经网络中每一层的算子的输出Tensor张量所占用的NPU内存分配为动态内存，那么如何在每一个计算图的输出节点上分配内存？怎样才能正确地分配内存？

内存优化算法

内存优化效果示例

下图为MobileNet v2模型在内存优化前后的内存占用示意图，上半所示为MobileNet v2在未经内存优化的情况下NPU内存申请的空间的示意图，可以发现在模型中存在许多碎片化的内存占用情况，模型总体占用的内存较多，而下半所示为经过内存优化后的MobileNet v2内存占用示意图。

可见经过内存优化之后，MobileNet v2模型整体的NPU内存消耗量明显减少，内存申请的碎片数量明显降低，有效的整合了不同算子的内存需求，因此内存优化具有重要的价值。那么应当如何进行内存优化呢？下文将介绍常见的节省内存的算法。

节省内存算法

常见的节省内存的算法分为四种类型，包括空间换内存、计算换内存、模型压缩、内存复用等。

以空间换内存的算法将GPU或NPU内存中的部分模块卸载到CPU内存中（CPU Offload），该类算法更多地常见于针对MoE的模型结构进行算法优化。

以计算换内存的算法将在部分场景下重新计算数据的效率比从存储空间中读取数据的效率高，那么可以选择不保存数据而在需要数据时进行重计算，例如使用梯度检查点（Gradient Checkpointing）技术。

模型压缩的算法在端侧推理场景下应用较多，包括但不限于量化训练Quantification、模型剪枝、模型蒸馏等压缩算法。

本文将介绍的节省内存的算法的重点为内存复用，利用AI编译器对计算图中的数据流进行分析，以允许重用内存。

AI编译器的内存分配算法与传统编译器的寄存器分配非常相似，可以借鉴许多思想。

内存替代

在内存复用中存在若干常见操作，例如替换操作Inplace Operation，即当一块内存不再被需要时，且下一个操作是element-wise时，可以在原地覆盖内存。

如下图中左图所示，在计算完成B算子之后，所分配的内存可以被算子C所复用，同理计算完成算子C之后，所分配的内存可以被算子E所复用。但是在右图中，因为在计算完成算子B之后算子F需要使用算子B的计算结果，因此算子C无法复用算子B所分配的内存。

内存共享

当两个数据所使用的内存大小相同、且有一个数据在参与计算后不再需要，那么后一个数据可以覆盖前一个数据，前一个数据所分配的内存无需重复申请和释放，只需要进行内存共享即可。

如下图所示，算子B的数据在算子C使用过之后不再需要，那么在算子E可以对算子B所使用的内存空间进行内存共享。

内存优化方法

从总体的角度看，上文中提到的两种内存优化方法Inplace Operation和Memory Sharing，如果要较为复杂的模型，如下图中右侧图所示计算图，进行内存优化的操作复杂度较高，为此研究者提出了一些新的内存优化算法。

本文将针对下图所展示的计算图，对一种内存优化算法进行解读，其方法与传统编译器的内存优化方法具有一定的相似之处。在下图中，最左侧为计算图，中间的虚线为计算图中数据的流动，在初始化时AI编译器将感知每一个算子被使用的次数并标记。

在第一步时将对算子B分配一个tag；在第二步为算子C分配新的tag时发现此时不能进行Inplace操作，因为算子B的生命周期仍存在；在第三步中计算算子F时，为其分配一个tag，此时算子F已经计算完成，B算子的生命周期结束，其分配的内存空间也释放到队列中；第四步中继续执行时，算子C已计算完成，其内存空间也将进入内存队列中，但是此时将执行算子E，因此会将之前存放在队列中的红色的tag对应的内存分配给算子E来进行Memory Sharing同时修改对应tag的标记；第五步中可以复用算子E所占用的内存，进行inplace操作。

最终可以看到在上图左下角为最终内存分配的方案，需要注意的是此处为预执行的操作而非模型真实运行中进行的内存分配，不同的颜色代表不同的内存分配的空间。

并行分配

在上文中介绍的内存分配算法主要为串行逻辑的算法，本文将探索并行逻辑的内存分配算法。

如下图所示，左右分别为两种不同的内存分配方案，以串行方式运行从A[1]到A[8]，那这两种分配方案是都有效的。然而左侧的分配方案引入更多的依赖，意味着不能以并行方式运行A[2]和A[5]的计算，而右边可以运行，因此右侧的分配方式更为高效。

本文原文来自cnblogs.com