近日,阿里云机器学习PAI主导的论文《机器学习访存密集计算编译优化框架AStitch》入选国际顶会ASPLOS 2022,论文通过编译优化的手段来自动化地提高机器学习任务的执行效率。此次入选意味着阿里云机器学习平台PAI自研的深度学习编译优化系统达到了全球业界先进水平,获得了国际学者的认可,展现了中国机器学习系统技术创新在国际上的竞争力。
ASPLOS是计算机系统领域的优异国际会议,涉及体系结构、编程语言和操作系统等多个方向,尤其重视不同方向之间的交叉,该会议曾推动了多核处理器、虚拟化、RAID、RISC、深度学习处理器等计算机系统领域的核心创新,在学术和工业界都有巨大的影响力。
访存密集型的计算已经成为当今流行的机器学习模型的一个性能瓶颈;然而,业界先进的工作中(TVM、XLA等),由于其自动生成代码的能力有限,难以针对复杂的访存密集算子子图进行高效的代码生成。
针对这一问题,AStitch提出了一种大粒度计算融合的编译优化手段,通过计算图的依赖关系特性、GPU多层次存储架构上的数据局部性、以及不同数据尺寸之下的线程并发性等三个方面的联合考虑,自动化地为大粒度的复杂访存密集算子子图生成高效的GPU代码,从而大幅减少GPU kernel调用及框架层算子调度的额外开销,避免不必要的重复计算,大幅减少片外访存,同时适配各种数据尺寸以得到非常好的并行效率。对比XLA,AStitch最高可以取得2.73倍的性能加速。
AStitch技术已经集成在阿里云机器学习PAI的编译优化引擎中,为阿里巴巴集团内部广泛的机器学习模型带来了性能加速,在集团内部GPU集群上,AStitch每周可以为数万的机器学习任务节省总计20000小时以上的GPU使用时长。
阿里云机器学习PAI已经被广泛应用于各行各业,实现企业自主可控的AI方案,全面提升机器学习工程效率。
