TPU的起源，Jeff Dean综述后摩尔定律时代的ML硬件与算法

过去十年我们见证了机器学习的显著进步，特别是基于深度学习的神经网络。机器学习社区也一直在尝试构建新模型，用于完成具有挑战性的工作，包括使用强化学习，通过和环境进行交互的方式完成难度较大的任务，如下围棋、玩电子游戏等。

机器学习对算力的需求无疑是庞大的，从计算机视觉到自然语言处理，更大的模型和更多的数据往往能够取得更好的性能。在摩尔定律时代，硬件进步带来的算力增长尚且能够满足机器学习的需求，但当摩尔定律被榨干后，怎样让硬件中的算力资源被机器学习模型充分利用成了下一个需要探讨的问题。

文章地址：https://arxiv.org/abs/1911.05289

在这篇文章中，Jeff Dean 首先讨论了机器学习的进步对计算设备的影响，特别是在后莫尔定律时代，他详细介绍了谷歌开发张量运算单元的初衷与效果。与此同时，Jeff Dean 还讨论了深度学习如何能够帮助设计电路与芯片。最终，Jeff Dean 描述了将来机器学习发展的方向——更大尺度的多任务模型，它能够根据任务要求选择性地激活，而且可以有更动态的调整过程。

整体文章比较长，机器之心简要概括了不同章节的主要思想，感兴趣的读者可以在 arXiv 上查阅原文章。

摩尔定律、后摩尔定律和机器学习的计算需求

自 20 世纪 60 年代至 90 年代，基于深度学习和人工神经网络的诸多关键思想和算法已经出现，并且在 80 年代末至 90 年代初，人们开始意识到神经网络能够以有趣的方式解决一些问题，并且它们的实质性优势在于既可以接受原始形式的输入数据并且在训练模型执行预测任务的过程中可以自动增强模型的层次表征，因而激发了 ML 和 AI 社区的研究兴趣。

但是，这种方法受限于当时不够强大的计算机性能。一些研究试图利用并行算法来扩展训练神经网络的计算量，但在大多数情况下，ML 和 AI 社区的关注重点已经不在基于神经网络的方法了。直到 2005 年以后，在摩尔定律推动计算性能提升的 20 年后，计算机性能才强大到可以训练大型神经网络来解决 ImageNet 等现实世界的实际问题，而不再局限于 MNIST 等小规模问题。其中，GPU 卡上的通用计算范式使得神经网络开始在解决棘手的实际问题上展示有趣的结果。

摩尔定律提出者戈登·摩尔。

然而令人遗憾的是，在计算性能开始足够强大到可以解决有趣的现实问题，并且机器学习规模的扩大和适用性的提升也迫切需要更多的计算资源来解决更大的问题，这时计算机行业却随着通用 CPU 性能的持续提升出现了整体性的大幅度放缓。下图 2 展示了近 40 年间的计算性能增长态势，其中自 1985 年至 2003 年，通用 CPU 性能每 1.5 年提升一倍；自 2003 年至 2010 年，通用 CPU 性能每 2 年提升一倍；而 2010 年以后，通用 CPU 性能预计每 20 年才能提升一倍。

图 2： 摩尔定律和后摩尔定律时代的计算需求增长态势。

下图 3 展示了近年来出现的一些重要的机器学习网络架构，如 AlexNet、GoogleNet、AlphaZero等，它们对计算的需求呈现越来越大的趋势。

图 3： AlexNet、GoogleNet、AlphaZero等重要的机器学习网络架构以及它们的计算需求增长态势。

下图 4 展示了 2009 年以来，机器学习领域 Arxiv 发表论文数量的增长情况，其中 2018 年的论文数量是 2009 年的 32 倍（论文数量每 2 年提升一倍以上）。现在，每天仍有 100 多篇论文在 Arxiv机器学习子板块发表，并且这一增长态势没有出现放缓的迹象。

图 4：自 2009 年以来，机器学习相关 Arxiv 论文发表数量的增长态势（蓝）和摩尔定律增长率（红）。

针对机器学习的硬件

早在 2011 年左右，谷歌研究者与系统工程师就构建了 DistBelief 这个早期分布式系统，在 12 年早期的一些深度神经网络中，模型使用 DistBelief 能大大提高准确性。然而，由于计算需求太大，要使用基于 CPU 的计算设备继续提升深度模型，谷歌数据中心的计算机数量需要翻一倍。为此，这样的现实驱动了谷歌开始考虑为神经网络构建专用硬件，用于推理乃至训练。

为什么深度学习模型需要定制化硬件

深度学习模型有三大属性，它们使得其与很多通用计算都不相同。首先，深度学习允许降低计算精度；其次，大多数模型执行的计算都是由少量运算组成，例如矩阵乘法、向量运算以及卷积运算等等；最后，过去 40 年开发了很多机制以令通用程序能在 CPU 上高效运行，但这些机制在机器学习运算中都是不必要的。

构建能进行密集低精度线性计算的硬件是非常有意义的一件事，为此，谷歌决定着手设计名为「张量处理单元」的加速器，从而加速深度学习的推理与训练。

为了能执行推断，Int8 的精度已经显示出足够强大的能力，目前研究界也在开展进一步的工作，以便采用更低的精度和更稀疏的权重来突破这一界限。TPUv1 的单个核心是由 65536 个 8 位乘加矩阵运算单元，其峰值运算量为 92 TeraOps/s。TPUv1 平均比现代 CPU 或 GPU 快 15 到 30 倍，平均能耗降低 30 到 80 倍。目前，这种芯片约占谷歌数据中心神经网络推理需求的 95%。

