AI系统-35大结局之AI框架与编程范式

39 阅读54分钟

image.png

AI系统系列,本篇要迎来大结局。从最开始想写点AI芯片的东西,到看到ZOMI酱的《AI系统》学习材料,就想着给大家分享出来:

image.png

本系列主要就是参考ZOMI酱的文章,摘取一些自己觉得概括性的入门内容,加上自己的一些理解。有想了解细节技术的可以去看看ZOMI酱的资料(比较多技术细节)。

AI系统系列目录: mp.weixin.qq.com/mp/appmsgal…

  • 首先介绍了AI的历史及算法: XXX
  • 然后重点介绍了AI芯片: XXX
  • 最后介绍了AI编译器和推理引擎: XXX

下面通过本篇内容进行一下回顾。

1. AI框架介绍

1.1 深度学习

image.png 深度学习的概念源于人工神经网络的研究,但是并不完全等于传统神经网络。在叫法上,很多深度学习算法中都会包含”神经网络”这个词,比如:卷积神经网络 CNN、循环神经网络 LSTM。所以,深度学习可以说是在传统神经网络基础上的升级,约等于神经网络。

虽然深度学习理论最初创立于上世纪八十年代,但有两个主要原因导致其直到近年来才得以发挥巨大作用:

  • 深度学习需要大量的标签化数据。例如,无人驾驶汽车模型训练需要数万亿张图片和数千万小时的视频进行学习;
  • 深度学习需要巨大的计算能力。例如,需要局别并行架构和集群组网能力的高性能 GPU/NPU 对于深度学习计算进行加速。

一个神经网络由多个神经元结构组成,每一层的神经元都拥有多个输入和输出,一层可以由多个神经元组成。例如,第 2 层神经网络的神经元输出是第 3 层神经元的输入,输入的数据通过神经元上的激活函数(非线性函数如 tanh、sigmod 等),来控制输出的数值。

数学上简单地理解,单个神经元其实就是一个x*w的矩阵乘,然后加一个激活函数 ,通过复合函数组合神经元,就变成一个神经网络的层。这种模拟生物神经元的数学计算,能够很好地对大规模独立同分布的数据进行非线性映射和处理,使得其能够应对到 AI 的不同任务。

神经网络与神经元关系

深度学习一般流程是:

  • 1)构建神经网络模型,
  • 2)定义损失函数和优化器(优化目标),
  • 3)开始训练神经网络模型(计算梯度并更新网络模型中的权重参数),
  • 4)最后验证精度,其流程如下图所示,前三步最为重要。

深度学习构建流程

因为 AI 框架已经帮我们封装好了许多功能,所以遇到神经网络模型的精度不达标,算法工程师可以调整网络模型结构、调节损失函数、优化器等参数重新训练,不断地测试验证精度,因此很多人戏称算法工程师又是“调参工程师”。

但是在这一过程中,这种机械的调参是无法触碰到深度学习的本质的,为了了解实际的工作原理,进行总结:训练的过程本质是进行反向求导(反向传播算法实现)的过程,然后通过迭代计算求得神经网络中的参数,调整参数是控制这一过程的前进速度和方向。

1.2 AI框架的发展及趋势

image.png

AI 框架在时间维度的发展大致可以分为四个阶段,分别为

  • 1)2000 年初期的萌芽阶段
  • 2)2012~2014 年的成长阶段、
  • 3)2015 年~2019 年的爆发阶段,
  • 4)2020 年以后深化阶段。

image.png

下面我们来看下AI框架发展的趋势:

1.2.1 全场景

AI 框架将支持端边云全场景跨平台设备部署

网络模型需要适配部署到端边云全场景设备,对 AI 框架提出了多样化、复杂化、碎片化的挑战。随着云服务器、边缘设备、终端设备等 AI 硬件运算设备的不断涌现,以及各类 AI 运算库、中间表示工具以及编程框架的快速发展, AI 软硬件生态呈现多样化发展趋势。

但目前主流 AI 框架仍然分为训练部分和推理部分,两者在不同厂商间不完全兼容。训练出来的模型也不能通用,学术科研项目间难以合作延伸,造成了 AI 框架的碎片化。目前业界并没有统一的中间表示层标准,导致各硬件厂商解决方案存在一定差异,以致应用模型迁移不畅,增加了应用部署难度。因此,基于 AI 框架训练出来的模型进行标准化互通将是未来的挑战。

1.2.1 易用性

AI 框架将注重前端便捷性与后端高效性的统一

AI 框架需要提供更全面的 API 体系以及前端语言支持转换能力,从而提升前端开发便捷性。AI 框架需要能为开发者提供完备度高、性能优异、易于理解和使用的 API 体系。

AI 框架需要提供更为优质的动静态图转换能力,从而提升后端运行高效性。从开发者使用 AI 框架来实现模型训练和推理部署的角度看,AI 框架需要能够通过动态图的编程范式,来完成在模型训练的开发阶段的灵活易用的开发体验,以提升模型的开发效率;通过静态图的方式来实现模型部署时的高性能运行;同时,通过动态图转静态图的方式,来实现方便的部署和性能优化。目前 PyTorch2.0 的图编译模式走在业界前列,不一定成为最终形态,在性能和易用性方面的兼顾仍然有待进一步探索。

1.2.3 大规模分布式

AI 框架将着力强化对超大规模 AI 的支持

OpenAI 于 2020 年 5 月发布 GPT-3 模型,包含 1750 亿参数,数据集(处理前)达到 45T,在多项 NLP 任务中超越了人类水平。随之谷歌不断跟进分布式技术,超大规模 AI 逐渐成为新的深度学习范式。

超大规模 AI 需要大模型、大数据、大算力的三重支持,对 AI 框架也提出了新的挑战,

  1. 内存:大模型训练过程中需要存储参数、激活、梯度、优化器状态,
  2. 算力:2000 亿参数量的大模型为例,需要 3.6EFLOPS 的算力支持,必要构建 AI 计算集群满足算力需求
  3. 通信:大模型并行切分到集群后,模型切片之间会产生大量通信,从而通信就成了主要的瓶颈
  4. 调优:E 级 AI 算力集群训练千亿参数规模,节点间通信复杂,要保证计算正确性、性能和可用性,手动调试难以全面兼顾,需要更自动化的调试调优手段
  5. 部署:超大规模 AI 面临大模型、小推理部署难题,需要对大模型进行完美压缩以适应推理侧的部署需求

1.2.4 科学计算的交融

AI 框架将进一步与科学计算深度融合交叉

