CN116150048A - 一种内存优化方法、装置、设备及介质 - Google Patents

Abstract

本发明公开了一种内存优化方法、装置、设备及介质。内存优化方法，包括：在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点；根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合；在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储；在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。本发明实施例的技术方案能够在片内分布式计算时，充分利用芯片的内存和计算力，提高芯片计算速度以及***性能。

Description

一种内存优化方法、装置、设备及介质

技术领域

本发明涉及存储技术领域，尤其涉及一种内存优化方法、装置、设备及介质。

背景技术

目前，分布式***在许多领域和行业中都有所应用，市场前景非常好。在分布式***中，内存是影响分布式***性能十分关键的因素。然而，现有分布式***中的内存资源的消耗量较大，影响了分布式***的性能表现。

现有的ASIC芯片和计算框架(如TensorFlow或pytorch等)对于片内分布式深度学习计算的支持优化主要是数据并行以及单个算子在多核的拆分。在超大模型计算时，由于模型参数、梯度以及优化器在卡内存储的冗余，无法充分提升芯片的存储力、计算力以及带宽的利用率。多核ASIC芯片由于受到算子的差异、存储亲和力、访问冲突等影响，很难充分利用计算和存储，造成浪费。

发明内容

本发明提供了一种内存优化方法、装置、设备及介质，以解决片内分布式计算时内存和计算力无法得到充分利用的情况。

根据本发明的一方面，提供了一种内存优化方法，包括：

在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点；

根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合；

在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储；

在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

根据本发明的另一方面，提供了一种内存优化装置，包括：

待布局节点确定模块，用于在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点；

目标计算图集合获取模块，用于根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合；

内存优化模块，用于在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储；在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

根据本发明的另一方面，提供了一种电子设备，所述电子设备包括：

至少一个处理器；以及

与所述至少一个处理器通信连接的存储器；其中，

所述存储器存储有可被所述至少一个处理器执行的计算机程序，所述计算机程序被所述至少一个处理器执行，以使所述至少一个处理器能够执行本发明任一实施例所述的内存优化方法。

根据本发明的另一方面，提供了一种计算机可读存储介质，所述计算机可读存储介质存储有计算机指令，所述计算机指令用于使处理器执行时实现本发明任一实施例所述的内存优化方法。

本发明实施例的技术方案，通过在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点，进而根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合，从而在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储，并在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。在本方案中，针对模型训练阶段以及模型推理阶段，对相应的数据进行等量拆分和亲和存储，可以避免将相同数据在每个片内分布式计算核上进行存储，从而可以充分利用芯片的带宽、计算和存储资源，加快深度学习计算的速度，提升***的性能，解决了现有技术中片内分布式计算时内存和计算力无法得到充分利用的问题，能够在片内分布式计算时，充分利用芯片的内存和计算力，提高芯片计算速度以及***性能。

应当理解，本部分所描述的内容并非旨在标识本发明的实施例的关键或重要特征，也不用于限制本发明的范围。本发明的其它特征将通过以下的说明书而变得容易理解。

附图说明

为了更清楚地说明本发明实施例中的技术方案，下面将对实施例描述中所需要使用的附图作简单地介绍，显而易见地，下面描述中的附图仅仅是本发明的一些实施例，对于本领域普通技术人员来讲，在不付出创造性劳动的前提下，还可以根据这些附图获得其他的附图。

图1为本发明实施例一提供的一种内存优化方法的流程图；

图2为本发明实施例二提供的一种内存优化方法的流程图；

图3是本发明实施例二提供的一种基于ASIC的片内分布式模型计算内存优化方法的流程图；

图4是本发明实施例二提供的一种两个计算核中部署的目标计算图的示意图；

图5为本发明实施例三提供的一种内存优化装置的结构示意图；

图6示出了可以用来实施本发明的实施例的电子设备的结构示意图。

具体实施方式

为了使本技术领域的人员更好地理解本发明方案，下面将结合本发明实施例中的附图，对本发明实施例中的技术方案进行清楚、完整地描述，显然，所描述的实施例仅仅是本发明一部分的实施例，而不是全部的实施例。基于本发明中的实施例，本领域普通技术人员在没有做出创造性劳动前提下所获得的所有其他实施例，都应当属于本发明保护的范围。

需要说明的是，本发明的说明书和权利要求书及上述附图中的术语“初始”、“目标”等是用于区别类似的对象，而不必用于描述特定的顺序或先后次序。应该理解这样使用的数据在适当情况下可以互换，以便这里描述的本发明的实施例能够以除了在这里图示或描述的那些以外的顺序实施。此外，术语“包括”和“具有”以及他们的任何变形，意图在于覆盖不排他的包含，例如，包含了一系列步骤或单元的过程、方法、***、产品或设备不必限于清楚地列出的那些步骤或单元，而是可包括没有清楚地列出的或对于这些过程、方法、产品或设备固有的其它步骤或单元。

