Skip to content

Latest commit

 

History

History
247 lines (155 loc) · 29 KB

pathways.md

File metadata and controls

247 lines (155 loc) · 29 KB

简单理解为目前的机器学习基础设施和框架,很难支持大规模、稀疏或者不规则的模型

  1. 在解决的是什么问题? 给加速器使用的大规模(上千万的卡)调度层
  2. 为何成功,标志/准是什么?很容易就能把 JAX 上已有代码跑成大规模的
  3. 在前人基础上的关键创新是什么?sharded(类似 MP,PP,TP) dataflow graph of a asynchronous operators that consume and produce futures. efficiently gang-schedules heterogeneous parallel computations。尽管有数据面的依赖,但控制流依然可以并行执行(asynchronous distributed dataflow)
  4. 关键结果有哪些?
  5. 有哪些局限性?如何优化?还不支持吹牛的部分更新模型权重
  6. 这个工作可能有什么深远的影响?

摘要

一个单一的控制器模型(a single-controller model)让表达复杂的新的并行模式更容易。

1. 介绍

本文认为目前的 ML系统设计的太特定于当前的负载,而无法扩展到未来的需求上。例如大部分都使用 一个程序,多份数据(Single Program, Multiple Data)的模型,这是受 MPI 启发出现的,所有的加速器都运行同样的计算,步调一致,通过 AllReduce 等集合通信来通信。如今有不少是需要流水并行以及 MoE的,虽然他们能构建在 MPI 之上,但是限制太多了。

另外,随着新一代加速器的产生,ML集群越来越异构了。把大量高速互联的同构加速器岛屿直接排他地提供访问是很贵的,经常是浪费了的,因为单个用户程序需要让所有加速器持续忙碌。这些限制让推动了研究员朝“Multiple program Multiple data”(MPMD)迈进,它可以通过把整个计算分为多个子部分,然后把子部分映射到小块岛屿的加速器上。为了增加使用率, 有很多 ML 硬件资源管理领域的研究。如在多个任务间,提供更细粒度的多路复用(multiplex),让负载更弹性,提高容错。

最终,研究员开始在一个基础的模型上开始干活。这样一个基础模型上接很多不同的下游任务,希望他们一起训练。

本文提出的 Pathways,是client server 架构,让运行时能在多个 client 情况下,执行多个程序。它是第一个设计来实现透明、高效地跨多个 TPU pod 运行程序的系统,通过使用新的数据流执行模型,它可以扩展到上千个加速器。它提供的编程模型 让表达非 SPMD 计算非常容易,而且让集中式资源管理和虚拟化提高了利用率。

2. 设计动机

当前 ML 系统的限制。

目前的机器学习基础设施和框架,很难支持大规模、稀疏或者不规则的模型。

给 SPMD 模型用的分布式的 ML 系统,经常利用多控制器的架构,相同的client executable被直接运行在所有的主机上,在程序执行期间,会排他性地拥有资源。比如 MPI,PyTorch,现在的 TensorFlow 和 Flax。这种设计的核心优势是分发加速器计算时延迟很低,因为所有加速器上的代码都一样, 分发涉及的通信只需要经过相对快的 PCIe 链路。其他通信只通过集合通信,利用专门的互联方法如 NVLink 来发生,不需要经过宿主机内存。然而对于使用流水或者稀疏计算的负载,这种设计不好。任何涉及到超过集合通信的方法,都需要用户自己实现协调的原语。多控制器经常假设对硬件资源的排他性使用。这样不仅 把保证高利用率的责任推给了用户,而且让实现一些特性比如资源虚拟化和多路复用复杂化,而这些特性是构建高效集群的基础。

单控制器系统如 TensorFlow v1,提供非常普适的分布式数据流模型,包括优化了的图内控制流。一个 TF python 客户端构建一个计算图,把它交给运行时的协调者,它会把图分为几个子图给到每个worker,委托worker上的本地运行时做子图的执行。在 worker 间的协调是通过数据中心网络(Data Center Network)来进行数据和控制侧的消息传递。 尽管单控制器设计提供了更弹性的编程模型和资源的虚拟化,但是带来了设计挑战。

