Pathways 论文精读【论文精读】

Pathways: Asynchronous Distributed Dataflow for ML
- Pathways 语言模型由 Google 发布,该模型的一大亮点是做到了 5400 亿个可学习的参数
作者在博客(https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html)中解释了一下为什么要使用这么大的模型



- 上面 3 张图展示了当参数从 8 亿开始的时候能做几个简单的任务,但是随着模型大小的增加,能做的任务就会变得越来越多
本视频主要是基于 Pathways 整个模型背后的系统进行讲解,原因如下
1、博客中所阐述的性能非常厉害的。整个模型用了 4096 块 TPU v3 的芯片,达到了 57.8% 的应用使用率
- 对于芯片来讲有一个非常重要的指标:每秒钟所能计算的浮点运算的数量,现在通常是几百 T (但是也不用特别迷信这个数字,因为从硬件制造的角度来说,这个数字可以通过硬件的堆积将这个数字增加到很高的程度,所以更加重要的是在实际任务中所能真正达到的计算的峰值)
- 这里的 57.8% 意味着能达到理论峰值的一倍以上。作为对比,在 GPU 上目前做的比较好的大概能达到 60% 或者稍微多一点。所以如果不考虑在 GPU 和 TPU 上计算利用率的细微差别的话,57.8% 这个数字已经是非常高了
2、这篇文章汇聚了 Google 的两位技术大神 Jeff 和 Sanjay
标题
“Pathways:”
- 系统文章取名最常见的方法,使用系统名字加冒号来告诉读者是个什么样的系统
“Asynchronous Distributed Dataflow for ML”
- 针对机器学习的异步的分布式的数据流
作者
Jeff Dean
- lead of Google AI
- 出名的工作包括:Bigtable,MapReduce 以及 TensorFlow 等
Sanjay Ghemawat
背景介绍
Introducing Pathways: A next-generation AI architecture
- 这篇博客由 Jeff 写于 2021 年10月
标题
介绍一下 pathways ,下一代的 AI 架构
核心思想