实施例一

图1为本发明实施例一提供的一种内存优化方法的流程图，本实施例可适用于在分布式计算时充分利用内存和计算力的情况，该方法可以由内存优化装置来执行，该内存优化装置可以采用硬件和/或软件的形式实现，该内存优化装置可配置于电子设备中。如图1所示，该方法包括：

S110、在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点。

其中，初始计算图可以是特定芯片的神经网络模型中正向或者反向传播计算的有向图。特定芯片可以是需要进行内存优化的芯片。特定芯片的类型可以包括但不限于ASIC(Application Specific Integrated Circuit，专用集成电路)芯片。片内分布式计算条件可以用于判断计算图是否可以进行分布式计算。片内分布式计算条件可以包括但不限于pipeline input流水线输入、主要计算节点部署在ASIC芯片上以及采用了至少一种优化器进行梯度更新等。待布局节点可以是初始计算图中存在的节点。

在本发明实施例中，可以先获取需要进行内存优化的芯片中神经网络模型计算所需的初始计算图，进而判断初始计算图是否符合片内分布式计算条件，若初始计算图符合片内分布式计算条件，则对初始计算图进行节点划分，通过对初始计算图的节点的分类标记，得到待布局节点。

S120、根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合。

其中，目标计算图集合可以是根据待布局节点的节点连接关系以及片内分布式计算核的数量，重构出的计算图集合。节点连接关系可以描述初始计算图中各节点之间的连接关系。片内分布式计算核可以还是芯片中进行分布式计算的计算核。

在本发明实施例中，可以根据片内分布式计算核的数量，对待布局节点的节点进行重构以及复制等处理，并根据待布局节点的节点连接关系，将完成上述处理的节点进行布局连接，得到目标计算图集合。

S130、在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储。

其中，模型参数可以是神经网络模型内部配置的变量。梯度变量可以用于描述模型参数的变化幅度。优化器状态变量可以是神经网络模型中优化器的状态变量。示例性的，当优化器为Adam优化器时，优化器状态变量可以包括参数的滑动平均以及参数的滑动平均方差等。等量拆分可以是按照数据量进行数据分组的操作。亲和存储可以按照片内分布式计算核与数据的亲和性进行数据存储的操作。

在本发明实施例中，当需要内存优化的芯片中的片内分布式计算核进行模型训练时，可以确定目标计算图集合中与模型训练相关的模型参数、梯度变量以及优化器状态变量，分别按照各自的数据量进行等量拆分，并将完成数据分组的数据，按照与片内分布式计算核的亲和性进行亲和存储。

S140、在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

在本发明实施例中，当需要内存优化的芯片中的片内分布式计算核进行模型推理时，可以确定目标计算图集合中与模型推理相关的模型参数，进而按照模型参数的数据量，对模型参数进行等量拆分，从而将分组好的数据按照与片内分布式计算核的亲和性进行亲和存储。

本方案对初始计算图的节点进行分析和分类，根据不同的输入方式、计算类型和优化器对计算图进行重构，并对于模型参数、梯度变量以及优化器状态变量在片内多卡进行拆分，实现片内分布式计算的片内相关参数的片内存储优化。

为了和分布式计算框架(如Pytorch和Horovod等)的兼容，可以通过计算框架API(Application Programming Interface,应用程序编程接口)或者全局环境变量的方式使用相关优化，还可以和多卡分布式计算框架兼容，进一步提升分布式计算的并行度。

片内内存优化技术适用于全浮点精度FP32，半浮点精度FP16、BF16以及FP8浮点精度的模型的训练与推理计算，适用于所有的数据精度(浮点与整数)。

本发明实施例的技术方案，通过在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点，进而根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合，从而在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储，并在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。在本方案中，针对模型训练阶段以及模型推理阶段，对相应的数据进行等量拆分和亲和存储，可以避免将相同数据在每个片内分布式计算核上进行存储，从而可以充分利用芯片的带宽、计算和存储资源，加快深度学习计算的速度，提升***的性能，解决了现有技术中片内分布式计算时内存和计算力无法得到充分利用的问题，能够在片内分布式计算时，充分利用芯片的内存和计算力，提高芯片计算速度以及***性能。

实施例二

图2为本发明实施例二提供的一种内存优化方法的流程图，本实施例以上述实施例为基础进行具体化，给出了根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合的具体的可选的实施方式。如图2所示，该方法包括：

S210、在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点。