传统科学计算领域亟需 AI 技术加持融合。计算图形可微编程,类似 Taichi 这样的语言和框架,提供可微物理引擎、可微渲染引擎等新功能。因此未来是一个 AI 与科学计算融合的时代,传统的科学计算将会结合 AI 的方法去求解既定的问题。至于 AI 与科学计算结合,看到业界在探索三个方向:

  1. 利用 AI 神经网络进行建模替代传统的计算模型或者数值模型,目前已经有很大的进展了,如拿了戈登贝尔奖的分子动力学模型 DeepMD。
  2. AI 求解,模型还是传统的科学计算模型,但是使用深度学习算法来求解,这个方向已经有一定的探索,目前看到不少基础的科学计算方程已经有对应的 AI 求解方法,比如 PINNs、PINN-Net 等,当然现在挑战还很大,特别是在精度收敛方面,如果要在 AI 框架上使用 AI 求解科学计算模型,最大的挑战主要在前端表达和高性能的高阶微分。
  3. 使用 AI 框架来加速方程的求解,科学计算的模型和方法都不变的前提下,与深度学习使用同一个框架来求解,其实就是把 AI 框架看成面向张量计算的通用分布式计算框架。

1.3 编程范式

image.png

编程范式、编程范型、或程式设计法(Programming paradigm),是指软件工程中的一类典型的编程风格。常见的编程范型有:函数式编程、命令式编程、声明式编程、面向对象编程等等,编程范式提供并决定了开发者对程序执行的看法。在开发者使用 AI 框架进行编程的过程中,

主要使用到的编程范式主要有 2 种:1)声明式编程与 2)命令式编程。

  • 命令式编程(Imperative):详细的命令机器怎么(How)去处理一件事情以达到想要的结果(What);
  • 声明式编程(Declarative):只告诉想要的结果(What),机器自己摸索执行过程(How)。

image.png

主流 AI 框架,无论 PyTorch 还是 TensorFlow 都使用都以 Python 为主的高层次语言为前端,提供脚本式的编程体验,后端用更低层次的编程模型和编程语言开发。后端高性能可复用模块与前端深度绑定,通过前端驱动后端方式执行。AI 框架为前端开发者提供声明式(declarative programming)和命令式(imperative programming)两种编程范式。

在主流的 AI 框架中,TensorFlow 提供了声明式编程体验,PyTroch 提供了命令式的编程体验。但两种编程模型之间并不存在绝对的边界,multi-stage 编程和即时编译(Just-in-time, JIT)技术能够实现两种编程模式的混合。随着 AI 框架引入更多的编程模式和特性,例如 TensorFlow Eager 模式和 PyTorch JIT 的加入,主流 AI 框架都选择了通过支持混合式编程以兼顾两者的优点

1.4 计算图

目前主流的 AI 框架都选择使用计算图来抽象神经网络计算表达,通过通用的数据结构(张量)来理解、表达和执行神经网络模型,通过计算图可以把 AI 系统化的问题形象地表示出来。

在真正的 AI 工程化过程中,我们会遇到诸多问题。而为了高效地训练一个复杂神经网络,AI 框架需要解决许多问题,例如:

  • 如何对复杂的神经网络模型实现自动微分?
  • 如何利用编译期的分析 Pass 对神经网络的具体执行计算进行化简、合并、变换?
  • 如何规划基本计算 Kernel 在计算加速硬件 GPU/TPU/NPU 上高效执行?
  • 如何将基本处理单元派发(Dispatch)到特定的高效后端实现?
  • 如何对通过神经网络的自动微分(反向传播实现)衍生的大量中间变量,进行内存预分配和管理?

为了使用用统一的方式,解决上述提到的挑战,驱使着 AI 框架的开发者和架构师思考如何为各类神经网络模型的计算提供统一的描述,从而使得在运行神经网络计算之前,能够对整个计算过程尽可能进行推断,在编译期间自动为深度学习的应用程序补全反向计算、规划执行、降低运行时开销、复用和节省内存。能够更好地对特定领域语言(DSL),这里特指深度学习和神经网络进行表示,并对使用 Python 编写的神经网络模型进行优化与执行。

因此派生出了目前主流的 AI 框架都选择使用计算图来抽象神经网络计算。

计算图/数据流图

我们会经常遇到有些 AI 框架把统一的图描述称为数据流图,有些称为计算图,这里可以统称为计算图。下面简单介绍为什么可以都统称为计算图的原因。

  • 数据流图(Data Flow Diagram,DFD):从数据传递和加工角度,以图形方式来表达系统的逻辑功能、数据在系统内部的逻辑流向和逻辑变换过程,是结构化系统分析方法的主要表达工具及用于表示软件模型的一种图示方法。在 AI 框架中数据流图表示对数据进行处理的单元,接收一定的数据输入,然后对其进行处理,再进行系统输出。
  • 计算图(Computation Graph):被定义为有向图,其中节点对应于数学运算,计算图是表达和评估数学表达式的一种方式。而在 AI 框架中,计算图就是一个表示运算的有向无环图(Directed Acyclic Graph,DAG)。

其两者都把神经网络模型统一表示为图的形式,而图则是由节点和边组成。其都是在描述数据在图中的节点传播的路径,是由固定的计算节点组合而成,数据在图中的传播过程,就是对数据进行加工计算的过程。下面以公式为例:

对上述公式转换为对应的计算图。

计算图/数据流图

1.5 分布式并行

大模型模型参数量实在太大,需要分布式并行训练能力一起来加速训练过程。分布式并行是在大规模 AI 集群上工作的,想要加速就需要软硬件协同,不仅仅要解决通信拓扑的问题、集群组网的问题,还要了解上层 MOE、Transform 等新兴算法。通过对算法的剖析,提出模型并行、数据并行、优化器并行等新的并行模式和通信同步模式,来加速分布式训练的过程。最小的单机执行单元里面,还要针对大模型进行混合精度、梯度累积等算法,进一步压榨集群的算力!

,包括并行训练,如:数据并行(Data Parallelism, DP)、模型并行(Model Parallelism, MP)、混合并行(Hybrid Parallel),可扩展的分布式训练组件,如:设备网格(Device Mesh)、RPC 分布式训练以及自定义扩展等。每种方法在特定用例中都有独特的优势。

