Sheng Xu's Blog

2022/08/10，许晟 加了三个月的班，我们的Trainium芯片和AWS上的trn1 instance终于要面世啦！毕业后加入AWS Annapurna labs的这一年半，我学到了很多，也越来越喜欢这份工作。最近经常和朋友们聊起AI芯片和编译器这个飞速发展、激动人心的领域。在这里总结记录一下这一年的收获和感想。 1/ 我们实验室是做什么的 我们是帮AWS(Amazon Web Services)搞芯片的，又快又省钱的芯片。 Annapurna labs在2011年由几个以色列工程师创立，一开始以做CPU芯片为主。2016年被Amazon收购以后，就一直是AWS的”急先锋”。先是和EC2的团队合作，推出了AWS Nitro这个虚拟化系统，把EC2上所有实例都大大的加速。然后又推出了Graviton系列芯片，目前仍然是AWS上price performance最好的CPU芯片。Price performance就是用价格除以性能，在单位性能下，价格越低越好。对于那些每天都得跑成千上万次计算的大公司，用我们的芯片能让他们每年省下几百万甚至几千万美金。 最近十年AI崛起，掀起了一波机器学习的潮流。尤其是深度学习这个分支，应用到了我们生活的方方面面，比如各种app和网站的推荐算法，语音识别，视觉识别等等。深度学习模型的大小也在爆炸性增长，从几年前的几百万参数，到最近新出模型的几百亿参数。这些都对底层的计算平台提出了很高的要求。所以这些年，许多科技大厂都开始搞自己的AI芯片，各种芯片startup更是如雨后春笋。 AWS坐拥世界上最大的云计算平台，有极好的用户基础，当然也不甘落后。2018年，在AWS的AI战略布局下，Annapurna labs开始转型做机器学习芯片。从2019初的AWS re:Invent宣布，到2020年正式推出Inferentia芯片和AWS inf1 instance，是当时云端上深度学习inference price performance最好的instance。Inferentia是专门为深度学习设计的芯片。这种专门为一个领域设计的芯片叫ASIC。ASIC牺牲了普适性，但是收获了高性能。 这里提一句，深度学习分为training和inference两步。深度学习模型先经过training把所有误差降到最低，然后就可以开始投入使用，也就是开始inference。training的过程是很难很慢的，大的模型需要训练几天甚至几周。好在过程是一次性的，一个train好的模型可以被不断的被使用。与此相比，大的商用模型每天要跑上万次inference。所以一般来说，inference才是公司支出的大头。Inferentia芯片就是专攻inference。 Inferentia面世不久就收获了许多大大小小的客户。在Amazon内部，我们最大的客户是Alexa，用我们的芯片把Alexa加速了25%，成本降低了30%。在Amazon外部，我们最开始的客户包括Airbnb, Snapchat, Pinterest等，后来又逐渐吸纳了Bytedance, Facebook, Apple等大公司。 Inferentia获得成功之后，团队又开始马不停蹄的开始下一代芯片Trainium的开发。Trainium不仅在inference上面比Inferentia强了许多，而且支持training。经过了两年的努力，Trainium芯片和相应的trn1 instance终于要正式上线了！ 这里必须提一下业界老大哥Nvidia英伟达。Nvidia从几十年前就开始造GPU，经验丰富。这些年他们接连推出T4, A10G, A100等针对机器学习的GPU。根据2020年的市场调查，Nvidia在AI芯片的市场占比在80%以上。有趣的是，Nvidia和AWS其实是紧密的合作伙伴，T4/A10G在AWS上叫G4dn/G5 instance, A100在AWS上叫P4d/P4de instance。 我们芯片的原理和GPU完全不同，不过我们必须得比GPU的性能更好才能吸引到顾客。我们的Inferentia在最常用的那些模型上比G5平均快了30%，便宜了50%。 而新推出的Trainium，跟A100在GPT模型上一样快，便宜了50%。比A100在BERT模型上快了50%，便宜了2.5倍（BERT和GPT分别是目前最常用和最火爆的深度学习模型）。鉴于A100是目前市面上最快最强的AI GPU，所以Trainium也算是跟上时代了！ 另外，光有好性能还远远不够，顾客的体验和易用性同样的重要。对此，我们有一整套软件的服务。顾客从GPU instance迁移到inf1/trn1上，只需要在tensorflow/pytorch里面加不到五行代码。 当然，比起Nvidia的GPU、Google的TPU，我们在普适性以及软件工具链的成熟度上都还有很长的路要走。可是考虑到Nvidia和Google起步比我们早了几年，而且在这个领域都有上千人的团队；相比来说，Annapurna labs上下不足200个工程师，我所在的编译器团队更是缺人手，只有不到20人。只用了4年的时间走到这一步，我觉得我们很棒了！ 可以说，比起不少芯片的初创公司，Annapurna labs目前发展的算是非常顺利了。除了团队的不懈努力外，必须得感谢AWS这个独一无二，全球第一的云计算平台，为我们提供了源源不断的客源。后续我们的目标是继续以2年一次的频率推出更强的芯片。 2/ 我的团队 – 深度学习编译器 什么是compiler（编译器）呢？简单来说，编译器就像是一个翻译机，把程序员写的代码翻译成底层芯片能懂的语言（也叫ISA)。有趣的地方在于，一份代码有多种翻译的方式，不同的方式生成的结果在性能上天差地别。 所以编译器有两大功能： 翻译。这里得考虑到芯片的各种限制，生成的ISA必须是可运行的 优化。优化做得好，能让性能有几十倍乃至百倍的提升 在传统的编译器领域，最著名的编译器就是LLVM。LLVM往上能接收几乎所有计算机语言，把它们翻译成LLVM自己的语言（也叫IR)，往下能接通绝大部分底层的ISA。LLVM提供的核心价值是基于IR的一整套复杂优化。…