在本发明的一个可选实施例中，对初始计算图进行节点划分，确定待布局节点，可以包括：对初始计算图中的节点进行划分，确定初始计算图中的输入节点、主计算节点以及共享节点；将输入节点、主计算节点以及共享节点，作为待布局节点。

其中，输入节点可以是待布局节点中为其它节点提供输入数据的节点。主计算节点可以是待布局节点中进行数据计算处理的节点。共享节点可以是待布局节点中可以为其他节点提供统一数据支撑或功能支撑的节点。

在本发明实施例中，可以按照节点功能，对初始计算图中的节点进行划分，得到初始计算图中的输入节点、主计算节点以及共享节点，从而将输入节点、主计算节点以及共享节点，作为待布局节点。

S220、根据片内分布式计算核的数量，对待布局节点中的输入节点进行重构操作，得到各重构输入节点，并对待布局节点中的主计算节点进行复制操作，得到各主计算复制节点。

其中，重构输入节点可以是重构出的输入节点。主计算复制节点可以是对主计算节点进行复制操作后得到的计算节点。

在本发明实施例中，可以根据片内分布式计算核的数量，确定输入节点的重构份数，对待布局节点中的输入节点进行重构操作，得到重构份数的重构输入节点，还可以根据片内分布式计算核的数量，确定每个主计算节点的复制份数，从而将每个主计算节点按照复制份数进行复制操作，得到各主计算复制节点。

示例性的，假设片内分布式计算核的数量为n,则对待布局节点中的输入节点进行重构操作，得到n个重构输入节点。假设待布局节点中包括主计算节点a以及主计算节点b，则对待布局节点中的主计算节点a进行复制操作，得到主计算复制节点a1、…、主计算复制节点an，对待布局节点中的主计算节点b进行复制操作，得到主计算复制节点b1、…、主计算复制节点bn。

S230、根据待布局节点的节点连接关系，对各重构输入节点、共享节点以及各主计算复制节点进行布局，得到目标计算图集合。

在本发明实施例中，可以根据待布局节点的节点连接关系，对各重构输入节点、共享节点以及各主计算复制节点进行布局连接，得到目标计算图集合。

在本发明的一个可选实施例中，根据待布局节点的节点连接关系，对各重构输入节点、共享节点以及各主计算复制节点进行布局，得到目标计算图集合，可以包括：对各主计算复制节点进行分组，确定各目标计算复制节点集合；根据待布局节点的节点连接关系、当前重构输入节点、当前目标计算复制节点集合以及共享节点，确定目标计算图集合中的当前目标计算图。

其中，目标计算复制节点集合可以是由待布局节点中的每个主计算节点分别复制出的一个主计算复制节点构成的节点集合。示例性的，假设片内分布式计算核的数量为2，对主计算节点a进行复制操作，得到主计算复制节点a1以及主计算复制节点a2，对主计算节点b进行复制操作，得到主计算复制节点b1以及主计算复制节点b2，通过对主计算复制节点a1、主计算复制节点a2、主计算复制节点b1以及主计算复制节点b2，进行分组，可以得到两个目标计算复制节点集合：{主计算复制节点a1，主计算复制节点b1}，{主计算复制节点a2，主计算复制节点b2}。当前目标计算图可以是根据待布局节点的节点连接关系、当前重构输入节点、当前目标计算复制节点集合以及共享节点构建出的计算图。

在本发明实施例中，可以先确定待布局节点中的全部主计算节点，按照全部主计算节点的标记对各主计算复制节点进行分组，得到各目标计算复制节点集合，进而从重构输入节点中选择出当前重构输入节点，并从多个目标计算复制节点集合中选择出当前目标计算复制节点集合，进而将当前重构输入节点、当前目标计算复制节点集合中的主计算复制节点，以及共享节点，按照待布局节点的节点连接关系进行布局连接，得到目标计算图集合中的当前目标计算图，以使当前目标计算图与初始计算图具备相同的模型训练与推理功能。

S240、在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储。

在本发明的一个可选实施例中，在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储，可以包括：获取底层调度时的可训练变量亲和数据以及优化器状态亲和数据；在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数进行等量拆分，得到各模型参数拆分结果，并将目标计算图集合上关联的梯度变量进行等量拆分，得到各梯度变量拆分结果；根据可训练变量亲和数据，将各模型参数拆分结果以及各梯度变量拆分结果，在相应片内分布式计算核进行亲和存储；将优化器状态变量进行等量拆分，得到各优化器状态变量拆分结果，并根据优化器状态亲和数据，将各优化器状态变量拆分结果，在相应片内分布式计算核进行亲和存储。

