BladeDISC: 支持动态Shape的深度学习编译器实践
邱侠斐
中文演讲 2022-07-30 16:00 GMT+8 #ai随着深度学习的不断发展,模型结构在不断演进的同时,底层计算硬件技术更是层出不穷。对于广大开发者来说,不仅要考虑如何在复杂多变的场景下有效的将算力发挥出来,还要应对计算框架的持续迭代。深度编译器是业界在实践中探索出的解决以上问题的技术方向,它让开发者仅需专注于上层模型开发,降低手工优化的人力成本,并进一步压榨硬件性能。
业内已经出现了以TensorFlow XLA、Apache TVM为代表的一批深度学习编译器,但是随着深度学习的发展,也来越来越多的模型呈现出了动态Tensor Shape的特性,例如,1)可以支持多种尺寸图片输入的CNN模型;2)以Bert为代表了NLP模型在BatchSize和SequenceLength维度均有动态性。然而前述编译器最初是面向静态Shape场景设计,需要输入及中间Tensor每个维度上具有固定的尺寸,因此并不能很好地适配动态Shape的场景。面临业务中的强需求,深度学习社区也出现了TVM Relay VM的工作。本演讲主要介绍阿里云PAI团队以BladeDISC为中心,在动态Shape编译器上的工作,主要包括如下内容:
- BladeDISC的主要架构:为何以及如何基于MLIR框架搭建支持动态Shape的编译器,为什么BladeDISC选择MHLO作为接入层Dialect,以及BladeDISC在为MHLO Dialect做了哪些改造。相较于Apache TVM,在技术路径上有哪些差异以及背后的设计考量。
- 动态Shape带来的挑战:许多静态Shape语义下比较确定性的问题,例如指令层的向量化,codegen模版选择,是否需要implicit broadcast等等在动态shape场景下都会面临更大的复杂性。部分在编译时期的决策需要移动到运行时进行。
- 大粒度算子融合:在动态Shape的语义下,如何实现通过低访存开销的硬件(Shared memory in GPU, Memroy Cache in CPU)实现更大尺度的算子融合。以及如何在未知Shape的情况下,推到Tensor之间在Shape上的约束关系进一步实现性能提升。
- 计算密集型算子:如何利用厂商库、Apache TVM算子生成、手写算子等手段提升计算密集型算子效率,以及非编译生成的算子如何在架构上实现与编译器的统一与互补。
- 如何支持多种深度学习框架,降低终端用户的使用门槛,以及BladeDISC在阿里云业务中的应用。
Speakers:
邱侠斐: 阿里云计算有限公司, 高级技术专家, 阿里云PAI团队是负责阿里集团内部和外部AI基础设施的团队,模型系统优化一直是团队的重点技术方向之一。编译器作为系统优化的重要手段,经过内部多年的沉淀打磨,目前已经在GitHub开源(https://github.com/alibaba/BladeDISC)。
团队其他分享包括:
- PAI-Blade 更易用、更可靠地使用 TensorRT Easier-to-use and More Robust TensorRT via PAI-Blade:https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41395/
- 阿里云基于 AI 编译器的通用透明性能解决方案:https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41073/
- AStitch: Enabling A New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures", Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu, Feiwen Zhu, Wenyi Zhao, Xiaoyong Liu, Jun Yang, Jidong Zhai, Shuaiwen Leon Song, and Wei Lin. ASPLOS 2022
- 模型优化(编译技术)介绍及极光广告系统最佳实践分享:https://yunqi.aliyun.com/2021/agenda/session147