为何Google当初决定自行开发TPU芯片?Google TPU主要设计者亲自披露原因

现今,Google许多服务,几乎都跟AI有关,但凡是搜索、地图、照片和翻译等等,这些AI应用服务,在训练学习和推论过程中,都使用到了Google的TPU。Google很早就在数据中心内大量部署TPU,用于加速AI模型训练和推论部署使用,甚至不只自用,后来更当作云计算服务或提供第三方使用,还将它变成产品销售。

在今年线上台湾人工智能年会上,Google研究部门软件工程师Cliff Young担任第一天的主题演讲,Cliff Young不只是Google Brain团队核心成员,更是Google TPU芯片主要设计者,一手设计和打造TPU,部署到Google数据中心内,作为AI硬件加速器,用于各种AI模型训练或推论。在加入Google前,他曾在D. E. Shaw Research和贝尔实验室,负责设计和建造实验室超级计算机。在整场演讲中,他不只亲自披露Google决定自行开发TPU的过程,针对深度学习革命对于未来AI发展影响,也提出他的最新观察。

Cliff Young表示,深度学习神经网络技术自2009年开始在语音识别大放异彩以来,几乎每年在不同领域应用上,我们都能看到因为深度学习而有了突破性的发展,从AI图片识别、Al下棋、到Al视网膜病变判读、语文翻译、机器人拣货等等,“这是我们以前从未想过的事。”

正因为,深度学习的出现,让人类在不同领域执行任务的方式产生重大变化,他也以美国知名科学哲学家Thomas Kuhn提出的科学革命的发展模式来形容,深度学习本身就是一种科学革命的典范转移,不只是常态科学。

Thomas Kuhn在《科学革命的结构》一书中提出两种科学发展模式,第一种是常态科学的模式,通过实验和证明来理解新事实的方法,当出现旧科学无法认同的新科学产生时,就会产生新旧科学之间的冲突,衍生发展出另一种模式,也就是第二种的科学革命的模式,在此模式下,新科学模式将彻底颠覆旧科学的做法。 “我认为深度学习革命就是这样一种转变,正在取代传统计算机科学。”Cliff Young说道。

更进一步来说,他指出,深度学习是一种数据驱动的决策过程,不同于传统的stored value或启发式(heuristic)决策方法,深度学习算法使用可观测的数据,来提供人类创建更好地决策的方式,比如运用在用户推荐,可以根据用户轮廓或网络行为,来推荐适合的产品或给出最佳搜索结果。

但他也坦言,不像数学原理可以被解释,深度学习模型运行原理目前仍难以解释,也因此,科学家无法从这些为何可行的原因中,找到提高效率的更好的做法。但如果想完全理解和解释深度学习运行原理,依照以往工业革命的发展经验,得等到合成神经动力学的出现后,才有机会得到解释,动辄可能耗费数十年之久,所以,他也说:“从事深度学习研究,比起问why,how更重要。”

Cliff Young回顾机器学习革命过程,可以2012年的AlexNet神经网络架构作为分水岭,由Alex Krizhesky等人提出的AlexNet运用GPU创建的深度学习模型,以85%准确度刷新世界纪录,在当年ImageNet图像分类竞赛中一举夺冠。

这项竞赛后来也引起Google的高度关注,认为深度学习技术大有可为,便开始投入研究。但他们投入后发现,深度学习模型在图像识别和分类的成效表现,高度依赖GPU的浮点运算能力,需要消耗大量计算资源供AI模型做学习训练,因模型训练运用GPU衍生出的运算成本十分昂贵。所以,Google才毅然决定自行开发深度学习专用的处理器芯片,也就是TPU(Tensor processing unit)。

在投入深度学习研究3年后,2015年时,Google开发出第一代TPU处理器,开始部署到自家的数据中心,用于深度学习的模型训练。

Google在2016年Google I/O大会首次披露TPU,与当时的CPU与GPU相比,Google的TPU不仅能提供高出30~100倍的每秒浮点运算性能,在整体运算性能上也有多达15到30倍的提升,甚至在性能/功耗比获得有将近30~80倍的改善,Cliff Young表示,TPU很可能是当时世上第一个实现以高内存容量的矩阵架构设计完成的处理器。