其中，可训练变量可以是神经网络模型中需要更新的深度学习变量。可训练变量亲和数据可以用于表示底层调度时遵循的可训练变量与片内分布式计算核的亲和性规范。优化器状态亲和数据可以用于表示底层调度时遵循的优化器状态变量与片内分布式计算核的亲和性规范。模型参数拆分结果可以是模型参数等量拆分的结果。梯度变量拆分结果可以是梯度变量等量拆分的结果。优化器状态变量拆分结果可以是优化器状态变量等量差分的结果。

在本发明实施例中，可以先读取底层调度时的可训练变量亲和数据以及优化器状态亲和数据，通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数进行等量拆分，得到模型参数拆分结果，并将目标计算图集合上关联的梯度变量进行等量拆分，得到梯度变量拆分结果，从而对可训练变量亲和数据进行解析，分别确定与各模型参数拆分结果亲和的片内分布式计算核，并分别确定与各梯度变量拆分结果亲和的片内分布式计算核，从而将各模型参数拆分结果在相应亲和的片内分布式计算核进行亲和存储，并将各梯度变量拆分结果在相应亲和的片内分布式计算核进行亲和存储。进一步，将优化器状态变量进行等量拆分，得到各优化器状态变量拆分结果，并通过解析优化器状态亲和数据，分别确定与各优化器状态变量拆分结果亲和的片内分布式计算核，从而将各优化器状态变量拆分结果在相应亲和的片内分布式计算核进行亲和存储。

可选的，由于模型参数更频繁的被多核同时获取使用时，会造成比较严重的访问冲突，因此为了解决多核同时获取一组模型参数引起的访存冲突，可以将模型参数进行不多于半数片内分布式计算核份数的拷贝，进而将一组模型参数间隔存储于不同的片内分布式计算核上。

S250、在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

在本发明的一个可选实施例中，在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储，可以包括：获取底层调度时的可训练变量亲和数据；在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分，得到模型参数拆分结果；根据可训练变量亲和数据，将模型参数拆分结果，在相应片内分布式计算核进行亲和存储。

在本发明实施例中，可以读取底层调度时的可训练变量亲和数据，在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分，得到模型参数拆分结果，进而对可训练变量亲和数据进行解析，分别确定与各模型参数拆分结果亲和的片内分布式计算核，从而将各模型参数拆分结果在相应亲和的片内分布式计算核进行亲和存储。

在本发明的一个可选实施例中，在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储之后，可以包括：在模型参数、梯度变量以及优化器状态变量更新后，返回执行将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储的操作；在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储之后，可以包括：在模型参数更新后，返回执行将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储的操作。

在本发明实施例中，片内分布式计算核基于目标计算图集合进行模型训练时，模型参数、梯度变量以及优化器状态变量会进行参数更新，当模型参数、梯度变量以及优化器状态变量更新后，可以返回执行将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储的操作。片内分布式计算核基于目标计算图集合进行模型推理时，若模型参数出现更新时，可以返回执行将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储的操作。

图3是本发明实施例二提供的一种基于ASIC的片内分布式模型计算内存优化方法的流程图。如图3所示，该方法包括如下步骤：

1)检查初始计算图是否符合片内分布式计算条件(pipeline input流水线输入、主计算节点部署在ASIC芯片上、采用了至少一种优化器进行梯度更新)。若符合片内分布式计算条件，则进行后续步骤，若不符合则结束流程。

2)对初始计算图上的节点进行分类和标记，将节点分为输入节点、可训练变量、主计算节点以及共享节点。其中，输入节点可以包括循环队列、迭代器，以及分阶区域等。

3)重构输入节点，使各重构输入节点能够从数据缓存区域正确取出数据。具体是，复制多份取数据的节点，得到各重构输入节点，并保留同一个缓存区节点。

4)处理可训练变量，在分布式训练时用片内共享方式读取。

5)复制主计算节点，并创建目标计算图集合中的目标计算图，复制份数由片内分布式计算核的数量决定。复制时要处理和共享节点之间的关系。如A为共享节点，B为主计算节点，且有关系：B的输入是A，则复制出的主计算复制节点B1、主计算复制节点B2以及主计算复制节点B3等，都要以A为输入。

6)内存优化：

a、在模型训练时，对初始计算图上相关的模型的可训练变量进行内存优化，将所有相关的模型参数在片内分布式计算核按照模型使用顺序进行等量拆分以及亲和存储，并将相关模型参数的梯度变量也按照同样的使用顺序进行等量拆分以及亲和存储，还可对初始计算图上相关模型的优化器状态变量进行内存优化，将所有模型参数对应的优化器状态变量也按照使用顺序进行等量拆分以及亲和存储。在实际需要使用时，从相应的片内分布式计算核上获取相关的模型参数，相应模型参数的梯度变量更新后进行多核间同一梯度变量的reduce求和操作，并在求和之后销毁当前片内分布式计算核的梯度变量的存储空间，只保存更新后的梯度变量在初始设置的片内分布式计算核进行亲和存储时的存储空间。在优化器状态变量进行更新时，只需要在当前片内分布式计算核进行优化器状态变量的更新操作。

