谷歌宣布了新的云TPU VMs,它通过提供对TPU主机的直接访问,使用户比以往更容易使用谷歌业界领先的TPU硬件,为在云TPU上开发和部署TensorFlow、PyTorch和JAX提供了新的和改进的用户体验。云TPU虚拟机不是通过网络远程访问云TPU,而是让用户在每个TPU主机上建立自己的互动开发环境。
现在你可以使用单个TPU VM逐行编写和调试一个ML模型,然后在云TPU Pod片上将其扩大,来利用超快的TPU互连。你对你创建的每一个TPU虚拟机都有root权限,所以可以安装和运行任何你希望与TPU加速器紧密循环的代码。你可以使用本地存储,在输入管道中执行自定义代码,并更容易将云TPU整合到你的研究和生产工作流程中。
除了云TPU与TensorFlow、PyTorch和JAX的集成,甚至可以通过虚拟机上新的libtpu
共享库来编写你自己的集成。
Cohere联合创始人兼首席执行官Aidan Gomez:
"对TPU虚拟机的直接访问完全改变了我们在TPU上的构建能力,并极大地改善了开发者的体验和模型性能。"
新的云计算TPU架构
在以前你只能远程访问云TPU。创建一个或多个虚拟机,然后通过网络使用gRPC与云TPU主机进行通信。
相比之下,云TPU虚拟机在直接连接到TPU加速器的TPU主机上运行,如下图所示。
这种新的云TPU系统架构更简单、更灵活。除了主要的可用性优势外,还可能实现性能提升,因为你的代码不再需要在数据中心网络中进行往返,才能到达TPU。此外,还能显著节约成本。如果你以前需要一队强大的计算引擎虚拟机来向云TPU Pod片中的远程主机提供数据,现在可以直接在云TPU主机上运行数据处理,不再需要额外的计算引擎虚拟机。
客户的评价
自去年10月以来,早期访问客户一直在使用云计算TPU虚拟机。从那时起,一些研究人员和工程师团队就已经密集地使用它们。以下是他们的评价:
Alex Barron是Gridspace的首席机器学习工程师。Gridspace提供了一个开箱即用的产品,用于实时观察、分析和自动处理100%的语音通话。该公司的软件为 USAA、Bloomberg 和 Square 等领先公司的语音操作提供支持。他表示:
"在Gridspace,我们一直在使用JAX和云TPU虚拟机来训练大规模的语音和语言模型。这些模型为我们最大的联络中心客户的高级分析和自动化能力提供动力。我们看到,在同样大小的TPU上训练运行时,比以前的云TPU产品立即提高了2倍的速度,并且能够在不改变代码的情况下扩展到32个主机的v3-256。我们对云TPU虚拟机的功能和易用性感到非常满意,我们期待着在未来继续使用它们。"
James Townsend是伦敦UCL皇后广场神经学研究所的一名研究员。他的团队一直在使用云TPU虚拟机上的JAX,将深度学习应用于医学成像。他说:
"谷歌云TPU虚拟机使我们能够以最小的实施复杂性从根本上扩大我们的研究规模。从在单个TPU设备上实现模型和调试,到多设备和多主机(pod规模)的训练,都有一个低摩擦的途径。在这种规模下,这种易用性是独一无二的,在研究的可能性方面对我们来说是一个游戏改变者。我真的很高兴看到这项工作能够产生的影响。"
Patrick von Platen是Hugging Face的研究工程师。Hugging Face是自然语言处理(NLP)技术的开源供应商,也是流行的变形金刚库的创建者。通过Hugging Face,研究人员和工程师只需几行代码就可以利用最先进的NLP模型。他说:
"在Hugging Face,我们最近将JAX与TensorFlow和PyTorch一起整合到我们的Transformers库。这使得NLP社区能够在云TPU虚拟机上有效地训练流行的NLP模型,如BERT。使用单一的v3-8,现在可以在一天之内使用高达2048的批处理规模预训练一个基本大小的BERT模型。在Hugging Face,我们相信,提供对云TPU虚拟机的便捷访问将使NLP社区的更多人有可能对大型语言模型进行预训练,包括小型初创公司以及教育机构。"
Ben Wang是一名独立的研究人员,致力于研究基于Transformer的语言和多模态应用模型。他已经发布了开源代码,用于在云端TPU虚拟机上训练大规模的转化器,以及用Ray在几个云端TPU虚拟机上协调训练。他评价:
"云TPU虚拟机上的JAX实现了对TPU的高性能直接访问,以及建立非常规训练设置的灵活性,例如使用Ray在可抢占的TPU pod片上进行管道并行训练。"
Keno Fischer是Julia编程语言的核心开发者,也是Julia Computing的联合创始人,他领导着一个将机器学习应用于科学建模和模拟的团队。他是Julia编译器重要部分的作者,包括Julia的原始TPU后端。
"新的TPU VM产品对于TPU在云上的可用性来说是一个巨大的进步。由于能够直接利用TPU硬件,我们不再受到中间网络连接的带宽和延迟的限制。这在我们的工作中至关重要,因为机器学习模型经常直接与在主机上运行的科学模拟相结合。"
Julia团队正在使用新的libtpu
共享库进行第二代云TPU集成。