具体来说,这些功能的实现可以分为三个主要组件:

  1. 分布式数据并行训练(DDP)是一种广泛采用的单程序多数据训练范式。在 DDP 中,模型会在每个进程上复制,每个模型副本将接收不同的输入数据样本。DDP 负责梯度通信以保持模型副本同步,并将其与梯度计算重叠以加速训练。
  2. 基于 RPC 的分布式训练(RPC)支持无法适应数据并行训练的通用训练结构,例如分布式流水线并行、参数服务器范式以及 DDP 与其他训练范式的组合。它有助于管理远程对象的生命周期,并将自动微分引擎扩展到单个计算节点之外。
  3. 提供了在组内进程之间发送张量的功能,包括集体通信 API(如 All Reduce 和 All Gather)和点对点通信 API(如 send 和 receive)。尽管 DDP 和 RPC 已经满足了大多数分布式训练需求,PyTorch 的中间表达 C10d 仍然在需要更细粒度通信控制的场景中发挥作用。例如,分布式参数平均,在这种情况下,应用程序希望在反向传播之后计算所有模型参数的平均值,而不是使用 DDP 来通信梯度。这可以将通信与计算解耦,并允许对通信内容进行更细粒度的控制,但同时也放弃了 DDP 提供的性能优化。

通过充分利用这些分布式训练组件,开发人员可以在各种计算要求和硬件配置下高效地训练大模型,实现更快的训练速度和更高的模型准确性。

当前的主流框架(如TensorFlow、Caffe和MXNet)提供了一些基本的并行技术,其中大多数框架提供了算子级别切分、流水线并行或者优化器切分的功能,支持的并行维度和功能的完整度欠佳。

  • 第一,这些框架通过手动的切分来切分神经网络模型来实现模型并行,配置难度非常大,对开发者的要求非常高,需要有丰富经验的专家来操作。实现混合并行(数据并行和模型并行同时进行)又极大增加了开发的复杂度。最近的研究成果提出了简化混合并行的方法,但这些方法在几个方面都存在局限性。
  • 第二,随着目前模型规模拓展到了万亿级别,训练卡数的规模也上升到了千级,以往的算子级别切分的并行不能满足目前大模型的需求。在大集群训练中,由于模型切分导致的通信占比在整个迭代耗时中升高,需要引入流水线并行等计算来降低通信占比,提升模型训练效率。另外,混合专家(MoE)技术能够在提升模型规模的同时,较少的提升计算量,是目前的一种主流技术。

2. 计算图

为了高效地训练一个复杂神经网络,框架需要解决诸多问题,例如:如何实现自动求导,如何利用编译期分析对神经网络计算进行化简、合并、变换,如何规划基本计算单元在加速器上的执行,如何将基本处理单元派发(dispatch)到特定的高效后端实现,如何进行内存预分配和管理等。用统一的方式解决这些问题都驱使着框架设计者思考如何为各类神经网络计算提供统一的描述,从而使得在运行神经网络计算之前,能够对整个计算过程尽可能进行推断,在编译期自动为用户程序补全反向计算,规划执行,最大程度地降低运行时开销。

目前主流的 AI 框架都选择使用计算图来抽象神经网络计算,《计算图》实际上,AI 框架主要的职责是把深度学习的表达转换称为计算机能够识别的计算图,计算图作为 AI 框架中核心的数据结构,贯穿 AI 框架的大部分整个生命周期,于是计算图对于 AI 框架的前端核心技术就显得尤为重要。

image.png 如上图中,计算图是神经网络模型的中间表示,非常的重要。

随着神经网络模型越来越复杂,TensorFlow 和 PyTorch以及包括混合专家模型 MOE、生成对抗网络 GAN、注意力模型 Attention Transformer 等。复杂的模型结构需要 AI 框架能够对模型算子的执行依赖关系、梯度计算以及训练参数进行快速高效的分析,便于优化模型结构、制定调度执行策略以及实现自动化梯度计算,从而提高 AI 框架训练的效率。所以目前主流的 AI 框架都选择使用计算图来抽象神经网络计算表达,通过通用的数据结构(张量)来理解、表达和执行神经网络模型,通过计算图可以把 AI 系统化的问题形象地表示出来。

2.1 计算图原理

image.png

AI 框架需要解决许多问题,例如:

  • 如何对复杂的神经网络模型实现自动微分?
  • 如何利用编译期的分析 Pass 对神经网络的具体执行计算进行化简、合并、变换?
  • 如何规划基本计算 Kernel 在计算加速硬件 GPU/TPU/NPU 上高效执行?
  • 如何将基本处理单元派发(Dispatch)到特定的高效后端实现?
  • 如何对通过神经网络的自动微分(反向传播实现)衍生的大量中间变量,进行内存预分配和管理?

为了使用用统一的方式,解决上述提到的挑战,驱使着 AI 框架的开发者和架构师思考如何为各类神经网络模型的计算提供统一的描述,从而使得在运行神经网络计算之前,能够对整个计算过程尽可能进行推断,在编译期间自动为深度学习的应用程序补全反向计算、规划执行、降低运行时开销、复用和节省内存。能够更好地对特定领域语言(DSL),这里特指深度学习和神经网络进行表示,并对使用 Python 编写的神经网络模型进行优化与执行。

image.png

计算图(Computation Graph):被定义为有向图,其中节点对应于数学运算,计算图是表达和评估数学表达式的一种方式。而在 AI 框架中,计算图就是一个表示运算的有向无环图(Directed Acyclic Graph,DAG)。

数据表达方式:

  • 标量 Scalar -标量(scalar),亦称“无向量”。有些物理量,只具有数值大小,而没有方向,部分有正负之分,物理学上指有大小而没有方向的量(跟「矢量」相区别)。物理学中,标量(或作纯量)指在坐标变换下保持不变的物理量。用通俗的说法,标量是只有大小,没有方向的量,如功、体积、温度等。

在 AI 框架或者计算机中,标量是一个独立存在的数,比如线性代数中的一个实数 488 就可以被看作一个标量,所以标量的运算相对简单,与平常做的算数运算类似。代码 x 则作为一个标量被赋值。

x = 488
  • 向量 Vector

向量(vector),物理、工程等也称作矢量、欧几里得向量(Euclidean vector),是数学、物理学和工程科学等多个自然科學中的基本概念。指一个同时具有大小和方向,且满足平行四边形法则的几何對象。理论数学中向量的定义为任何在稱為向量空间的代數結構中的元素。

在 AI 框架或者计算机中,向量指一列顺序排列的元素,通常习惯用括号将这些元素扩起来,其中每个元素都又一个索引值来唯一的确定其中在向量中的位置。其有大小也有方向,以公式为例,其代码 x_vec 则被作为一个向量被赋值。