b、在模型推理时，由于可训练变量已经固定，推理时只需要将相关模型参数在片内分布式计算核按照计算时间等量拆分以及亲和存储，由于在模型推理时没有优化器，忽略优化器状态变量的存储优化。

7)处理其它辅助功能。如重置数据汇总(summary)节点，可以汇总不同分布式部分的数据给客户。其中，重置数据汇总属于一类共享节点。

图4是本发明实施例二提供的一种两个计算核中部署的目标计算图的示意图。如图4所示，片内分布式计算核1的目标计算图中的前向主计算复制节点a1，与片内分布式计算核2的目标计算图中的前向主计算复制节点a2复制于同一个主计算节点。片内分布式计算核1的目标计算图中的前向主计算复制节点b1，与片内分布式计算核2的目标计算图中的前向主计算复制节点b2复制于同一个主计算节点。同理，主计算复制节点b1的反向计算节点b1_backward与主计算复制节点b2的反向计算节点b2_backward，复制于同一个主计算节点，主计算复制节点a1的反向计算节a1_backward与主计算复制节点a2的反向计算节点a2_backward，复制于同一个主计算节点。模型参数W1以及W1的梯度变量与片内分布式计算核1亲和，并存储于片内分布式计算核1。模型参数W2以及W2的梯度变量与片内分布式计算核2亲和，并存储于片内分布式计算核2。片内分布式计算核1的优化器状态变量存储于片内分布式计算核1的优化器内，片内分布式计算核2的优化器状态变量存储于片内分布式计算核2优化器内。

本发明实施例的技术方案，通过在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点，从而根据片内分布式计算核的数量，对待布局节点中的输入节点进行重构操作，得到各重构输入节点，并对待布局节点中的各主计算节点进行复制操作，得到各主计算复制节点，进一步根据待布局节点的节点连接关系，对各重构输入节点、共享节点以及各主计算复制节点进行布局，得到目标计算图集合，从而在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储，并在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。在本方案中，针对模型训练阶段以及模型推理阶段，对相应的数据进行等量拆分和亲和存储，可以避免将相同数据在每个片内分布式计算核上进行存储，从而可以充分利用芯片的带宽、计算和存储资源，加快深度学习计算的速度，提升***的性能，解决了现有技术中片内分布式计算时内存和计算力无法得到充分利用的问题，能够在片内分布式计算时，充分利用芯片的内存和计算力，提高芯片计算速度以及***性能。

实施例三

图5为本发明实施例三提供的一种内存优化装置的结构示意图。如图5所示，该装置包括：待布局节点确定模块310、目标计算图集合获取模块320以及内存优化模块330，其中，

待布局节点确定模块310，用于在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点；

目标计算图集合获取模块320，用于根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合；

内存优化模块330，用于在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储；在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

本发明实施例的技术方案，通过在初始计算图符合片内分布式计算条件时，对初始计算图进行节点划分，确定待布局节点，进而根据待布局节点的节点连接关系以及片内分布式计算核的数量，对待布局节点进行布局，得到目标计算图集合，从而在通过片内分布式计算核进行模型训练时，将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储，并在通过片内分布式计算核进行模型推理时，将目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。在本方案中，针对模型训练阶段以及模型推理阶段，对相应的数据进行等量拆分和亲和存储，可以避免将相同数据在每个片内分布式计算核上进行存储，从而可以充分利用芯片的带宽、计算和存储资源，加快深度学习计算的速度，提升***的性能，解决了现有技术中片内分布式计算时内存和计算力无法得到充分利用的问题，能够在片内分布式计算时，充分利用芯片的内存和计算力，提高芯片计算速度以及***性能。

可选的，待布局节点确定模块310，具体用于对所述初始计算图中的节点进行划分，确定所述初始计算图中的输入节点、主计算节点以及共享节点；将所述输入节点、所述主计算节点以及所述共享节点，作为所述待布局节点。

可选的，目标计算图集合获取模块320包括重构节点创建单元以及目标计算图集合获取单元，重构节点创建单元，用于根据所述片内分布式计算核的数量，对所述待布局节点中的输入节点进行重构操作，得到各重构输入节点，并对待布局节点中的主计算节点进行复制操作，得到各主计算复制节点。目标计算图集合获取单元，用于根据所述待布局节点的节点连接关系，对各所述重构输入节点、共享节点以及各所述主计算复制节点进行布局，得到所述目标计算图集合。

