Pathways – Asynchronous Distributed Dataflow for ML

Posted on by remaper in LLM 训练
内容 隐藏
1. 论文背景、动机、贡献
2. 设计动机
3. Pathways 的编程模型
4. Pathways 的系统设计
4.1. Resource Manager
4.2. Client
4.3. Coordination implementation
4.4. Gang-scheduled dynamic dispatch
4.5. Parallel asynchronous dispatch
4.6. Data management
5. 性能评估
其他参考：
xx

1. 论文背景、动机、贡献

在探讨现有 AI 模型的局限时，Jeff Dean 曾经说过，今天的人工智能系统总是从头开始学习新问题，我们为数千个单独的任务开发了数千个单独的模型，这种方式需要大量的数据和算力，是非常低效的
理想的发展模式应该是训练1个模型来干多个事情
0
而导致这个现象的根本原因，是因为过去10年的机器学习，总是 算法&硬件&系统 一起演进的，这种 co-evolution 的方式有一定的隐患：系统过于针对当前的任务，而不能很好的适应未来的需求
未来 ML 的需求是什么？
  1. 模型会越来越大：
    1. 多模态：复杂，参数多。纯粹的数据并行已经不行了，开始引入流水线并行、模型并行等新的并行训练方式
    2. 计算稀疏：大模型特有的一个问题
  2. 硬件异构：SPMD 不适用了，因为我们需要在不同的“硬件”上运行不同的代码，需要转向 MPMD
  3. 训推一体：提升硬件利用率
为了实现这个愿景，论文提出了一种叫 [Pathways] 的通用AI架构。Pathways 旨在用一个架构同时处理多项任务，并且拥有快速学习新任务、更好地理解世界的能力

2. 设计动机

为什么说当今的ML系统都不适合用来处理大、稀疏、不规则的模型？
我们来看下现有ML系统的一些优缺点，从下面2个架构思路来进行分类：
  1. SPMD 还是 DPMD
    1. SPMD：所有节点都运行相同的代码，它的问题在于不支持流水线负载
    2. DPMD：不同节点运行不同的代码，它的问题在于 gang-scheduler（群调度）容易死锁
  2. 单控制器（single controller）还是多控制器（multi controller）
    1. 多控制器：最常见，每个节点有自己的控制器
      1. 优点：低延迟，因为所有交互都是通过PCIe（同机通信）或者NVLink（跨机通信）
      2. 缺点：
        1. Any communication beyond standard collectives in multi-controller systemsrequires users to implement their own coordination primitives. (备注：这句话其实我不理解）
        2. 独占硬件，资源利用率不高
    2. 单控制器：控制器是中心化的，节点只有runtime
      1. 优点：灵活性很强
      1. 缺点：
        1. 通信慢（关键）：几乎所有的通信都是 farther away
        2. 调试困难：由于图可能被优化，所以 debug 起来很困难
        3. 调度延迟，导致资源利用率不高
如下：
0
我们对系统进行分类
  1. SPMD
    1. multi controller：MPI，Pytorch ，Tensorflow，JAX
0
    1. single controller：Tensorflow v1
0
  1. non-SPMD
    1. single controller：Tensorflow v1
0
    1. multi controller：由于 non-SPMD 本身就依赖一个中心式的控制器，所以就没有所谓的 multi controller 了
未来的趋势一定是 non-SPMD 的，因为模型会越来越复杂，具备流水线、计算稀疏等特点的ML负载，很难用 SPMD 架构来实现了
然后虽然单控制器有debug困难、调度延迟等一些问题，但这都不是关键的（都是局部可优化问题，而不是结构性问题），论文认为 single controller 这种 dataflow 的设计仍然是最理想的设计
言外之意：tensorflow 还是最棒的！
剩下的就是怎么解决 non-SPMD + single controller 的缺点问题：
  1. gang-scheduler 死锁问题
  2. 性能问题：异步数据流 + 资源集中调度

3. Pathways 的编程模型

Pathways 目前可以无缝对接 tensorflow 以及 JAX（可以理解为是 JAX 或者 tf 的一个后端），用户代码不需要任何修改
支持2种方式：
  1. jax.pmap(fn, devices=get_devices())：把一个函数映射到具体的TPU上执行，每个fn会编译成一个computation
  2. @pw.program：把整个函数打包编译成一个大的computation（如果函数内包括pmap，也会一起编译进来，pmap不再单独编译了）
0
除此之外，pathways 支持 DCN 通讯，能够在多 TPU Pod 内运行

4. Pathways 的系统设计

为了实现对业务代码的最小侵入，pathways大量复用了现有的系统能力：
  1. 支持 JAX 和 tensorflow API
  2. 使用 tensorflow graphs 和 executors 来执行 CPU 计算
  3. 使用 XLA 来执行 TPU 计算
再看一下系统架构
0
从这个架构图里面，我们大概知道：
  1. Resource Manager 是全局唯一的
  2. 每个 island 有独立的 scheduler
  3. 每个 island 有多个 host，每个 host 对应1~N个device（TPU）
  4. 每个 device 有独立的 executor

4.1. Resource Manager

xx

4.2. Client

xx

4.3. Coordination implementation

xx

4.4. Gang-scheduled dynamic dispatch

xx

4.5. Parallel asynchronous dispatch

xx

4.6. Data management

xx

5. 性能评估

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注