首先,多控制器系统只需要通过 PCIe 来分发加速器计算图,而通过DCN来分发的延迟,通常是比 PCIe 要慢多个数量级的。 其次,为了支持并发执行 MPMD 程序里的 SPMD 子计算图,每个子图都跨越一个共享集群里的部分加速器,运行时必须支持加速器上的 gang-scheduling(成组)调度. 这个对 TPU 非常关键,因为 TPU 都是单线程的,而且只能运行非抢占式内核,所以如果通信和计算如果不是用一致的顺序入队列,系统会死锁。即使对于 GPU 和其他能执行并行计算的加速器, 成组调度会有更高效的集合执行效果。因此单controller的系统,需要一个分布式的调度机制来安排不同程序的计算顺序。

最后,系统需要能在上千个加速器上执行计算,把共享表达和数据结构(shared representations and data structures)当作一等公民。例如,一个简单的表示 M 路分片计算和 N 路分片计算之间的数据流图,会需要 M+N个节点,M*N个边,很快就会笨拙、运转失灵。

从上图看到,Host 和 Device 上都有 controller,而 Host 通过 PCIe 与Device进行dispatch和通信(read wait)。而 TF1 的区别是它是单 controller,是个单独的。虽然 Host 上也有控制逻辑,但这个只是子图的,不涉及总的。所以这单个 controller 需要通过 DCN 消息与其他加速器通信,当子图执行完,需要通知 controller

TF v1 实现的选择,是专门用于单个,小规模,排他性拥有加速器的场景。这种过度导致很难用它来做当代或者未来的ML负载。虽然 TF 能够执行需要跨主机协调或者通过send 和 recv 算子的计算,host 侧的工作,比如分发加速器计算,只能在传输完成后,才能触发(?)。在设计大量跨主机传输的程序里,例如流水并行的程序,这些分发的延迟会累加,比如上图中的 step k 和 k+1 之间的空闲? 最终导致加速器利用率不高效。虽然 TF v1 的用户可以(低效地)通过使用控制边,实现单个程序上的一致顺序来做成组调度。缺乏中心的调度器,导致无法确保多个程序之间计算的一致性顺序(?). TF 会把整个分片的计算图 materialize,在shard的数量达到上千规模后,涉及引入潜在的计算图序列化和执行方面的开销,导致子计算图之间有上百万的图之间的边。

上图是 TF1 上跑 MPMD 程序的例子,每个 Host 上的跑的子图不一样,他们之间通过 Send/Receive 的 DCN 来通信

Pathways 结合了单个 controller 的灵活性和多 controller 的高性能(咋就没有对应的图了?有多 controller 的典范嘛). 选择单 controller 是因为觉得它能更好服务号创新和高效的ML计算:能利用好计算稀疏和异构,能够让集群管理系统提升分享和虚拟化资源的能力了。 我们的涉及和老的单 controller ML 系统的最大区别是使用了异步分发,来达到多 controller 的性能,支持集中资源管理和对成组调度的一等公民支持,使用了shared dataflow(啥意思?) 系统来高效协调。

3. 编程模型

提供的灵活编程模型。

我们实现了对 TensorFlow 和JAX 的支持,单本文只讨论 JA心 相关。JAX 用户可以用 Python 装饰器,显式包装标准的Python代码,来说明需要被编译为(可能是 SPMD)XLA 计算的片段。这些 XLA 计算通常是知道输入、输出类型和形状的,bounded loops, 有少数条件语句(看附录B),因此很容易提前估算计算的资源。 我们把这类已知资源需求的计算叫做“编译后/好的函数”。Pathway 程序中,每个这类函数映射到单个(sharded)计算节点上。JAX无法扩展到多个 TPU pod 上,因为 JAX 程序运行在多controller 配置上,会使用 XLA 集合通信来传递数据(这是 JAX 的限制?),这些目前在 TPU 上只能通过 ICI 实现(意思是无法跨Pod 使用?)。PW 可以作为 JAX 后端的替代品,让 JAX 代码可以无需修改,就能运行,只不过 SPMD 计算现在不仅能访问 locally connected TPU cores,也能访问系统里更多的核。因为PW 能通过 ICI 和 DCN 来通信,他让 JAX 程序第一次可以扩展到多个 POD 上(DCN 不是用来做集合通信的吧?),包含上千个 TPU 核