x_vec = [1.1, 2.2, 3.3]
  • 矩阵 Matrix

矩阵(Matrix)是一个按照长方阵列排列的复数或实数集合,最早来自于方程组的系数及常数所构成的方阵。这一概念由 19 世纪英国数学家凯利首先提出。矩阵是高等代数学中的常见工具,也常见于统计分析等应用数学学科中。

在机器学习领域经常被使用,比如有 N 个用户,每个用户有 M 个特征,那这个数据集就可以用一个 NxM 的矩阵表示,在卷积神经网络中输入模型的最初的数据是一个图片,读取图片上的像素点(Pixel)作为输入,一张尺寸大小为 256x256 的图片,实质上就可以用 256*256 的矩阵进行表示。

以公式为例,其代码 x_mat 则被表示为一个矩阵被赋值。

x_mat = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

图中对标量、向量、矩阵进行形象化表示:

标量、向量与矩阵示例

  • 张量

张量(tensor)理论是数学的一个分支学科,在力学中有重要应用。张量这一术语起源于力学,它最初是用来表示弹性介质中各点应力状态的,后来张量理论发展成为力学和物理学的一个有力的数学工具。张量之所以重要,在于它可以满足一切物理定律必须与坐标系的选择无关的特性。张量概念是矢量概念的推广,矢量是一阶张量。

在几何代数中,张量是基于向量和矩阵的推广,通俗一点理解的话,可以将标量是为零阶张量,向量视为一阶张量,矩阵视为二阶张量。在 AI 框架中,所有数据将会使用张量进行表示,例如,图像任务通常将一副图片根据组织成一个 3 维张量,张量的三个维度分别对应着图像的长、宽和通道数,一张长和宽分别为 H, W 的彩色的图片可以表示为一个三维张量,形状为 (C, H, W)。自然语言处理任务中,一个句子被组织成一个 2 维张量,张量的两个维度分别对应着词向量和句子的长度。

一组图像或者多个句子只需要为张量再增加一个批量(batch)维度,N 张彩色图片组成的一批数据可以表示为一个四维张量,形状为 (N, C, H, W)。

计算图表示AI框架:

image.png z=x+y为例如上图,计算图是用来描述运算的有向无环图,有两个主要元素:节点 (Node) 和边 (Edge)。节点表示数据,如向量、矩阵、张量;边表示具体执行的运算,如加、减、乘、除和卷积等。

image.png

上面 (a) 中计算图具有两个节点,分别代表卷积 Conv 计算和激活 ReLU 计算,Conv 计算接受三个输入变量 x 和权重 w 以及一个偏置 b,激活接受 Conv 卷积的输出并输出一个变量。(b)为对应(a)的反向计算图,在神经网络模型训练的过程当中,自动微分功能会为开发者自动构建反向图,然后输入输出完整一个完整 step 计算。

总而言之,AI 框架的设计很自然地沿用了张量和张量操作,将其作为构造复杂神经网络的基本描述单元,开发者可以在不感知复杂的框架后端实现细节的情况下,在 Python 脚本语言中复用由后端优化过的张量操作。而计算 Kernel 的开发者,能够隔离神经网络算法的细节,将张量计算作为一个独立的性能域,使用底层的编程模型和编程语言应用硬件相关优化。

2.2 自动微积分

自动求导应用链式法则求某节点对其他节点的雅可比矩阵,它从结果节点开始,沿着计算路径向前追溯,逐节点计算雅可比。将神经网络和损失函数连接成一个计算图,则它的输入、输出和参数都是节点,可利用自动求导求损失值对网络参数的雅可比,从而得到梯度。

image.png

神经网络模型可以用来近似(拟合)一个无法求解的复杂函数 ,其中输入变量 x 和输出变量 y 皆为高维变量。具体训练过程主要分为以下三个部分:

  1. 前向计算:定义这神经网络模型的前向传播过程,即网络训练的 forward 部分,张量数据输入神经网络模型,模型输出具体的预测值,类似 y=fun(x)。这里的前向传播(forward propagation 或 forward pass)指的是按顺序(从输入层到输出层)计算和存储神经网络中每层的结果。
  2. 计算损失:根据损失函数的定义,一般为真实样本的 , 和神经网络模型的预测  的比较函数。在损失函数中分别对每个维度的参数求其偏导数,得到每个参数的偏导数值即 x_i.grad()
  3. 更新权重参数:根据优化器(Optimizer)的学习策略,小幅通过反向计算图更新网络模型中的各个权重参数的梯度,即反向传播的过程(backward propagation 或 backward pass)。先看其梯度的 grad 正负,再根据正负方向对原参数值加减一定比例的梯度值。假设更新公式为 , 如果梯度值为正,网络模型的权重参数就会减小;如果梯度值为负,网络模型的权重参数值就会增大。

在训练神经网络时,前向传播和反向传播相互依赖。对于前向传播,沿着依赖的方向遍历计算图并计算其路径上的所有变量。然后将这些用于反向传播,其中计算顺序与计算图的相反。

以上述简单网络为例:一方面,在前向传播期间计算正则项取决于模型参数和的当前值。它们是由优化算法根据最近迭代的反向传播给出的。另一方面,反向传播期间参数的梯度计算,取决于由前向传播给出的隐藏变量的当前值。

在计算图的概念中会经常提到自动微分功能,通过自动微分来构建反向的计算图。而在神经网络模型的训练流程和训练原理中,主要会提及反向传播算法,优化器对损失函数进行求导后的值,通过反向传播算法把传递给神经网络的每一层参数进行更新。那么在 AI 框架中,自动微分和反向传播之间的关系是什么呢?

首先,自动微分是将复合函数分解为输出变量(根节点)和一系列的输入变量(叶子节点)及基本函数(中间节点),构成一个计算图,并以此计算任意两个节点间的梯度:

  • 加法法则:任意两个节点间的梯度为它们两节点之间所有路径的偏微分之和;
  • 链式法则:一条路径的偏微分为路径上各相邻节点间偏微分的连乘。

而在神经网络中,只要各个组件以及损失函数都是可微的,那么损失函数就是关于各个输入变量的可微的复合函数。这个时候就可以使用自动微分的方式去计算神经网络模型里面的输入变量梯度,从而使用梯度下降算法减小误差。

因此,反向传播算法实际上就是自动微分,只不过在 AI 框架中,实际上计算图中的根节点为度量误差的损失函数,因而把节点间的偏导称为误差项。