当年,Google击败韩国棋王李世石的AI计算机围棋程序AlphaGo,背后功臣就是使用TPU运算的服务器机柜,Google Deepmind团队在AlphaGo中采用48个TPU用于AlphaGo的AI下棋推论,与人类进行棋力比赛。

到目前为止,Google TPU一共历经4代发展演进,从初代TPU仅能应用于推论,到第二代TPU开始加入深度学习模型训练处理能力,对于网络吞吐量需求增高,而随着运算能力的提升,考虑到散热问题,所以新一代TPU开始在散热机制上结合液冷设计,也就是第3代TPU,也因此增加TPU密度,到了最新第4代TPU,则推出无液冷的TPU v4i和采用分布式液冷的TPU v4两种不同版本。

这几年,深度学习硬件加速器越来越红火,Cliff Young认为,在AI训练和推论硬件发展上将出现转变。他预测,未来推论硬件设计上将更具多样性,来发展出不同推理解决方案,来对应不同场景的使用需求,从微瓦的超低功耗,到高性能计算HPC以及超级计算机应用。

另在AI训练硬件方面,他表示,融合式硬件架构将成为发展主流,尤其现在不少新推出的AI训练硬件,都有不少相似之处,像是在设计高密度运算芯片Die时,会采用HBM(高带宽内存)集成设计,还有创建高性能互联网络,用于传输训练资料,如TPU就有使用ICI(Inter-Core Interconnect)来与其他TPU高速互联,其他还有如Nvidia的NVLink高速互联接口等。虽然这些技术,都来自不同团队,但他们都有个公用之处,就是在研究共同问题想办法找答案,他表示,通过这些技术融合,有机会可以找到好的解决方案。

Google数据中心内目前部署了许多TPU Pod集群,以上图TPU Pod为例,分上下二层创建TPU Pod集群,每层配置多台机架式机柜,每台机柜中安装了数十个TPU,包括TPU v2(上排)与TPU v3(下排),再通过网络线来与其他TPU高速互联。

另一方面,他也观察到,近几年,全球AI竞赛进到白热化阶段,虽然加速深度学习在自然语言模型的突飞猛进,但也使得需要训练的AI模型越来越庞大,像是为了完成使用1,750亿个神经参数的GPT-3文本产生器模型的训练,OpenAI使用1万个GPU创建运算集群,以petaflops算力花了3,640天来训练该模型。

为了训练出像GPT-3这样的超大AI模型,Google也以多个TPU互联创建TPU Pod集群,来打造超级计算机集群,放在自己的数据中心内加速AI模型的训练。过去几年Google TPU Pod也从一开始256个TPU、增加到1,024个,到现在一个Pod拥有多达4,096个运算节点。以上图在Google数据中心使用的TPU Pod为例,分上下二层创建TPU Pod集群,每层配置多台机架式机柜,每台机柜中安装了数十个TPU,包括TPU v2与TPU v3,再通过网络线来与其他TPU高速互联。

但想要跟上深度学习发展脚步,Cliff Young认为,不能仅靠加大训练用的运算机器,现有的软硬件架构设计也得跟着转变才行,他提出materials -application codesign协同设计的概念,认为未来深度学习架构设计,需要结合包含从物理到应用所有层面的协同设计,他认为这是打破摩尔定律瓶颈,找到深度学习发展新出路的方法。

他进一步说明,在传统协同设计中,硬件和软件之间仅靠一层单薄的ISA指令集架构作为联系沟通,但在以DSA(Domain-specific Architecture)专用领域架构为主的协同设计中,则由包含许多不同软件层,架构层,以及不同接口组合而成。其中软件层方面,包括函数库、编译器、应用程序、模型、算法、Numerics等,硬件架构方面包括物理设计、半导体材料、架构及微架构等。这些软硬件协同设计,未来可以运用到深度学习架构设计优化上,像是在内存技术中,能大幅降低模型训练过程对于位元(bits)复写使用,以及加入采用较慢的内存读取速度的设计等。