def get_devices(n):
    device_set = pw.make_virtual_device_set()
    return device_set.add_slice(tpu_devices=n).tpus
    
a = jax.pmap(lambda x: x*2., devices=get_devices(2))
b = jax.pmap(lambda x: x + 1., devices=get_devices(2))
c = jax.pmap(lambda x: x / 2., devices=get_devices(2))

@pw.program # Program tracing (optional)
def f(v):
    x = a(v)
    y = b(v)
    z = a(c(x))

    return (y, z)
    
print(f(numpy.array([1., 2.])))

Figure 2. Python user code example for PW running sharded computations(如何体现?是说2个设备,算一个算子?) across multiple islands(多个TPU 之间有数据流关联) of TPU

能够不用修改就运行 JAX 很方便,但无法解锁PW 的完整性能。PW 用户可以请求“虚拟设备”,可以限制设备类型,位置 或者互联拓扑,可以把特定的编译后函数放到这些设备上。系统可以自动处理数据移动和关联计算图的 resharding(啥意思?因为是运行时动态申请的虚拟设备,所以会涉及动态划分图?)

默认情况下,我会把每个编译好的函数转换为一个单独的 PW 程序,包含一个(sharded)计算,意味着一个用户想背靠背运行多个函数,那么需要 client 给每个函数单独的调用和RPC给协调者。我们因此也实现了一个 program tracer,上文中 Figure 2, 用户可以把一段调用了多个编译后函数的 Python 代码包裹起来。tracer 会产生一个单独的 PW 程序,每个单独的编译后的函数代表数据流图里的一个计算节点。

JAX的支持 traced 代码的转换哲学是一个我们想探索的研究方向的较佳匹配。例如,JAX 有个伙伴库叫 FLAX,用来表达分层的 DNN 模型,我们写了一个库来自动转换一个 FLAX 模型到流水化的 PW 程序上。另外 JAX 支持转化为向量化的"per-example" Python代码,产出高效的批量代码,这种转换对探索新的 数据独立的向量控制流比较友好,后面会简单讲到。

4. 架构

如何用 shared dataflow 和 asynchronous gang-scheduling 来克服 ML 系统的限制。

PW 基于已有的系统构建,包括 XLA 来表达和执行 TPU 计算,TF 图和执行器来代表和执行分布式的 CPU 计算,Python ,JAX 和 TensorFlow APIs。通过使用这些技术,我们能关注在创新、协调的部分,同时可以让已有的代码以最小的改动跑起来。

4.1 Resource Manager

上面的 Gang scheduling 可以看出,每个 island 上的 scheduler,负责做岛上的 gang scheduling

PW 后端包含一堆加速器集合,他们分为多个紧密关联的岛屿(上述图里每个岛屿有一个调度器),顺序地通过 DCN 连接起来。PW 有一个"资源管理"者(全局唯一),负责集中地管理跨岛屿的硬件。客户端可以申请特定2D或3D网格的虚拟分片,满足自己的通信模式。每个虚拟分片包含虚拟设备来让客户端表达计算如何布局到网格里。resource manager 为虚拟设备动态分配满足互联拓扑,内存大小要求的 物理资源设备

第一版 RM 实现只用了简单的启发式方法,尝试通过把计算铺平到所有可用设备上,记录一对一的虚拟和物理设备关系来静态平衡负载。未来可以使用更复杂的分配算法,例如考虑所有client计算的资源需求,同时兼顾当前系统的状态,来达到更好的分配策略。

PW 容许后端计算资源动态添加和删除,通过RM记录可用设备。通过单控制器设计容许的虚拟和物理设备之间的间接层,可以未来支持透明地挂起/恢复和迁移的特性,client的虚拟设备可以临时被回收或者重新分配而不需要用户程序的配合。只能说还没执行到,就可以替换为其他设备把?并不能适用于当前正在跑的计算。