这个 AI 系统能够处理非常多的任务(multimodality)
但是有一点区别是系统中其实有非常多的模型,当接到一个任务之后可以通过稀疏调用来调用这些模型(就像是人的大脑一样有很多不同的区块,每个区块可能是一个比较大的深度学习模型,在接到一个任务的时候触发不同的模型,在更新的时候也可以根据特定的数据更新特定的模型,或者使用大批的数据对模型进行更新)
和现在主要关心的深度学习的模型区别主要有两点:
1、系统中可以有非常大和非常多的模型
- 整个模型的参数可以达到万亿级别或者更大的规模
- 每个模型可能处理不一样的数据
2、系统中模型的连接是比较稀疏的
- 现在一般的模型在接到一个任务之后所有的神经元都会触发,而 Pathways 系统只会触发特定的神经元
这是 Jeff 的一个早期的畅想,也可能是这个项目最终的目标,本视频所讲解的系统也叫做 Pathways 但是跟博客中所描述的还是有一定的区别
摘要
论文所讲述的是针对加速器使用的一个新的大规模的编排层的设计
- 编排层:在语音上有很多服务,当接到服务的计算任务后怎样将它映射到下面的计算资源进行调度和管理
Pathways 是被设计用来使得探索新的系统和机器学习研究中新的想法变得更加方便,同时也能达到很好的性能
导言
近年来,整个机器学习的模型是跟硬件和背后的系统一起演进的
- 算法必须要适合当前的系统和当前的硬件,如果不适用的话,就处理不了很大的数据,做不出什么效果
- 同理,硬件也要随着模型的变更来产生变革,比如说内存、带宽、计算量等
- 软件系统是将任务和硬件融合在一起,需要与两者适配,所以随着硬件的进步和模型的进步,每隔几年软件系统也需要进行更新
- 做系统方向一般不用担心失业,因为硬件和任务总是不断变化的。具体来讲就是,现在的软件系统可能与当前的这些任务过度优化,但是不能很好地适用于未来的需求(任何硬件或者软件系统的升级都是因为这个原因)
现在的机器学习任务都是用的单程序多数据(SPMD)模型,这个模型来自于 MPI
1、假设要在一台机器上跑 pytorch 模型,通常情况下,将程序所需要的包在这台机器上安装完成之后,再将代码(.py 文件)复制到这台机器上,然后再启动一个进程来运行代码,这样就完成了
2、假设这台机器上有 GPU ,而且代码确实需要使用 GPU 做运算的话,当 CPU 在运行的时候看到一个 python 的函数需要调用 GPU 来进行计算的话,CPU 就会将这个指令发送给 GPU ( PCIE 连接),等到 GPU 运算完成之后再将结果返回回来继续往下计算
3、分布式:假设要在两台机器上同时计算一个模型。其实跟在一台机器上计算没有太大的区别,在第二台机器上同时将所要的包都安装好,而且需要保证两台机器上所安装的包的版本相同,同时将运算的文件(.py 文件)也都复制一份到第二台机器上而且需要保证在每台机器上的代码是一样的。运行时,在每台机器上分别启动自己的进程来运行代码。
- 最简单的情况下,两台机器各自单独进行训练,要完成分布式的话需要做数据的交换。
- 上一期参数服务器中讲过一个最简单的分布式----数据并行,就是每个机器拿到各自的,然后在每个小批量中计算好梯度之后将每个机器中所计算得到的梯度全部加起来就完成了
- 这些梯度是当时每台机器自己的数据上运行的,需要将将这两台机器的计算结果全部相加,然后分别告诉所有的机器,从逻辑上来讲,就是将两台机器运算的结果放到另外的一个地方进行相加,然后将相加之后的结果分别返回给两台机器,这个操作在 MPI 中叫做 AllReduce ,对应的是 MPI 中的函数调用:目标是高效地将不同机器中的数据整合(reduce)之后再把结果分发给各个机器
- 当要将一个单机的代码拓展到分布式的时候,唯一需要做的就是在每一次梯度计算完成(backward)之后,调用 AllReduce 函数将每个机器计算的梯度在各个机器之间进行相加,完成这项操作之后就能够继续向下进行运算了,所以和单机的区别不是很大,只是分布式上面多加了几个函数
- 核心是每台机器上都有同样的代码,这些代码中插入了所需要的 MPI 的函数,启动时,每台机器启动各自的进程,在计算到中间的时候,所有的机器停下来做一次数据的同步,然后再各自继续进行运算,这个也叫做 BSP 模型
- 这比上一期参数服务器中所讲的要简单一些,因为深度学习早期的 CNN 甚至是早期的 Transformer 中,整个模型相对来说不是特别大,只是计算比较复杂,计算的时候比较耗时间
在通讯不多的情况下,简单的操作其实也是挺好的。但是作者说现在这些简单的 BSP 模型在新的任务上会有新的问题
1、现在的语言模型变得越来越大,参数规模都是上百亿、千亿了,可能导致整个模型在一个加速器上放不下,这是纯粹的数据并行就不行了,就需要使用一个叫做 pipelining 的方式来进行任务的切割
2、现在有的模型在进行稀疏性的探索,整个模型里面不是一个全连接,而是一个稀疏连接的过程,所以导致在通讯和任务调度上会有不同
MPI 假设的是时的一个数的通讯,一次性将所有的东西都发出去,所以 8 年前当深度学习崛起的时候,它相比于分布式来说确实是比较简单的,因为模型足够小,而且没有稀疏性,所以导致 8 年前在参数服务器上所做的各种优化其实对于深度学习来讲没有必要。但是现在模型已经足够大了,一块加速卡已经放不下了,而且在模型具有稀疏性的时候,又需要重新考虑这两点
本视频中可能会经常提到上一期所讲的参数服务器做对比,主要有两个原因
1、假设已经看过上一期,再次提及可以加深理解
2、在系统中,没有什么东西是完全新的,对于一个新的硬件、一个新的任务将其中的困难点进行拆解之后,可能跟之前所做过的东西是一样的。之所以不断地引用之前参数服务器中所讲的东西,并不是说这些东西在之前的系统已经被做过了,这篇文章只是将旧的东西重新拿来炒一炒,这篇文章处理的是一个新的任务、新的硬件环境,然后提出了一些解决方案,也许和前面的系统有类似的东西,并不是因为延用了之前文章中的一些概念,而是说这些方法是在这个特定的场景之下解决这些问题最适合的方法
先去补上一期,补完再回来接着更新
----to be continued----