对应计算图 AI 框架基于反向模式(Reverse Mode)的自动求导机制中,根据反向计算图的构建时机,又分为基于对偶图的自动求导、基于表达式追踪和图层 IR 的自动求导,基于对偶图的求导方式通常与图层 IR 相结合,而使用动态图的 AI 框架多采用基于表达式追踪的自动求导方式。

2.3 调度与执行

基于计算图的 AI 框架中,训练的过程阶段中,会统一表示为由基础算子构成的计算图,算子属于计算图中的一个节点,由具体的后端硬件进行高效执行。

目前 AI 框架的前端负责给开发者提供对应的 API,通过统一表示把开发者编写的 Python 代码表示为前向计算图,AI 框架会根据前向计算图图,自动补全反向计算图,生成出完整的计算图。神经网络模型的整体训练流程,则对应了计算图的数据流动的执行过程。算子的调度根据计算图描述的数据依赖关系,确定算子的执行顺序,由运行时系统调度计算图中的节点到设备上执行。

实际上,计算图的执行方式,可以分为两种模式:1)逐算子下发执行的交互式方式,如 PyTroch 框架;2)以及整个计算图或者部分子图一次性下发到硬件进行执行,如 TensorFlow 和 MindSpore。无论采用哪种模式,其大致架构如下所示。

计算图执行通用架构

计算图的调度主要是指静态图。在静态图中,需要先定义好整个计算流,再次运行的时就不需要重新构建计算图,因此其性能更加高效。之所以性能会更高效,是因为会对计算图中的算子的执行序列进行调度优化。

什么是算子?

AI 框架中对张量计算的种类有很多,比如加法、乘法、矩阵相乘、矩阵转置等,这些计算被称为算子(Operator),它们是 AI 框架的核心组件。为了更加方便的描述计算图中的算子,现在来对算子这一概念进行定义:

  • 狭义的算子(Kernel) :对张量 Tensor 执行的基本操作集合,包括四则运算,数学函数,甚至是对张量元数据的修改,如维度压缩(Squeeze),维度修改(reshape)等。
  • 广义的算子(Function) :AI 框架中对算子模块的具体实现,涉及到调度模块,Kernel 模块,求导模块以及代码自动生成模块。

无论是大模型还是传统的神经网络模型,实际上最后执行都会落在单台设备环境上执行对应的算子。对单设备执行环境,制约计算图中节点调度执行的关键因素是节点之间的数据流依赖和具体的算子。

假设继续以简单的复合函数为例子:

image.png

AI 框架根据上述计算图的数据流的依赖关系,在单设备环境下,依次调用具体的算子可以如下所示:

# 正向执行算子
Log(v_(-1), 2) -> v1
Mul(v_(-1), v0) -> v2
Sin(v0) -> v3
Add(v1, v2) -> v4
Sub(v4, v3) -> v5

# 反向执行算子
Sub_grad(v5, v5_delta) -> v4_delta
...

由于计算图准确的描述了算子之间的依赖关系,运行时的调度策略可以变得十分直接。根据计算图中的数据流依赖关系和计算节点函数,通过先进先出队列来执行具体的计算逻辑:

  1. 初始状态下,AI 框架会在运行时将计算图中入度为 0 的节点加入到 FIFO(First-In-First-Out)队列中
  2. 从 FIFO 队列中选择下一个节点,分配给线程池中的一个线程执行计算;
  3. 当前节点执行结束后,会将其后继节点加入就绪队列,当前节点出队;
  4. AI 框架在运行时继续处理 FIFO 队列中的剩余节点,直到遍历完所有的节点,队列为空。

图中按照数据流约束执行对应的计算图的一个可能调度序列。其中蓝色为正向计算时候用到的算子,红色为反向计算时候用到的算子。这种调度方式主要以 PyTorch 的默认执行方式,TensorFlow 的 eager 模式,以及 MindSpore 的 PyNative 模式为主。

执行队列时间轴

目前大模型非常的火,典型代表如下图 Transformer Decoder 堆叠的模型结构,这时如果后端有多个执行算子 Kernel 的硬件加速设备,因为模型结构太大,参数量太多,没有办法在一张 AI 加速卡上放下整个计算图,因此在 AI 框架的运行时在调度执行计算图前,可以对网络模型进行切分,按照模型结构层数进行切分,把 2/3 层 Transformer 结构模块放在同一设备上。

下面以简单的模型并行对神经网络模型的计算图进行切分,对模型按层数来切分,也可以按照模型单一层横向来切分出不同的子图。

针对 Transformer 模型并行

多计算设备环境下执行计算图,AI 框架的运行时需要解决,如何将计算图中的具体计算,放置到不同设备上以及如何管理跨设备数据传输两个问题:

  1. 计算图切分:给定一个计算图,并将计算图切分为不同的子图或者单算子后,放置到多个计算设备上,每个设备拥有计算图的一部分。
  2. 跨设备通信:子图被放置不同设备上,此时 AI 框架会为计算图新增一些跨设备的链接和通信节点(All Reduce 或 All Gather 等集合通信),实现跨设备数据传输。

基本的 Inception 模块执行策略

实际上做好计算图切分,并把计算图映射到多设备是一个复杂的组合优化问题,目前针对大模型在千卡集群规模下同时进行训练的最优并行策略寻优,称为自动并行。

PyTorch 的函数是一个非常复杂核心的模块,其大部分代码都是由 PyTorch tool 根据模板文件自动生成。如果想要查看其源代码,无法直接在 PyTorch 的 GitHub 代码库中搜索到,必须要将代码下载到本地并进行编译。当调用函数时,就会接触到 PyTorch 的调度模块。

以 PyTorch 的加法为例,假设调用 torch.add 函数 API 时,AI 框架总共会经历两次调度:

PyTorch 算子执行

第一次调度会根据执行张量的设备(device)和布局(layout)动态选择对应的实现函数,比如 <CPU, Strided> Tensor<CPU, Sparse> Tensor或者<GPU, , Strided> Tensor。不同设备布局的实现,可能会编译在不同的动态链接库里。

第二次调度则会根据张量元素的数据类型,通过 switch 分支的方式进行一次轻量级的静态选择,最终选出合适的 Kernel 来执行对张量的操作。

Kernel 主要是算子的计算模块,但是别忘记了在深度学习中,算子还包含求导模块。计算模块主要定义了 Kernel 的计算步骤,需要先在 aten/src/ATen/native/native_functions.yaml 中声明 Kernel 计算模块的函数签名,然后在 native/ 目录下实现该函数。

