[发明专利]调度计算图上的操作

说明书

方法、系统和装置，包括在计算机存储介质上编码的计算机程序，用于调度在计算图上表示的操作。方法中的一个包括：由计算图系统接收生成用于处理计算图的调度的请求，获得表示计算图的数据，生成计算图的分隔器；以及生成调度以执行计算图中表示的操作，其中生成调度包括：利用零节点来初始化调度；对于分隔器中的每个节点：确定节点在计算图中是否具有任何前趋节点，当节点具有任何前趋节点时，将该前趋节点添加到调度，以及将该节点添加在调度中，以及将每个子图中的既不在分隔器中也不是计算图上的分隔器中的任何节点的前趋的每个节点添加到调度。

相关申请的交叉引用

本申请要求2019年7月17日提交的美国专利申请No.62/875,433的优先权，其全部内容通过引用的方式并入本文中。

技术领域

本说明书涉及在执行计算图时优化峰值存储器使用。

背景技术

神经网络是采用非线性单元的一层或多层来预测所接收的输入的输出的机器学习模型。除了输出层之外，一些神经网络还包括一个或多个隐藏层。每个隐藏层的输出被用作网络中的下一层的输入，即，下一隐藏层或输出层。网络的每一层根据相应参数集的当前值从所接收的输入生成输出。

为了允许神经网络对各种任务，例如，语音、语言和视觉处理任务，更好地执行，神经网络被设计为具有更大尺寸的层、更大量的层和/或越来越复杂和错综的数据相关性。然而，由于计算成本和所需的存储器容量高，所以难以训练大且复杂的神经网络。为了促进计算速度，在训练复杂神经网络模型时通常采用诸如GPU和AI加速器的专用硬件。在越来越近的神经网络中，当训练这些大模型时，硬件的峰值存储器容量达不到对存储器的峰值需求，从而将存储容量变成训练的瓶颈。

发明内容

本说明书总体上描述用于调度计算图的操作以由一个或多个设备执行以优化峰值存储器使用的技术。

根据一个方面，一种调度计算图的操作的方法包括：由计算图系统接收生成用于处理计算图的调度的请求，获得表示计算图的数据，生成计算图的分隔器，并且生成调度以执行计算图中表示的操作。该调度定义由计算图表示的根据排序来执行的操作序列。计算图包括多个节点和有向边，其中每个节点表示相应操作。从相应第一节点到相应第二节点的每个有向边表示由相应第二节点表示的第二操作需要通过执行由相应第一节点表示的第一操作而生成的至少一个输出作为输入。该分隔器满足从计算图中移除分隔器中的节点和将其他节点连接到分隔器的边使得计算图的剩余节点和边形成多个连接分量的特性。生成所述调度包括：(1)利用零节点来初始化调度；(2)对于分隔器中的每个节点，确定节点在计算图中是否具有任何前趋节点，当节点具有任何前趋节点时，将前趋节点添加到调度，以及将该节点添加在调度中，以及(3)将在计算图上既不在分隔器中也不是分隔器中的任何节点的前趋的每个子图中的每个节点添加到调度。

本说明书中描述的主题能够在特定实施例中实现以便实现以下优势中的一个或多个。

所描述的系统能够被用于在计算有向或非有向非循环图时，即，当在一个或多个计算设备上执行由该图表示的操作时，优化存储器使用。例如，所描述的技术能够被用于在计算用于训练大型和复杂神经网络的图形时优化存储器使用。为了在计算图时优化存储器使用，系统生成按照执行顺序指定图中表示的操作的序列的调度。系统生成的调度是有效的，并且通过利用重新具体化减少了计算图的峰值存储器使用，使得在执行期间的任何给定时间仅必要的操作输入和输出被保持在存储器中。一些实施方式当计算有向或非有向非循环图到给定存储器大小时，例如在特定硬件上可用的存储器时，能够被用于满足峰值存储器要求。

当生成调度时，系统仅重新排列执行每个操作的顺序，决定在存储器中保留哪个操作，如果这样的话，则在何时和多长时间保持操作。因为图中所表示的操作都没有被修改，所以这种调度操作的技术避免了在使用诸如重复使用存储器区域以及CPU与GPU存储器之间的通信等减少存储器使用的其他技术时存在的降低计算图的最终输出的精度的风险。