4.2 客户端

用户想执行一个 traced program,会调用PW客户端库,它会首先把虚拟设备分配给还没有被执行过的计算(为啥不是RM),注册计算到 resource manager,触发服务器后台编译计算。客户端会给程序构建一个设备位置不感知的 PW 中间表达(IR),是一种定制的 MLIR 方言。这种 IR 会逐渐通过一系列标准的编译器pass 来lowered,最终产出为包含了物理设备位置信息的低级别表示。种低级别的程序考虑了物理设备之间的网络连接,包含了传输从一个计算源头分片到目标位置分片的数据传输操作,包含 scatter 和 gather 操作来交换数据。重复执行低级别的虚拟设备位置不会改变的程序是高效的,如果 RM 修改了虚拟和物理设备之间的映射关系,那么程序可以被重新 relowered。

client在老的单 controller 系统里很快会成为瓶颈,当他需要协调上千个单独的计算,而每个计算之间的数据传输分布在上千个加速器上。PW 使用分片的buffer 抽象(shared buffer) 来代表逻辑buffer,可以分不到多个设备上。这种抽象帮助 client 可扩展,可以把 bookkeeping 任务(包含引用计数)的开销摊薄到逻辑buffer粒度而非单独的分片。

4.3 协调的实现

这个组件是图里的 Scheduler ?是的,后面的 dataflow program 会和scheduler 通信,来决定岛上的一致性的函数执行顺序

这个主要是为了实现成组调度?

PW 依赖 PLAQUE 来做DCN之上的跨主机协调。它是现成的闭源的生产级别分片数据流(sharded dataflow)系统,用在很多面向用户的服务上,需要高扇出/入通信的场景,扩展性和延迟非常关键。low-level PW IR 直接转换为一个 PLAQUE 程序,代表一个数据流的图。PW 对协调有严格的要求,而 PL 都满足。

首先,用做 PW IR 的表示,必须包含一个单节点,用来给每个分片的计算,这是为了保证垮了多个分片的计算的表示是紧凑的(为啥能做到?),比如一个串联了2个计算A和B,有N个分片,每个包含4个节点:Arg -> Compute(A) -> Compute(B) -> Result ,而不需要关注 N 有几个。在 PLAQUE 的运行时实现里,每个节点产出输出数据的tuples,用目标shard做标签,这样在执行数据并行时,N个数据tuple可以流动,IR节点里相邻的节点间各一个。

其次,协调机制的运行时,必须保证沿着分片的边,稀疏数据的传输,在动态选择的分片子集间,消息可以传输。使用标准的进度追踪机制,来检测一个分片的所有的消息都到齐了。高效的稀疏通信是为了避免 DCN 成为数据依赖的控制流的瓶颈的一个要求。这是我们希望 PW 能做到的核心能力。

协调基座用来发送 DCN 消息,是传输调度消息和数据处理时的关键路径,因此它必须以低延迟发送关键消息,当希望在相同主机间达到高吞吐时,批量发送消息。

如果使用可扩展的,通用的数据流引擎来处理 DCN 通信,也是很方便的,因为这样 PW 能用它来做后台的打扫任务,比如配置信息分发,监控程序,清理他们,失败后传递错误等。

我们认为用其他分布式框架比如 Ray 而非 PLAQUE 来实现完整的 PW 实现是可以的。这种实现下,PW 执行器和调度可以被长期运行的 Ray actors 代替,actors 会在 Ray 集群调度之上实现 PW 的调度,executors 可以用 PyTorch 来做 GPU 计算和集合通信。额外需要做一些事来获得类似的高性能,因为 Ray 缺乏这个,比如 HBM 对象存储 (object store), 或者在GPU网络上高效传输远程对象的原语。

4.4 成组动态分发