在低功率移动端设备上做推断也非常重要，它可以极大降低网络延迟并提升隐私保护。谷歌采用与 TPUv1 相同的设计原理，应用于更低功耗的环境而设计出 Edge TPU，它以 2W 的功耗提供了 4TOPS 的计算力。这样的边缘设备可以更好地在农业和日常生活中应用机器学习模型。

图 5：谷歌张量运算单元 v2（TPUv2）的结构图。

前面 V1 只能做推断，但要设计定制化的机器学习训练硬件要复杂得多。主要原因在于单块训练芯片无法在合理的时间内解决大部分训练问题，因为单块芯片无法提供足够的算力。因此设计训练系统实际上就是设计大规模的整体计算系统，它需要考虑各个加速芯片间的高性能互联，从而形成及紧密耦合的超级计算机。

谷歌的第二代和第三代 TPU 旨在支持训练和推理，TPUv2 和 TPUv3 每一个 Device 包含四块芯片，不同的 Device 整合在一起就形成了一个整体 Pod。图 5 展示了具有两个核心的 TPUv2，每个核心的主要计算能力有大型的矩阵乘法单元提供，其每个时钟周期能计算 128×128 的矩阵运算。

图 6： 谷歌 TPUv3 Pod，它包含 1024 块 TPU 芯片。

低精度数值运算

TPUv2 和 TPUv3 使用一种名为 bfloat16 的定制化浮点运算格式，该格式与 IEEE 半精度 float 16 不同，它对机器学习更有效。自 2015 年以来，它一直是 TPU 的主要精度，英特尔 2018 年也宣布计划为下一代处理器增加 bfloat16 的支持。下图展示了 IEEE fp32 单精度浮点格式、IEEE fp16 半精度浮点格式和 bfloat16 格式，它们的划分与数值范围都展示在内。

图 7： 三种浮点精度。

  因为 bfloat16 格式的乘加器需要最少的电路，因此有可能在相同的芯片面积和功率预算中配置更多的乘法器，这能令芯片提供更强的算力与更低的能耗。与此同时，降低到 16 位精度还能减少内存的占用与权重的体积，从而进一步提升模型的实用性。

ML 专用芯片设计的挑战

目前机器学习研究领域正快速发展，而如今芯片设计项目通常需要 18 个月到 24 个月才能完成设计、制造与部署。为了节省成本，一般这些硬件至少需要工作三年。因此，构建 ML 硬件的计算机架构师需要预测 2 到 5 年内机器学习快速发展的领域。Jeff Dean 的经验是，将计算机架构师、高级软件系统构建者与机器学习研究者汇聚一堂，共同设计相关的主题，例如，「那个时期可能需要的硬件是什么，有什么有意思的研究趋势，它们需要什么样的 ML 硬件」。

机器学习在硬件领域中的其他应用

机器学习设计芯片

机器学习另一个可以和硬件领域结合的方向是进行自动化的芯片设计，如设计定制化的 ASICs。由于芯片领域目前已有大量的人力投入，因此贸然改变整体的芯片设计是不太可能的。然而，在芯片的布局和布线方面机器学习是可以参与的。例如，可以将一个序列的布局和布线结合起来，用一个整体的评价标准——如芯片区域大小、耗时和走线长度进行评价。

如果有强化学习算法去「玩」这样一个布局布线的游戏，则可以产生很多不同的 ASIC 设计，或只有一种特定的 ASIC 设计。这只需要设计一个奖励函数——其中包括不同的奖励属性。这样有可能能够更快且更高效地设计出更好的布局布线。现在，谷歌内部已经在尝试这些方法，而且已经有了一些初级但有效的成果。此外，针对不同的目标优化评价指标，奖励函数中的属性也可以进行调整，使得机器学习算法设计出满足不同需求的芯片。

机器学习解决半导体制造问题

在半导体制造的过程中，对晶片的视觉质量检查是可以转移到机器学习上的，可以让这一过程变得更自动化，或者可以提升现有方法的精确度。通过更早或更好地发现瑕疵，制造企业可以提升生产率、降低成本。

计算机系统中的学习型启发式算法

另一个机器学习可以参与的领域则是使用学习式的启发法，应用于计算机系统中，如编译器、操作系统、文件系统、网络堆栈等。计算机往往需要很多手工设计的启发式算法用于在系统运行中寻找解决方法（如资源调用等）。

现在，机器学习也许可以替代这些算法了。基于机器学习的启发式算法能够考虑更多的语境信息，使得系统的运行方式更符合当前的运行模式，而不是根据一般情况设计的方法。其他机器学习可以替代的地方包括对哈希表、B 树等传统数据结构的替代。机器学习可以获得这些数据的真实分布，并使用系统进行处理，使得数据量减少、效率提升。

未来的机器学习发展

将来的机器学习模型可能和现在有一些明显区别。将来的机器学习系统可能是在大规模机器学习加速硬件上运行的，而且单一模型可以被训练用来完成上千甚至上百万的任务。该模型由不同的组件和结构组成，样本之间的数据流动可能是动态的，每种样本都不一样。模型可能使用类似于「稀疏权重门（sparsely-gated）」结构，混合了专家知识和学习路径，具有很强的能力。但是对于给定的任务或样本，模型只会激活其中一部分。

图 8：描述这一大型、具有稀疏权重且可以完成多任务的模型。

后摩尔定律时代，简单地压榨硬件算力不一定能够让机器学习再进步了。相反，通过设计专业的机器学习硬件设备，让机器学习解决现有的硬件设计、制造和运行方面的痛点才是出路。将来的机器学习模型可能更大，但是会朝着多任务的方向继续发展。