可选的，目标计算图集合获取单元，用于对各所述主计算复制节点进行分组，确定各目标计算复制节点集合；根据所述待布局节点的节点连接关系、当前重构输入节点、当前目标计算复制节点集合以及所述共享节点，确定所述目标计算图集合中的当前目标计算图。

可选的，内存优化模块330包括第一亲和存储单元，第一亲和存储单元用于获取底层调度时的可训练变量亲和数据以及优化器状态亲和数据；在通过所述片内分布式计算核进行模型训练时，将所述目标计算图集合上关联的所述模型参数进行等量拆分，得到各模型参数拆分结果，并将所述目标计算图集合上关联的梯度变量进行等量拆分，得到各梯度变量拆分结果；根据所述可训练变量亲和数据，将各所述模型参数拆分结果以及各所述梯度变量拆分结果，在相应片内分布式计算核进行亲和存储；将所述优化器状态变量进行等量拆分，得到各优化器状态变量拆分结果，并根据所述优化器状态亲和数据，将各所述优化器状态变量拆分结果，在相应片内分布式计算核进行亲和存储。

可选的，内存优化模块330包括第二亲和存储单元，用于获取底层调度时的可训练变量亲和数据；在通过所述片内分布式计算核进行模型推理时，将所述目标计算图集合上关联的模型参数进行等量拆分，得到模型参数拆分结果；根据所述可训练变量亲和数据，将所述模型参数拆分结果，在相应片内分布式计算核进行亲和存储。

可选的，内存优化装置还包括第一更新存储模块以及第二更新存储模块，第一更新存储模块，用于在所述模型参数、所述梯度变量以及所述优化器状态变量更新后，返回执行将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储的操作。第二更新存储模块，用于在所述模型参数更新后，返回执行将所述目标计算图集合上关联的模型参数进行等量拆分以及亲和存储的操作。

本发明实施例所提供的内存优化装置可执行本发明任意实施例所提供的内存优化方法，具备执行方法相应的功能模块和有益效果。

实施例四

图6示出了可以用来实施本发明的实施例的电子设备的结构示意图。电子设备旨在表示各种形式的数字计算机，诸如，膝上型计算机、台式计算机、工作台、个人数字助理、服务器、刀片式服务器、大型计算机、和其它适合的计算机。电子设备还可以表示各种形式的移动装置，诸如，个人数字处理、蜂窝电话、智能电话、可穿戴设备(如头盔、眼镜、手表等)和其它类似的计算装置。本文所示的部件、它们的连接和关系、以及它们的功能仅仅作为示例，并且不意在限制本文中描述的和/或者要求的本发明的实现。

如图6所示，电子设备10包括至少一个处理器11，以及与至少一个处理器11通信连接的存储器，如只读存储器(ROM)12、随机访问存储器(RAM)13等，其中，存储器存储有可被至少一个处理器执行的计算机程序，处理器11可以根据存储在只读存储器(ROM)12中的计算机程序或者从存储单元18加载到随机访问存储器(RAM)13中的计算机程序，来执行各种适当的动作和处理。在RAM 13中，还可存储电子设备10操作所需的各种程序和数据。处理器11、ROM 12以及RAM 13通过总线14彼此相连。输入/输出(I/O)接口15也连接至总线14。

电子设备10中的多个部件连接至I/O接口15，包括：输入单元16，例如键盘、鼠标等；输出单元17，例如各种类型的显示器、扬声器等；存储单元18，例如磁盘、光盘等；以及通信单元19，例如网卡、调制解调器、无线通信收发机等。通信单元19允许电子设备10通过诸如因特网的计算机网络和/或各种电信网络与其他设备交换信息/数据。

处理器11可以是各种具有处理和计算能力的通用和/或专用处理组件。处理器11的一些示例包括但不限于中央处理单元(CPU)、图形处理单元(GPU)、各种专用的人工智能(AI)计算芯片、各种运行机器学习模型算法的处理器、数字信号处理器(DSP)、以及任何适当的处理器、控制器、微控制器等。处理器11执行上文所描述的各个方法和处理，例如内存优化方法。

在一些实施例中，内存优化方法可被实现为计算机程序，其被有形地包含于计算机可读存储介质，例如存储单元18。在一些实施例中，计算机程序的部分或者全部可以经由ROM 12和/或通信单元19而被载入和/或安装到电子设备10上。当计算机程序加载到RAM 13并由处理器11执行时，可以执行上文描述的内存优化方法的一个或多个步骤。备选地，在其他实施例中，处理器11可以通过其他任何适当的方式(例如，借助于固件)而被配置为执行内存优化方法。