第二章讨论过,在共享的加速器集合上运行 SPMD 程序的一个要求是支持高效的成组调度。PW 运行时包含集中式的,每个岛屿一个的调度器,负责协调岛上所有计算的一致性顺序。当PW 让一个程序入队列,PL 数据流程序负责:

  1. 把本地编译好的函数在每个加速器上入队列,把buffer future 当作输入。
  2. 把函数执行的输出作为 buffer future 放到网络发送队列里,给到远程的加速器。
  3. 和调度器通信来决定岛上所有程序的一个一致性的函数执行顺序

调度器必须实现在ms级别时间单位分配加速器的策略。我们的当前策略是简单的用 FIFO 顺序入队列,也可以基于预估的执行时间,来做更复杂的重新安排顺序。

4.5 并发异步分发

Figure 4: 对于一个三个节点的程序,线性和并行分发机制的对比。当一个计算,它在设备上执行的时间比花在调度,资源分配,协调上的时间要短,那么异步的流水线就阻塞了,因为 host 侧工作是顺序分发的。并行的异步分发机制可以通过执行并行的host侧分发机制来解决,利用了常规编译好的函数的资源是使用是静态的机制。上图中为了简单,没有画调度器。 在加速器上运行计算时,系统需要利用异步 API 的优势来做把计算和协调重叠。考虑图4a 里三个节点的图,正方形代表三个节点 A、B 和 C,他们运行在主机 A、B、C 上的加速器里。所有节点计算都是常规的编译好的函数。Host A 把节点 A 入队列,接收到一个 A 的输出的 future,然后把这个 future 传输给主机 B。B分配输入,把输入buffer 的地址给到主机 A,执行给启动B函数做的大部分的准备 性工作。当 A 完成,输出通过加速器间直接互联来发送给 B 的输入 buffer(类似 RDMA),然后 B 开始了。一个节点完成到下一个节点开始之间的延迟,可以做到比数据传输时间更短。

上述设计在一个在一个加速器节点的计算比调度、资源分配、主机间协调的时间长时也能work。但如果计算时间特别短,异步的pipeline就会阻塞,host侧工作会成为整个计算顺序执行的关键瓶颈。给的编译好的函数都是常规的,一个节点的前驱输入形状可以在实际中,计算还没入队列之前,就提前计算出来。如果不是常规的呢?

因此提出了创新的并行异步分发的设计,如上图 4b 展示,利用了常规的编译好的函数,它的资源使用是静态的,提前可知道的,来并行执行计算节点在 host 侧的工作(分配资源、调度),而不是串行执行前驱节点已经入队列之后节点需要做的工作。由于只要有函数是常规的情况下,才能并行调度,所以 PW 会把并行调度当作一种优化方案,当出现只有前驱计算完成后(比如数据依赖的控制流),才能知道资源消耗的 情况,就退化到传统的同步/串行模式。当计算的子图可以静态调度,程序会发送一个单个消息(描述了整个子图)给 scheduler,scheduler 可以把子图里的所有分片的执行背靠背顺序化。单个消息是为了减小网络带宽占用,但不需要scheduler实际把所有子图分片对当作一批来全部入栈:计算依然和其他并行提交的程序一起交替执行。第5张评估了不同分发方式的开销。

4.6 数据管理

每个主机管理了一个共享对象存储(shared object store),类似于 Ray 的 object stores,但是可以扩展到也追踪每个分片上加速器 HBM 里的缓存。Client 程序可以拿着在远程主机上的对象引用或者加速器内存,client和server可以用handle 来引用他们,让系统可以在必要时迁移他们。程序中间结果也保存在对象存储里,例如当整个系统在等待加速器之间的数据传输完成,或者把他们传递给计算子序列时。(类似临时内存) 对象是通过拥有者来标签的,这样可以在程序或者 client 失败后,进行垃圾回收。当因为其他计算的buffer在临时占用 HBM 而导致当前计算无法被分配内存,可以使用简单的背压机制来阻塞。

5 实验

6 讨论

6.1 PW design vs implementation

