动态张量运算自动优化技术解析

自动优化动态张量运算的执行

深度学习模型核心依赖于涉及张量（矩阵的高维类比）的代数运算，这些运算可能重复数万次。高效学习需要优化频繁重复的张量运算。但涉及不同形状张量（如32x32、64x64、128x128等）的运算必须单独优化。

自动调度器是学习优化形状的程序，这些形状在当前的张量运算库中实现可能不够优化。然而，现有的自动调度器在处理形状变化的工作负载时表现不佳。例如，许多自然语言处理应用程序接受任意长度的输入，这意味着任意形状的张量。

在今年机器学习与系统会议（MLSys）上，我们和同事展示了一个名为DietCode的新型自动调度器，它比前代产品更高效地处理动态形状工作负载。现有的自动编码器必须单独优化每个可能的形状，而DietCode构建了一个形状通用的搜索空间，使其能够同时优化所有可能的形状。

我们在一个自然语言处理任务上测试了方法，该任务可以接受大小从1到128个令牌不等的输入。当使用反映合理真实世界分布的输入大小随机抽样时，相对于先前最佳的自动调度器，我们将优化过程加速了近六倍。当我们考虑所有可能的形状时，加速比增加到94倍以上。

尽管速度更快，DietCode还提高了生成代码的性能，相对于先前的自动调度器提高了70%，相对于现有张量运算库中手动优化的代码提高了19%。因此，它有望加速客户的动态形状机器学习工作负载。

动态工作负载

处理任意长度文本字符串的自然语言处理模型是动态设计模型的例子，这些模型允许可变大小的输入。但其他应用程序也需要动态工作负载。

例如，神经架构搜索通过从不同形状的组件构建它们来尝试不同的深度学习架构，这需要不同形状张量的运算。一些模型（例如BERT语言模型）在网络的不同层应用相同的运算，这些层具有不同数量的节点。

微内核

自动调度器通常依赖于计算内核——程序模板，其使用大大加速了评估不同候选优化的速度。然而，奇怪形状的工作负载可能无法精确适应内核。例如，如果张量在其中一个维度上有513个元素，但内核容量只有512，那么必须将两个内核拼接在一起以适应张量。

然而，拼接后的内核在相关维度上的组合容量为1,024，而输入张量只有513。因此，必须对输入张量进行填充以填满内核。这种填充可能显著减慢优化过程，因为它导致不必要的计算，然后必须从结果中修剪掉。

DietCode使用根据可用硬件而非输入形状调整大小的微内核，这有助于针对该硬件进行优化。对于给定的硬件配置，DietCode还可以生成一系列不同的微内核形状和大小，这些可以组合使用。

微内核足够小，通常可以平铺在输入上，以更精确地适应其形状。这可能仍然需要在边缘进行一些填充，但比更大的内核需要的填充少得多。

然而，微内核的真正优势是使DietCode能够同时优化多个形状的运算符。标准的自动调度器将获取工作负载形状，根据需要填充以适应其拼接的内核，然后使用提取程序特征（如循环结构和内存访问模式）的成本模型来估计不同实现的效率。然后它将为下一个形状重复该过程。

相比之下，DietCode将运算符分解到微内核上。成本模型有两个组成部分：一个评估分配给每个微内核的部分运算的特征，另一个评估将这些部分运算拼接在一起形成完整运算符的成本。

在这里，我们实现了最大的效率增益，因为每个部分运算都是多个工作负载形状的运算符的组成部分。与评估运算的计算成本（这是一个涉及实际硬件测量的机器学习过程）相比，将部分运算拼接在一起的成本很低。

有了优化的微内核，我们训练了一个高效的决策树模型，将工作负载形状映射到微内核。该决策树被合并到执行张量运算的二进制文件中，以将任意形状的输入路由到适当的微内核进行处理。

有关实验结果和更多细节，请参阅论文。

致谢：Cody Yu, Yizhi Liu, Gennady Pekhimenko
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）

公众号二维码

公众号二维码