在前面的函数调用中,主要就通过 Kernel 对张量进行操作。求导模块主要是对计算模块的一个反向求导,需要直接在 tools/autograd/derivatives.yaml 中声明定义求导的过程,剩下就可以交给 Autograd 代码生成模块自动生成对应的代码。

2.4 控制流实现

在计算机科学中,控制流(Control Flow)定义了独立语句,指令,函数调用等执行或者求值的顺序。例如,根据函数 A 的输出值选择运行函数 B 或者 C 中的一个。

image.png

随着神经网络算法的快速发展,一些新颖的神经网络结构很难自然地表示为纯计算图。

以 Transformer 结构的神经网络为例,来看看使用最自然地方式描述这些算法对 AI 框架会带来什么新的表示要求。Transformer 结构的神经网络算法中,图的左侧是一个通用 Transformer 结构的中关键步骤的计算示意图,通过堆叠 Transformer 结构使得网络模型层数更深,右侧对应了使用最自然的方式描述这一算法计算过程的伪代码。

Transformer 结构

Transformer 是一种基于注意力机制的神经网络结构,由多个 Encoder 和 Decoder 堆叠而成,可以应用于各种自然语言处理任务。在具体应用时,可以根据任务的特点和需求,选择不同的 Transformer 架构来构建模型。从伪代码描述中可以看到,想要以通用的方式,自然地描述出 Transformer 的算法框架,均依赖于循环控制逻辑 for

引入控制流将会使得计算图的构建以及前向传播带来很大的差异。

首先,计算图将变为动态的方式,分支选择以及循环控制流只有在真实运行的时候,才能够依据其依赖的数据输入来判断走哪个分支、是否结束循环。

其次,控制流引入的另一个难点在于循环控制流的实现。引入循环之后,原本的计算图在逻辑上出现了环,从而无法进行有效的拓扑排序。所以对于有控制流的计算图,前向计算和反向传播的实现要么抛弃拓扑排序这一思路,要么就要通过其他手段将循环进行拆解。

为了能够支持含有控制流结构的神经网络模型,AI 框架需要引入了对动态控制流这一语言结构(language construct)的支持。目前基于计算图,在控制流解决方案上,主要采用了三类设计思路:

  • 复用宿主语言:复用前端宿主语言的控制流语言结构,用前端语言中的控制逻辑驱动后端计算图的执行;
  • 支持控制流原语:AI 框架的后端对控制流语言结构进行原生支持,计算图中允许计算流和控制流混合;
  • 源码解析:前端对高级编程语言的代码进行解析称计算图,后端对控制流语言结构解析成子图,对计算图进行延伸。

复用宿主语言以 PyTorch 为典型代表,支持控制流原语以 TensorFlow 为典型代表,源码解析的方式则以 MindSpore 为典型代表。

2.5 挑战与未来

通过计算图可以把 AI 系统化的问题形象地表示出来。

AI 框架架构图

计算图的基本组成有两个主要的元素:1)基本数据结构张量和 2)基本计算单元算子。节点代表 Operator 具体的计算操作(即算子),边代表 Tensor 张量。整个计算图能够有效地表达神经网络模型的计算逻辑和状态。

  • 基本数据结构张量:张量通过 shape 来表示张量的具体形状,决定在内存中的元素大小和元素组成的具体形状;其元素类型决定了内存中每个元素所占用的字节数和实际的内存空间大小;
  • 基本运算单元算子:具体在加速器 GPU/NPU 中执行运算的是由最基本的代数算子组成,另外还会根据深度学习结构组成复杂算子。每个算子接受的输入输出不同,如 Conv 算子接受 3 个输入 Tensor,1 个输出 Tensor。

AI 框架的设计很自然地沿用了张量和张量操作,将其作为构造复杂神经网络的基本描述单元,开发者可以在不感知复杂的框架后端实现细节的情况下,在 Python 脚本语言中复用由后端优化过的张量操作。而计算 Kernel 的开发者,能够隔离神经网络算法的细节,将张量计算作为一个独立的性能域,使用底层的编程模型和编程语言应用硬件相关优化。

开发者编写神经网络模型代码的时候,需要 AI 框架提供自动微分功能,构建反向的计算图。基于反向模式(Reverse Mode)的自动求导机制中,根据反向计算图的构建时机,又分为基于对偶图的自动求导、基于表达式追踪和图层 IR 的自动求导实现的静态图,而使用动态图的 AI 框架多采用基于表达式追踪技术的自动求导方式。动态图和静态图的区别又引起不同的编程范式,和对编程领域引入的控制流表达不同方式。

动态图与静态图区别

  • 图表示

随着图、图结构、图数据、图算法、图计算、图学习、图神经网络等信息技术的快速发展,各行业数字化水平的逐步提高,企业业务环境和计算场景日益复杂,数据间的关联关系也正在变得更加交错。在面对需要深度挖掘数据间复杂关联关系的场景时,传统的关系型数据往往计算效率低下,难以满足应用需求。

针对图神经网络(GNN)模型的结构类型,大量稀疏的结构下,会适合使用计算图来表示图吗?如何通过计算图更加高效地计算图稀疏图结构?

  • 大数据融合

大数据技术演化的总体目标是高效收集、存储、处理与分析大规模、多源数据,并满足业务需求。近年来,大数据技术路线从批处理架构,逐渐演化为内存计算架构、流处理架构、批流融合处理架构、图数据处理架构等。

在大数据技术发展的早期,仅仅是在计算框架 MapReduce 中提供简单的作业调度算法,随着资源管理的需求,在 Hadoop 2.0 时代,Yarn 作为单独组件负责分布式计算框架的资源管理。但是,一方面,Yarn 仅仅能够管理调度计算框架的资源;另一方面,资源的管理粒度较为粗放,不能做到有效的资源隔离,越来越不能满足企业客户的需求。

未来 AI 框架的计算图如何与数据的处理进行融合与表示?大数据和 AI 融合,计算图会产生什么新的表示和需求?

图表示与大数据与 AI 融合

  • 部署推理

深度学习不仅包括使用 AI 框架模型训练过程,对训练好的模型,通常需要进行一些优化和加速,如网络层合并、量化、剪枝等,得到精简的推理引擎,然后部署上线。通常我们把模型优化和推理框架合称为推理引擎。在推理部署环节的优化进度较快,主要思路包括调整模型架构实现计算并行度提升,或通过捕捉参数结构实现混合精度推理,降低计算复杂度。