PW 是给大量 TPU 上运行设计的。TPU而非GPU的使用影响了很多底层设计的决策。最大的差别是运行时长(far longer-running) 和更多复杂的计算可以融合到一个单独的 TPU kernel里,因为TPU支持丰富的控制逻辑和通信原语,而 GPU 系统里,这些东西需要被驱动程序执行。GPU相比而言,是更紧密地和宿主机内存系统和 DCN 集成到一起的。TPU 和 PW 很好的匹配,因为 XLA 可以编译出包含融合后集合通信的高性能的函数, 大量高性能 TPU 之间互联的岛屿,可以允许做大量不同尺寸的灵活的计算调度。尽管这样,我们认为高层架构选择是依然对大规模 GPU 系统有效的。

6.2 Resource management

PW 是设计用于种类多样的细粒度动态资源管理策略。初始研究关注在TPU计算的高效动态时间上多路复用。对于更复杂的未来多租户场景,PW 需要处理更多样的资源类型,比如host memory,ICI(是啥?),DCN 和 PCIe 带宽。 PW 的单个控制器模式保证系统有能力跟踪所有可用资源,大规模分配资源。我们计划探索通用的多租户需求如优先级, 性能隔离,访问控制,资源审计等。但时间粒度会更小,资源池维度更大(例如上千核,TB的加速器内存)

6.3 数据依赖的向量化控制流

目前绝大多数 ML 模型会在每一步里每个训练样本跑完之后更新每个模型权重。我们希望让研究员可以做更细粒度控制,便于每个样本后,可以更新不同的模型权重,甚至是子样本(图片的一部分,句子的一部分单词)。MoE 和 路由的胶囊网络通过路由不同(子)样本到有不同模型子权重的加速器上来利用好计算稀疏。路由策略也是训练过程中 学习到的。这种路由需要节点间支持更细粒度数据依赖的数据交换。我们的 ML 研究员希望更高效利用稀疏,有更多的任务,更大的模型。但当前啊的框架限制了他们。这个用一个干净的编程模型和高性能的支持数据依赖的向量化控制流,是未来研究方向。

一些目前 AI 的短板

  1. 目前的模型只能训练来做一件事。而 Pathways 容许训练一个模型来做成千上万种事。我们希望一个模型有不同的能力,按需调用,组合到一起来执行新的,更复杂的任务 -- 更接近大脑在多任务之间泛化的过程。

  2. 目前的模型只关注于一个场景。而 Pathways 容许多场景。人们依赖多场景来感知世界。可以输入图片、文字或者语音。Pathways 可以容许多模态的模型,把上述三者结合起来。

  3. 模型是稠密和不高效的。Pathways 可以让他们稀疏并高效起来。稠密意味着为了完成任务,整个神经网络都被激活,而不管任务是简单或复杂的。这个也不像人处理问题的方法。人脑里有不同的部分,都是给特定任务的,在特定事件下,我们只需要调用特定部分。即大脑里有上千亿神经元,但只需要少数一部分来干活。这样就可以高效和大容量了。比如 GShard 和 Switch Transformer

A 加速器设计考量

下面的方法是加速器设计中,为了达到好的性能而做的常用的方法

A.1 Batching

需要解决带宽瓶颈,需要 high-bandwidth memory (HBM)

A.2 异步编程

同步的会由于 PCIe的延迟,kernel 调度的开销、中断的延迟而浪费加速器的很多计算资源。计算被放入到 streams 来在未来某个时间点执行。这种方法掩盖了小操作的的分发开销

A.3 高性能互联

现在的神经网络维度比加速器内存还大,所以放到多个加速器上进行分片、并行,因此加速器之间的高速互联非常关键。GPU之间,在岛屿之内是通过 NVLink 做几个节点间的高速互联,而在岛屿之间,通过以太网 RDMA 和 Infiniband NIC(GPUDirect) 技术来快速通信。

TPUs 上已经有了定制的网格网络,内建在芯片上,芯片之间可以直接互联,无需宿主机或者 DCN 的介入。GPU 和 TPU 之间的互联通常是通过有30年历史的 MPI 原语(例如 AllReduce)来做的,他是成组的调度,每个程序都必须同时进入相同的原语。当更大的计算在运行时(例如训练更大的神经网络, 训练固定大小的网络,在很多加速器上,通过数据并行的方法),更快的集合操作和网络带宽是必须的,这样才能维持集群资源的高效利用。这激发了很多 chip 网络拓扑的实验,包括 hyper cubes,2-D 和 3-D mesh tori

