动态张量运算自动优化技术突破

动态张量运算的自动优化

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

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

在今年机器学习与系统会议（MLSys）上，团队提出了名为DietCode的新型自动调度器，其处理动态形状工作负载的效率远超前代方案。现有自动编码器必须单独优化每种可能形状，而DietCode构建了形状通用搜索空间，能够同时优化所有可能形状。

在自然语言处理（NLP）任务测试中，输入尺寸范围从1到128个token。当采用反映真实世界分布的随机输入尺寸采样时，优化过程相比先前最佳自动调度器加速近6倍。若考虑所有可能形状，加速比提升至94倍以上。

尽管速度大幅提升，DietCode还改善了生成代码的性能：相比先前自动调度器提升最高达70%，相比现有张量运算库中手动优化代码提升最高达19%。这有望加速客户的动态形状机器学习工作负载。

动态工作负载

处理任意长度文本字符串的NLP模型是动态设计模型的典型代表，允许可变尺寸输入。但其他应用也需要动态工作负载，例如神经架构搜索通过不同形状组件构建深度学习架构，这需要操作不同形状张量。某些模型（如BERT语言模型）在网络不同层应用相同操作，这些层具有不同数量的节点。

微内核技术

自动调度器通常依赖计算内核——程序模板能极大加速不同候选优化的评估速度。然而非常规形状工作负载可能无法精确匹配内核。例如，若张量某维度有513个元素，但内核容量仅为512，则必须拼接两个内核以适应张量。

但拼接后内核在相关维度上的组合容量达到1024，而输入张量仅需513。因此必须对输入张量进行填充以填满内核。这种填充会显著拖慢优化过程，因为它导致不必要的计算，后续还需从结果中修剪。

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

微内核足够小，通常可以平铺覆盖输入，更精确地适应其形状。这可能仍需要在边缘进行少量填充，但远少于较大内核的需求。

然而微内核的真正优势在于使DietCode能够同时优化多种形状的算子。标准自动调度器会获取工作负载形状，按需填充以适应拼接内核，然后使用提取程序特征（如循环结构和内存访问模式）的成本模型评估不同实现的效率。随后对下一个形状重复此过程。

相比之下，DietCode将算子分解到微内核中。成本模型包含两个组件：一个评估分配给每个微内核的部分操作特征，另一个评估将这些部分操作拼接形成完整算子的成本。

在此实现最大效率提升，因为每个部分操作都是多种工作负载形状算子的组成部分。与评估操作的计算成本（涉及实际硬件测量的机器学习过程）相比，拼接部分操作的成本较低。

获得优化的微内核后，训练高效决策树模型将工作负载形状映射到微内核。该决策树被集成到执行张量操作的二进制文件中，将任意形状的输入路由到适当的微内核进行处理。

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

致谢：Cody Yu, Yizhi Liu, Gennady Pekhimenko
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
公众号二维码