既然最终形态是面向推理部署,那么一定需要 AI 框架按照计算图设计后的逻辑明确的分层解耦?计算图不能解决哪些 AI 业务的问题?部署场景 AI 作为流程中的一部分,能否泛化计算图解决部署流程复杂问题?

  • 科学计算

自然界中诸多现象如材料损伤与破坏、流体湍流、核爆炸过程、生物大分子等均呈现出巨大的尺度效应, 并伴随着不同尺度上的物理多样性和强耦合性以及多个时间与空间尺度的强关联。这些典型的多尺度问题的求解一直是非常有挑战性的课题。科学计算曾经并正在为求解多尺度问题发挥重要作用,但目前也碰到了瓶颈。 AI 为解决多尺度问题提出了新的思路。

3. 分布式并行

分布式训练是一种模型训练模式,它将训练工作量分散到多个工作节点上,从而大大提高了训练速度和模型准确性。虽然分布式训练可用于任何类型的 AI 模型训练,但将其用于大模型和计算要求较高的任务最为有利。

三个主要组件:

  1. 分布式数据并行训练(DDP)是一种广泛采用的单程序多数据训练范式。在 DDP 中,模型会在每个进程上复制,每个模型副本将接收不同的输入数据样本。DDP 负责梯度通信以保持模型副本同步,并将其与梯度计算重叠以加速训练。
  2. 基于 RPC 的分布式训练(RPC)支持无法适应数据并行训练的通用训练结构,例如分布式流水线并行、参数服务器范式以及 DDP 与其他训练范式的组合。它有助于管理远程对象的生命周期,并将自动微分引擎扩展到单个计算节点之外。
  3. 提供了在组内进程之间发送张量的功能,包括集体通信 API(如 All Reduce 和 All Gather)和点对点通信 API(如 send 和 receive)。尽管 DDP 和 RPC 已经满足了大多数分布式训练需求,PyTorch 的中间表达 C10d 仍然在需要更细粒度通信控制的场景中发挥作用。例如,分布式参数平均,在这种情况下,应用程序希望在反向传播之后计算所有模型参数的平均值,而不是使用 DDP 来通信梯度。这可以将通信与计算解耦,并允许对通信内容进行更细粒度的控制,但同时也放弃了 DDP 提供的性能优化。

通过充分利用这些分布式训练组件,开发人员可以在各种计算要求和硬件配置下高效地训练大模型,实现更快的训练速度和更高的模型准确性。

3.1 数据并行

数据并行是一种广泛应用于分布式 AI 系统中的技术,旨在通过将数据集划分为多个子集并在不同计算节点上并行处理这些子集,以提高计算效率和速度。通过这种方法,计算任务可以被分摊到多个节点上,从而显著提高处理速度和效率。

数据并行的实现方式多种多样,按照同步方式进行分类,包括同步数据并行异步数据并行。同步数据并行要求所有计算节点在每一轮迭代后同步其参数,确保模型的一致性。而异步数据并行则允许节点独立进行计算和参数更新,从而减少等待时间,但也可能带来参数不一致的问题。按照实现方式进行分类,包括数据并行分布式数据并行完全分片的数据并行异步的数据并行弹性数据并行以及参数服务器

数据并行只能在单台机器上运行,采用单进程、多线程的实现方式,将原本在 NPU0 上进行的数据训练过程,扩展到多个 NPU 并行训练。在某 NPU 上随机初始化模型和优化器后,就可进行数据并行的训练,算法主要分为下面两个步骤:

  • 前向传播:将 mini-batch 数据平均分配到每个 NPU 上。接下来进行分布式初始化,将模型和优化器复制到每个 NPU 上,保证各 NPU 的模型、优化器完全相同。初始化完成后,各 NPU 根据分配到的数据和模型同时进行前向传播。
  • 损失计算与反向传播:前向传播完成后,每个 NPU 分别计算模型损失并进行反向传播。得到梯度后,将梯度传递到某 NPU 进行累加,更新模型的参数和优化器状态。更新后的模型参数和优化器将会在下一轮的前向传播中被复制到每个 NPU 上。

如图所示,不断上述步骤重复进行,直到模型收敛或者达到预定的训练轮数。

数据并行

但由于数据并行相对来说还不够完善,造成了许多性能的浪费。如在语言层面,使用作为最热门的深度学习开发语言 Python,在数据并行中采用的单进程、多线程并行方式往往受到 GIL(全局解释器锁)限制,CPU 的性能瓶颈使得多线程不能良好的利用 NPU 集群的资源。

另外在算法层面,全局的梯度累积和参数更新发生在一个 NPU 上,会出现明显的单个 NPU 利用率更高,其他 NPU 空闲的情况,造成了资源的浪费。同时如果在数据并行中的 mini-batch 设置过小,将导致 NPU 内并行度不足,从而降低训练速度;在通信开销的影响下,甚至可能出现比单 NPU 慢的情况。

分布式数据并行(Distributed Data Parallel, DDP)是数据并行的一种高级形式,它综合了多种优化,是当前应用最广的并行算法之一,通常用于大型 NPU AI 集群和 AI 系统中。

3.2 张量并行

模型并行(Model Parallelism, MP)是一种重要的方法。模型并行的基本思想是将模型的计算任务拆分到不同的设备上执行,以提高训练效率和处理更大规模的模型。下面将重点介绍模型并行中的张量并行。

张量并行(Tensor Parallelism,TP)是一种更细粒度的模型并行方法,它将单层内部的参数和计算任务拆分到不同的设备上执行,这种方法特别适合于具有大量参数的大规模模型。它最初是在 Megatron-LM 论文中提出的,它是一种高效的模型并行技术,可用于训练大型 Transformer 模型。

通过张量并行,可以将矩阵乘法等计算操作的矩阵按行或按列切分,然后在不同设备上并行执行部分计算,最后通过集合通信操作合并结果。张量并行可以分为 MatMul 并行、Transformer 并行、Embedding 并行、Cross Entropy Loss 并行。

其中序列并行(Sequence Parallel,SP)也是张量并行的一种变体,它在序列维度上对 nn.LayerNorm 或 RMSNorm 进行分割,以进一步节省训练过程中的激活内存。当模型变得越来越大时,激活内存就会成为瓶颈,因此在张量并行训练中,通常会将序列并行应用于 LayerNorm 或 RMSNorm 层。

张量并行的主要挑战在于如何切分参数和计算任务,以保证计算的一致性和通信的高效性。例如,在进行矩阵乘法时,必须确保各设备上的部分结果在数学上是一致的。此外,通信开销也是一个重要考虑因素,需要在计算和通信之间找到平衡点,以达到最佳性能。