A.4 单租户

和计算机里的大多数资源不同,加速器经常是独占使用的。因为DL 模型通过增加参数或者 batch size,可以很容易地扩展到使用更多的内存,所以实践中会使用大部分的 HBM。而 PCIe 带宽比 HBM 和加速器之间互联的带宽要低很多。这样细粒度的上下文切换下(HBM 里大部分数据会通过 PCIe 被换出到宿主机 DRAM),导致出现浪费加速器周期的情况。 因此一个宿主机程序没有完全利用加速器,计算资源就会被搁置,无法被有效利用。而且,实践中加速器资源抢占是尽量避免的,因为会导致大规模、共享的集群里异构负载没有被最优使用。很难分配大量物理上相近的设备来有效利用网络局部性。这是什么原因?

A.5 对比 GPU 和 TPU

有很多相似性,但也有很多重要的不同之处。GPU系统倾向于使用小的岛屿规模:使用 NVLink 连接的设备(比如一机8卡),通过 IB 或者数据中心网络技术来互联、聚合到更大规模。GPU 通常是把很多小的预编译好的函数分发到加速器上的编程模型,因为他们是预编译好的,kernel 必须支持动态形状(dynamic shapes)。 GPU间的任何互联通信,无论通过 NVLink 还是 DCN,都是通过 NCCL 库,由宿主机来发起。

TPU 系统有上千块设备,每个岛屿有上百个宿主机,互相之间是全联接的。TPU 包含一个“标量核”,可以协调 TPU 里向量计算但愿,让 TPU 可以执行 XLA 写出的长期运行的函数,而不需要host介入,这些函数可能包含了通过特定 ICI 网络进行的集合通信。因此,TPU上,ML 框架通常构建一个大的 XLA 程序,它会被 JIT 即席编译并分发到加速器上。 单个 XLA 程序可能比一个 GPU kernel 要执行时间上是维度级别更长的,让编译器可以增加很多优化,比如:静态buffer分配和动态中间结果的重计算。由于这类静态buffer分配,TPU 对动态形状只有有限的支持。让她成为 PW 里常规编译好的函数概念是相匹配的。

B 典型的 ML 程序结构

ML 负载在加速器上执行的计算是由我们叫“编译好的计算”的。这些是有如下特点的子计算:

  • 输入和输出的类型和tensor形状,都是在输入数据还没被计算时就知道的
  • 循环的次数是预先知道的,而且会提前推出
  • 条件语句是“函数式”,每个分之的资源是提前分配好的

上述“编译好的计算” 的限制,大部分是因为 ML 模型和硬件的相互演进带来的。上述 A 里详细讨论了。下面主要说计算资源的需求是提前知道的部分的影响。

大部分如今的 高性能 ML 计算都是以一个长序列的编译好的函数来表示,只会偶尔根据数据有分支。由于系统可以提前给函数执行资源分配,所以 ML框架利用这个特性,可以异步地在函数的前驱还没执行之前,就把函数放入到队列里,让 host 侧工作可以和加速器计算并行进行。

当框架提交编译函数的图给JIT编译器,可以利用优化来做比如布局分配和融合来提高加速器代码效率 因为需要优化编译好的函数的图(所以本文里编译好的函数,只是一个IR)来达到峰值性能,意味着框架通常要跟踪(trace)高级别(Python)代码的执行片段,来lowered 到编译好的函数。因此client代码可能是高级别语言,用复杂的状态来写的,高性能的结算通常是通过降低到一个内部表示(IR)上来序列化然后更容易地发送到远程主机上执行。

输入数据处理

继续用 tensorflow/datasets

TODO

  1. 原来 NVLink 也有论文可看

参考资料

  1. Jeff Dean 在 Google Blog 上的:Introducing Pathways: A next-generation AI architecture