本文中以上描述的***和技术的各种实施方式可以在数字电子电路***、集成电路***、场可编程门阵列(FPGA)、专用集成电路(ASIC)、专用标准产品(ASSP)、芯片上***的***(SOC)、负载可编程逻辑设备(CPLD)、计算机硬件、固件、软件、和/或它们的组合中实现。这些各种实施方式可以包括：实施在一个或者多个计算机程序中，该一个或者多个计算机程序可在包括至少一个可编程处理器的可编程***上执行和/或解释，该可编程处理器可以是专用或者通用可编程处理器，可以从存储***、至少一个输入装置、和至少一个输出装置接收数据和指令，并且将数据和指令传输至该存储***、该至少一个输入装置、和该至少一个输出装置。

用于实施本发明的方法的计算机程序可以采用一个或多个编程语言的任何组合来编写。这些计算机程序可以提供给通用计算机、专用计算机或其他可编程数据处理装置的处理器，使得计算机程序当由处理器执行时使流程图和/或框图中所规定的功能/操作被实施。计算机程序可以完全在机器上执行、部分地在机器上执行，作为独立软件包部分地在机器上执行且部分地在远程机器上执行或完全在远程机器或服务器上执行。

在本发明的上下文中，计算机可读存储介质可以是有形的介质，其可以包含或存储以供指令执行***、装置或设备使用或与指令执行***、装置或设备结合地使用的计算机程序。计算机可读存储介质可以包括但不限于电子的、磁性的、光学的、电磁的、红外的、或半导体***、装置或设备，或者上述内容的任何合适组合。备选地，计算机可读存储介质可以是机器可读信号介质。机器可读存储介质的更具体示例会包括基于一个或多个线的电气连接、便携式计算机盘、硬盘、随机存取存储器(RAM)、只读存储器(ROM)、可擦除可编程只读存储器(EPROM或快闪存储器)、光纤、便捷式紧凑盘只读存储器(CD-ROM)、光学储存设备、磁储存设备、或上述内容的任何合适组合。

为了提供与用户的交互，可以在电子设备上实施此处描述的***和技术，该电子设备具有：用于向用户显示信息的显示装置(例如，CRT(阴极射线管)或者LCD(液晶显示器)监视器)；以及键盘和指向装置(例如，鼠标或者轨迹球)，用户可以通过该键盘和该指向装置来将输入提供给电子设备。其它种类的装置还可以用于提供与用户的交互；例如，提供给用户的反馈可以是任何形式的传感反馈(例如，视觉反馈、听觉反馈、或者触觉反馈)；并且可以用任何形式(包括声输入、语音输入或者、触觉输入)来接收来自用户的输入。

可以将此处描述的***和技术实施在包括后台部件的计算***(例如，作为数据服务器)、或者包括中间件部件的计算***(例如，应用服务器)、或者包括前端部件的计算***(例如，具有图形用户界面或者网络浏览器的用户计算机，用户可以通过该图形用户界面或者该网络浏览器来与此处描述的***和技术的实施方式交互)、或者包括这种后台部件、中间件部件、或者前端部件的任何组合的计算***中。可以通过任何形式或者介质的数字数据通信(例如，通信网络)来将***的部件相互连接。通信网络的示例包括：局域网(LAN)、广域网(WAN)、区块链网络和互联网。

计算***可以包括客户端和服务器。客户端和服务器一般远离彼此并且通常通过通信网络进行交互。通过在相应的计算机上运行并且彼此具有客户端-服务器关系的计算机程序来产生客户端和服务器的关系。服务器可以是云服务器，又称为云计算服务器或云主机，是云计算服务体系中的一项主机产品，以解决了传统物理主机与VPS服务中，存在的管理难度大，业务扩展性弱的缺陷。

应该理解，可以使用上面所示的各种形式的流程，重新排序、增加或删除步骤。例如，本发明中记载的各步骤可以并行地执行也可以顺序地执行也可以不同的次序执行，只要能够实现本发明的技术方案所期望的结果，本文在此不进行限制。

上述具体实施方式，并不构成对本发明保护范围的限制。本领域技术人员应该明白的是，根据设计要求和其他因素，可以进行各种修改、组合、子组合和替代。任何在本发明的精神和原则之内所作的修改、等同替换和改进等，均应包含在本发明保护范围之内。

Claims (10)

1.一种内存优化方法，其特征在于，包括：

在初始计算图符合片内分布式计算条件时，对所述初始计算图进行节点划分，确定待布局节点；

根据待布局节点的节点连接关系以及片内分布式计算核的数量，对所述待布局节点进行布局，得到目标计算图集合；