3.3 流水并行

流水线并行(Pipeline Parallelism,PP)是一种将模型的不同层(layer)按顺序分配到不同设备上的方法。不同于朴素的模型并行,流水线并行通过将输入数据切分成多个微批次(micro-batch),使得每个设备可以在处理完当前批次后立即处理下一个批次,从而提高设备利用率。

主要集中在 Gpipe 流水线并行和 PipeDream 流水线并行上(基于 F-then-B 策略与 1F1B 策略),不过还有很多优秀的流水线并行实现方式,例如:PipeDream-2BW、PipeDream-Flush、PipeDream-Megatron-LM 等,但他们一般都在大规模分布式深度学习训练框架中使用,如:Megatron-LM 和 Deepspeed,而不是 AI 框架,因此并不作为讨论范围。

在 Gpipe 中,模型被分割成多个阶段,每个阶段在不同的设备上执行。输入数据也被切分成多个微批次,每个设备同时处理不同的微批次,从而提高并行效率。此外,Gpipe 也可以使用重计算策略,在前向和反向传播过程中节省内存。

模型并行

通过这种方法,可以更好地平衡各设备的负载,减少空闲时间。然而,由于任务分段的传递顺序,某些设备在等待前一任务完成时会有空闲时间。这种空闲时间被称为“气泡”。通过优化分段和任务分配,可以最小化气泡的影响,提高整体效率。

Gpipe 流水线并行提供了多项显著优势。它可以高效地利用计算资源。通过将模型分段并分配到多个设备上,充分利用各设备的计算能力,从而提高整体计算效率。其次可以减少内存需求。由于模型被分段,每个设备只需要存储当前分段的参数和激活值。这显著降低了每个设备的内存需求,使得可以在内存较小的设备上训练大模型。在启动激活检查点后,通过在流水线反向传播时重新计算激活,可以进一步压缩内存需求。

3.4 混合并行

混合并行(HybridParallel)是一种用于分布式计算的高级策略,它结合了数据并行和模型并行的优势,以更高效地利用计算资源,解决深度学习中的大模型训练问题。混合并行不仅能提高计算效率,还能在有限的硬件资源下处理更大的模型和数据集。在深度学习中,数据并行和模型并行各自有其适用的场景和局限性。数据并行适用于训练样本较多而模型较小的情况,通过将数据集分割成多个子集并在不同的设备上同时训练来提高训练速度。而模型并行则适用于模型较大无法单独放入一个设备内存的情况,通过将模型切分成不同的部分分别在多个设备上进行计算。混合并行通过将这两种并行方式结合,加速计算和处理超大模型,从而在各种硬件条件下实现高效的神经网络模型训练。现主流的混合并行为 3D 混合并行,但由于他们一般都在大规模分布式深度学习训练框架中使用,如:Deepspeed 和 Colossal AI,而不是 AI 框架,因此只进行简单讨论。

模型并行

混合并行是由多个并行方式组合而成的:

  • 数据并行:将数据集分割成多个子集,在多个设备上分别处理这些子集。这种方式能显著提高数据处理速度,但需要确保子集之间处理结果一致。
  • 张量并行:将模型的不同部分分配给不同的设备进行处理。这种方式能充分利用各种设备的计算能力,但需要注意设备之间的通信开销。
  • 流水线并行:将模型按层分割成若干块,每块都交给一个设备进行处理。这种方式能提高设备的利用率,但需要确保流水线中各阶段正确传递中间结果。

3D 混合并行 DP+PP+TP

3D 混合并行是一种在深度学习训练中常用的混合并行策略,它将数据并行、模型并行和流水线并行三种并行方式结合起来,以优化资源利用率和训练效率。这种策略尤其适用于处理大规模的神经网络模型。由于每个维度至少需要 2 个 GPU,因此在这里至少需要 8 个 GPU 才能实现完整的 3D 并行。

模型并行

数据、模型和流水线并行各自在提高内存和计算效率方面发挥特定作用。

  • 内存效率:模型的层被分为流水线阶段,每个阶段的层通过模型并行进一步划分。这种 2D 组合同时减少了模型、优化器和激活函数消耗的内存。然而,不能无限制地分割模型,否则会因通信开销而限制计算效率。
  • 计算效率:为了让 Worker 数量在不牺牲计算效率的情况下超越模型和流水线并行,使用 ZeRO 驱动的数据并行(ZeRO-DP)。ZeRO-DP 不仅通过优化器状态分区进一步提高了内存效率,还通过利用拓扑感知映射,使 GPU 数量的扩展具有最小的通信开销。
  • 拓扑感知 3D 映射:3D 并行中的每个维度都被仔细映射到 Worker 上,通过利用两个关键的架构属性实现最大计算效率。

通过流水线和模型并行,数据并行组通信的梯度大小线性减少,因此总通信量比纯数据并行减少。此外,每个数据并行组在一部分本地 Worker 之间独立且并行地执行通信。因此,数据并行通信的有效带宽通过通信量减少和本地化及并行性的结合而得到放大。

模型并行是三种策略中通信开销最大的,所以优先将模型并行组放置在一个节点中,以利用较大的节点内带宽。其次,流水线并行通信量最低,因此在不同节点之间调度流水线,这将不受通信带宽的限制。最后,若张量并行没有跨节点,则数据并行也不需要跨节点;否则数据并行组也需要跨节点。

值得注意的是 ZeRO,它是 DP 的超级可伸缩增强版,在完全分片的数据并行一节中已经讨论过了。通常它是一个独立的功能,不需要 PP 或 TP。但它也可以与 PP、TP 结合使用。当 ZeRO-DP 与 PP (以及 TP) 结合时,它通常只启用 ZeRO 阶段 1,只对优化器状态进行分片(ZeRO 阶段 2 还会对梯度进行分片,阶段 3 也对模型权重进行分片)。虽然理论上可以将 ZeRO 阶段 2 与流水线并行一起使用,但它会对性能产生不良影响。每个 micro batch 都需要一个额外的 reduce-scatter 通信来在分片之前聚合梯度,这会增加潜在的显著通信开销。根据流水线并行的性质,一般会使用小的 micro batch ,并把重点放在算术强度 (micro batch size) 与最小化流水线气泡 (micro batch 的数量) 两者间折衷。因此,增加的通信开销会损害流水线并行。

后记: AI的东西分享还是挺吃力,最近比较忙也停更了挺久。下面继续分析一下OS和SoC芯片中的知识,敬请期待。