研究组论文Enabling Tensor Language Model to Assist in Generating High-Performance Tensor Programs for Deep Learning被中国计算机学会推荐A类国际会议 OSDI 2024(18th USENIX Symposium on Operating Systems Design and Implementation)录用!本届会议共收稿282篇,录用49篇,录用率17.4%。
翟祎, 杨思嘉, 潘柯宇, 张任伟, 刘硕, 刘超, 叶子纯, 吉建民, 赵捷, 张昱*, 张燕咏*. Enabling Tensor Language Model to Assist in Generating High-Performance Tensor Programs for Deep Learning. 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), Santa Clara, CA, USA, July 10-12, 2024.
代码见:https://github.com/zhaiyi000/tlm
论文简介见微信推文,并已申请发明专利。
北京时间2024年7月10日线下在OSDI2024进行论文汇报和交流。
论文摘要:
高效获得高性能张量程序⼀直存在挑战。倾向于提高搜索效率的方法通常通过启发式约束来限制探索空间,该类方法往往缺乏普适性;而追求高性能的方法通常会扩大探索空间,但所采用的探索策略却不够高效。
本文提出了⼀种用于深度学习应用的张量程序搜索框架。其核心思想是在确保高性能的同时,借助语言模型进行更高效的搜索,以便高效生成高性能张量程序。论文将张量程序的搜索任务转化为语言模型生成任务。为此,设计了⼀种语言模型友好的张量语言,用于记录决策信息进而表示张量程序。在编译目标工作负载时,张量语言模型(TLM)结合离线学习的知识和之前的决策,通过概率采样的方式在当前决策空间中选择最佳决策。
相比于先前常用的随机采样方法,这种方式更为高效地探索搜索空间。TLM在效率和性能方面均表现出色。与完全调优的Ansor和MetaSchedule相比,TLM在匹配其性能的同时,编译速度提高了61倍。在与Roller相同的编译时间下,TLM的性能提升了2.25倍。