在通过所述片内分布式计算核进行模型训练时，将所述目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储；

在通过所述片内分布式计算核进行模型推理时，将所述目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

2.根据权利要求1所述的方法，其特征在于，所述对所述初始计算图进行节点划分，确定待布局节点，包括：

对所述初始计算图中的节点进行划分，确定所述初始计算图中的输入节点、主计算节点以及共享节点；

将所述输入节点、所述主计算节点以及所述共享节点，作为所述待布局节点。

3.根据权利要求2所述的方法，其特征在于，所述根据待布局节点的节点连接关系以及片内分布式计算核的数量，对所述待布局节点进行布局，得到目标计算图集合，包括：

根据所述片内分布式计算核的数量，对所述待布局节点中的输入节点进行重构操作，得到各重构输入节点，并对待布局节点中的主计算节点进行复制操作，得到各主计算复制节点；

根据所述待布局节点的节点连接关系，对各所述重构输入节点、共享节点以及各所述主计算复制节点进行布局，得到所述目标计算图集合。

4.根据权利要求3中所述的方法，其特征在于，所述根据所述待布局节点的节点连接关系，对各所述重构输入节点、共享节点以及各所述主计算复制节点进行布局，得到所述目标计算图集合，包括：

对各所述主计算复制节点进行分组，确定各目标计算复制节点集合；

根据所述待布局节点的节点连接关系、当前重构输入节点、当前目标计算复制节点集合以及所述共享节点，确定所述目标计算图集合中的当前目标计算图。

5.根据权利要求4所述的方法，其特征在于，所述在通过所述片内分布式计算核进行模型训练时，将所述目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储，包括：

获取底层调度时的可训练变量亲和数据以及优化器状态亲和数据；

在通过所述片内分布式计算核进行模型训练时，将所述目标计算图集合上关联的所述模型参数进行等量拆分，得到各模型参数拆分结果，并将所述目标计算图集合上关联的梯度变量进行等量拆分，得到各梯度变量拆分结果；

根据所述可训练变量亲和数据，将各所述模型参数拆分结果以及各所述梯度变量拆分结果，在相应片内分布式计算核进行亲和存储；

将所述优化器状态变量进行等量拆分，得到各优化器状态变量拆分结果，并根据所述优化器状态亲和数据，将各所述优化器状态变量拆分结果，在相应片内分布式计算核进行亲和存储。

6.根据权利要求4所述的方法，其特征在于，所述在通过所述片内分布式计算核进行模型推理时，将所述目标计算图集合上关联的模型参数进行等量拆分以及亲和存储，包括：

获取底层调度时的可训练变量亲和数据；

在通过所述片内分布式计算核进行模型推理时，将所述目标计算图集合上关联的模型参数进行等量拆分，得到模型参数拆分结果；

根据所述可训练变量亲和数据，将所述模型参数拆分结果，在相应片内分布式计算核进行亲和存储。

7.根据权利要求1所述的方法，其特征在于，在通过所述片内分布式计算核进行模型训练时，将所述目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储之后，包括：

在所述模型参数、所述梯度变量以及所述优化器状态变量更新后，返回执行将目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储的操作；

在通过所述片内分布式计算核进行模型推理时，将所述目标计算图集合上关联的模型参数进行等量拆分以及亲和存储之后，包括：

在所述模型参数更新后，返回执行将所述目标计算图集合上关联的模型参数进行等量拆分以及亲和存储的操作。

8.一种内存优化装置，其特征在于，包括：

待布局节点确定模块，用于在初始计算图符合片内分布式计算条件时，对所述初始计算图进行节点划分，确定待布局节点；

目标计算图集合获取模块，用于根据待布局节点的节点连接关系以及片内分布式计算核的数量，对所述待布局节点进行布局，得到目标计算图集合；

内存优化模块，用于在通过所述片内分布式计算核进行模型训练时，将所述目标计算图集合上关联的模型参数、梯度变量以及优化器状态变量进行等量拆分以及亲和存储；在通过所述片内分布式计算核进行模型推理时，将所述目标计算图集合上关联的模型参数进行等量拆分以及亲和存储。

9.一种电子设备，其特征在于，所述电子设备包括：

至少一个处理器；以及

与所述至少一个处理器通信连接的存储器；其中，

所述存储器存储有可被所述至少一个处理器执行的计算机程序，所述计算机程序被所述至少一个处理器执行，以使所述至少一个处理器能够执行权利要求1-7中任一项所述的内存优化方法。

10.一种计算机可读存储介质，其特征在于，所述计算机可读存储介质存储有计算机指令，所述计算机指令用于使处理器执行时实现权利要求1-7中任一项所述的内存优化方法。

Images