撰文|袁进辉
1
为什么讨论Pathways?
近两年 TensorFlow 被斜刺里杀出的 PyTorch 打了个措手不及,整个行业都在期待 Jeff Dean 力挽狂澜,祭出一记大杀器扭转乾坤。
去年,Jeff Dean 发了一篇有点万众瞩目的博客(https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/)。一方面,写博客的人是 Jeff Dean,另一方面,博客的标题《Pathways: 下一代AI架构》耸人听闻,毕竟 Google 和 Jeff Dean 引领了互联网时代和人工智能时代基础架构的发展,别人嘴里说出来“下一代架构”可能是吹牛,Jeff Dean 说出来你不得不重视。这篇博客罗列了 Jeff Dean 对深度学习算法未来演进的判断:多模态、稀疏、动态路由等,这样特征的负载确实会对深度学习底层架构带来影响。
前几天,谷歌先是在arXiv上传了题为《Pathways: asynchronous distributed dataflow for ML》的论文(https://arxiv.org/abs/2203.12533), 介绍了行业期待已久的Pathways系统的设计理念。没过几天,谷歌又上传了基于 Pathways 系统训练的5400亿参数的大模型《PaLM: scaling language modeling with Pathways(https://arxiv.org/abs/2204.02311)》。
说实话,Pathways 论文不是很好懂。我认为这有几方面原因:首先,全世界范围内只有极少数从业者在从事分布式深度学习系统的研究和开发,即使在这个群体里,也只有极少数人曾经思考过文章中讨论的问题,如果不理解 WHY,就很难理解 HOW;其次,这篇论文写得并不好。我觉得,如果一个团队写出来的论文很容易被接收,论文质量就得不到保障,Jeff Dean 的论文就属于这一类,我认为当年TensorFlow的论文(https://www.usenix.org/system/files/conference/osdi16/osdi16-abadi.pdf)也很一般,但就是能中 OSDI。几年前,我就是看了 TensorFlow 论文才有信心做一个更好的深度学习框架。
我可能是比较能理解 Pathways 的少数人之一。Pathways 讨论的这些问题我们都思考过,好几年前就思考了,而且还研发出 OneFlow,也写过论文探讨这些问题。遗憾的是,我们在论文里讨论这些问题时,几乎从来没被理解过。今天,从 Google 论文讲出这些道理,这就是不需要证明的真理了。
刚才提到,这篇论文不太容易理解,我来解释一下,希望能把论文里令人费解的事讲明白。尤其当我们理解了 WHY 之后,可能会发现 Pathways 的设计和实现仍欠火候,如果可以更进一步,那恰恰是 OneFlow 早已实现的思路。所以,我在朋友圈开玩笑说要写一篇博客,题目是《Pathways: 向前一步是 OneFlow》或《Pathways 的尽头是 OneFlow》。我计划分两篇文章来讨论 Pathways,第一篇主要讨论背景,也就是 Pathways 的 design motivation,第二篇讨论 Pathways 的设计和实现。
论文花了不少笔墨在讨论 single-controller 和 multi-controller 的问题,没有思考过这个问题的人可能会觉得有点摸不着头脑:“SPMD, MPMD, single-controller (client) 和multi-controller (client) 是什么意思?有这么重要吗?这个话题明明很无聊啊”让我们在这篇博客里把这个问题讨论清楚。论文先讨论了后出现的multi-controller,后讨论先出现的single-controller,像“倒叙”一样,我们还是按照时间轴来介绍。
2
TF v1 和 single-controller
先来介绍一下TF v1里的一些基本概念,用户通过一个 Python 客户端(client)来描述计算图(包括神经网络的结构和可使用的资源信息),然后这个计算图通过众所周知的 session.run 交给运行时(runtime)系统, 运行时系统分成一个 master 和若干