AI学习笔记整理(61)——大模型分布式训练并行策略

大模型参数计算

参考链接:https://andyguo.blog.csdn.net/article/details/131971176

模型参数单位

"10b"、"13b"、"70b"等术语通常指的是大型神经网络模型的参数数量。其中的 "b" 代表 "billion",也就是十亿。表示模型中的参数量,每个参数用来存储模型的权重和偏差等信息。例如:

"10b" 意味着模型有大约 100 亿个参数。

"13b" 意味着模型有大约 130 亿个参数。

"70b" 意味着模型有大约 700 亿个参数。

例如:Meta 开发并公开发布的 Llama 2 系列大型语言模型 (LLM),这是一组经过预训练和微调的生成文本模型,参数规模从 70 亿(7b)到 700 亿(70b)不等。经过微调的LLMs(称为 Llama-2-Chat)针对对话场景进行了优化。

meta-llama/Llama-2-7b-hf

meta-llama/Llama-2-13b-hf

meta-llama/Llama-2-70b-hf

输入:仅输入文本

输出:仅生成文本

模型架构:Llama 2 是一种使用优化的 Transformer 架构的自回归语言模型。调整后的版本使用监督微调(SFT)和带有人类反馈的强化学习(RLHF)来适应人类对有用性和安全性的偏好。

训练显存计算

1.模型权重

模型权重是模型参数中的一部分,通常是指神经网络中连接权重(weights)。这些权重决定了输入特征与网络层之间的连接强度,以及在前向传播过程中特征的传递方式。

2.梯度

在训练过程中,计算梯度用于更新模型参数。梯度与模型参数的维度相同。

3.优化器参数

一些优化算法(如带有动量的优化器)需要保存一些状态信息,以便在每次更新时进行调整。这些状态信息也会占用一定的显存。比如:

  • 采用 AdamW 优化器:每个参数占用8个字节,需要维护两个状态。意味着优化器所使用的显存量是模型权重的 2 倍;
  • 采用 经过 bitsandbytes 优化的 AdamW 优化器:每个参数占用2个字节,相当于权重的一半;
  • 采用 SGD 优化器:占用显存和模型权重一样。

4.输入数据和标签

训练模型需要将输入数据和相应的标签加载到显存中。这些数据的大小取决于每个批次的样本数量以及每个样本的维度。

5.中间计算

在前向传播和反向传播过程中,可能需要存储一些中间计算结果,例如激活函数的输出、损失值等。

6.临时缓冲区

在计算过程中,可能需要一些临时缓冲区来存储临时数据,例如中间梯度计算结果等。减少中间变量也可以节省显存,这就体现出函数式编程语言的优势了。

7.硬件和依赖库的开销

显卡或其他硬件设备以及使用的深度学习框架在进行计算时也会占用一些显存。

【以 Llama-2-7b-hf 为例】

数据类型:Int8

模型参数: 7B * 1 bytes = 7GB

梯度:同上7GB

优化器参数: AdamW 2倍模型参数 7GB * 2 = 14GB

LLaMA的架构(hidden_size= 4096, intermediate_size=11008, num_hidden_lavers= 32, context.length = 2048),所以每个样本大小:(4096 + 11008) * 2048 * 32 * 1byte = 990MB

A100 (80GB RAM)大概可以在int8精度下BatchSize设置为50

综上总现存大小:7GB + 7GB + 14GB + 990M * 50 ~= 77GB

Llama-2-7b-hf模型Int8推理显存大小6.5GB, 由此可见,模型训练需要的显存是至少推理的十几倍。

推理显存计算

  • 模型结构: 模型的结构包括层数、每层的神经元数量、卷积核大小等。较深的模型通常需要更多的显存,因为每一层都会产生中间计算结果。
  • 输入数据: 推理时所需的显存与输入数据的尺寸有关。更大尺寸的输入数据会占用更多的显存。
  • 批处理大小BatchSize: 批处理大小是指一次推理中处理的样本数量。较大的批处理大小可能会增加显存使用,因为需要同时存储多个样本的计算结果。
  • 数据类型DType: 使用的数据类型(如单精度浮点数、半精度浮点数)也会影响显存需求。较低精度的数据类型通常会减少显存需求。
  • 中间计算: 在模型的推理过程中,可能会产生一些中间计算结果,这些中间结果也会占用一定的显存。

GPU计算背景知识

GPU计算是指利用图形处理单元(GPU)的强大并行处理能力来加速计算密集型任务,而不仅仅局限于传统的图形渲染。其核心背景在于GPU独特的架构设计,使其在处理大量可并行化数据时,远超传统的中央处理器(CPU)。

  • 并行计算架构‌:现代GPU包含成百上千个较小、更高效的核心,这些核心被设计为同时执行大量相似的计算任务。这与CPU的少量、功能强大的核心(专注于顺序处理和复杂逻辑)形成鲜明对比。这种"人海战术"式的并行架构,使得GPU在处理矩阵运算、向量计算等任务时具有压倒性优势。‌

  • ‌从图形渲染到通用计算‌:GPU的起源是加速3D图形渲染,但其高度并行的特性很快被发现适用于其他领域。自2007年NVIDIA推出CUDA(统一计算设备架构)编程模型以来,开发者得以使用GPU进行通用计算(GPGPU),从而开启了其在科学模拟、金融分析、数据分析等领域的广泛应用。‌

  • ‌异构计算成为主流‌:现代高性能计算系统普遍采用CPU与GPU协同工作的异构架构。CPU负责任务调度、逻辑判断和I/O控制,而GPU则承担大规模并行计算任务。这种分工协作能最大化系统整体性能。例如,世界顶尖的超级计算机Summit就配备了数万个NVIDIA Tesla GPU。‌

  • 内存和带宽:GPU有自己的专用内存,通常称为显存或VRAM。这种内存具有非常高的带宽,使得GPU能够快速地处理大量数据,这对于图形渲染和其他内存密集型任务至关重要。

  • 专用硬件加速‌:为应对人工智能等领域的爆发式需求,现代GPU(如NVIDIA的Tensor Core)集成了专门用于加速深度学习矩阵运算的硬件单元,进一步提升了在训练和推理大型神经网络时的效率。‌‌

  • 软件生态的支撑‌:NVIDIA通过CUDA及其丰富的软件库(如cuBLAS、cuDNN),为开发者提供了相对易用的编程接口。这使得主流深度学习框架(如TensorFlow、PyTorch)能够底层调用GPU加速,无需开发者深入掌握硬件细节,极大地推动了GPU计算的普及。‌

总而言之,GPU计算的背景是硬件架构的创新(海量并行核心)、软件生态的成熟(CUDA等)以及计算需求的演变(AI、大数据)共同作用的结果,使其成为驱动当前人工智能和高性能计算发展的关键基石。

GPU和CPU有什么区别?

GPU(图形处理单元)和CPU(中央处理单元)是现代计算机系统中的两种关键处理器,下表列举了CPU和GPU在关键特性上的对比,以帮助理解它们之间的主要差异。

计算指标与矩阵运算

参考链接:http://www.uml.org.cn/Computer/202407111.asp

1、指标
FLOPS

"FLOPS" 是 "Floating Point Operations Per Second" 的缩写,意为每秒浮点运算次数。它是衡量计算机系统或处理器性能的一种常见指标,特别是在科学计算、工程领域和高性能计算中。例如,一个处理器的性能为 1 TFLOPS,意味着它可以每秒执行 1 万亿次浮点运算。

MACs

Multiply--Accumulate Operations,乘加累积操作。1MACs包含一个乘法操作与一个加法操作(Ax+y)

优化:

去掉没有意义的MACs

降低每次MACs的计算时间

增加MACs并行计算能力

增加PE利用率

Accuracy

1.计算精度 (FP32/FP16 etc.)

2.模型结果精度 (ImageNet 78%)

优化:

能够处理各类型的无规则数据 >> 异构平台

能够应对复杂网络模型结构 >> 计算冗余性

Throughput

高维张量处理 (high dimension tensor)

实时性能 (30 fps or 20 tokens)

优化:

除了峰值算力,看计算单元的平均利用率 >> 负载均衡

SOTA网络模型的运行时间 >> MLPerf

Latency

生成第一个 Token 的时间(Time To First Token (TTFT)

生成每一个输出 Token 的时间(Time Per Output Token (TPOT)

Latency = TTFT + TPOT *(要生成的 Token 数 - 1)

Latency 可以转换为 Tokens Per Second (TPS):TPS = (the number of tokens to be generated) / Latency。

优化:

通信时延对 MACs 的影响 >> 优化带宽

Batch Size 大小与内存大小 >> 多级缓存设计

2、矩阵运算

常见矩阵运算:

  • 乘加
  • 重排列降维:convolution,Matrix Multiply
  • 分块-Tiling:
    • 根据 Cache 大小来对矩阵进行分块 Tiling,最大程度重用数据和利用空间换时间
    • GPU 上各种内存的访问速度为 Global memory << shared memory < register (local memory)。Global memory 大而慢, shared memory 小而快,因此减少内存访问延迟的一个重要方向就是要尽量减少 global memory 的访问,其中一个常见的策略就是 Tiling ------ 将数据分片,然后将每个小分片缓存到 shared memory 中。
    • Tiling 技术通过将 global memory 中的元素加载到 shared memory 中以便多次使用,从而减少了对 global memory 的访问次数;一般情况下,如果分片大小为 K×K 个元素,则 global memory 的访问次数会减少 K 倍
  • GEMM :通用矩阵乘,是一种广泛用于深度学习神经网络模型的计算操作,通常在软件层面实现,并利用现代处理器的FMA指令来加速运算。
  • FMA: 融合矩阵乘加,通过单个指令实现矩阵乘加。D=AB+C(tensor core)
  • 硬件
    • 现有库: Matrix Multiplication (GEMM)
      • CPU: OpenBLAS, Intel MKL, etc.
      • GPU: cuBLAS, cuDNN, etc.
    • 实现逻辑:
      • Lib 感知相乘矩阵的 Shape
      • 选择最优的 Kernel 实现来执行
    • 实现方法:
      • Loop 循环优化 (Loop tiling)
      • 多级缓存 (memory hierarchy)
        1. Macro-Kernel(宏内核): Macro-kernel 是指高级别的循环结构,负责管理多个 micro-kernel 的执行。在矩阵乘法中,它通常用于处理大规模的矩阵计算,例如管理整个矩阵乘法的批处理过程。
        2. Micro-Kernel(微内核): Micro-kernel 是指低级别的循环结构,负责执行实际的矩阵乘法操作。它通常用于管理单个矩阵块的计算,以实现更细粒度的并行化和优化。

GPU工作原理

GPU (Graphics Processing Units, GPUs),正如其全称"图形计算单元",GPU的初衷主要是为了接替CPU进行图形渲染的工作。因为图像上的每一个像素点都需要处理,这项任务计算量相当大。尤其遇上一个复杂的三维场景,就需要在一秒内处理几千万个三角形顶点和光栅化几十亿的像素。不过,由于每个像素点处理的过程和方式相差无几,这项艰巨的任务可以靠并行计算来化解。

GPU的性能取决于其内部的核心数量、时钟频率以及RAM容量。GPU包含多个核心,每个核心都可以独立执行命令。每个核心具有自己的存储器,用于存储指令和数据。

GPU的工作原理流程如下:

  1. GPU会接收来自CPU的指令,并把它们分发到多个核心中进行处理。
  2. GPU会把处理后的数据传输回 CPU,以便CPU可以使用它。
线程原理

kernel

主设概念: CUDA引入主机端(host)和设备(device)概念。CUDA 程序中既包含host程序,又包含device程序。

互相通信: host与device之间可以进行通信,这样它们之间可以进行数据拷贝。

CUDA 执行流程中最重要的一个过程是调用CUDA的核函数来执行并行计算,kernel是CUDA中一个重要的概念。

在 CUDA 程序构架中,Host 代码部分在CPU上执行,是普通C代码;当遇到数据并行处理的部分,CUDA 就会将程序编译成GPU能执行的程序,并传送到GPU,这个程序在CUDA里称做核(kernel)。

CUDA Kernel函数:是数据并行处理函数(核函数),在GPU上执行时,一个Kernel对应一个Grid,基于GPU逻辑架构分发成众多thread去并行执行。kernel 用 __global__符号声明,在调用时需要用 <<<grid, block>>> 来指定kernel要执行及结构。

CUDA Stream流:Cuda stream是指一堆异步的cuda操作,他们按照host代码调用的顺序执行在device上。

thread

在GPU编程中,线程是最基本的执行单元。一个GPU程序会同时启动成千上万个线程来执行任务

warp

GPU 的 每一行由1个控制单元加上若干计算单元所组成,这些所有的计算单元执行的控制指令是一个 。这其实就是个非常典型的 "单指令多线程机制(SIMT)" 。

单指令多线程机制是说:多个线程同时执行相同的指令序列,但是每个线程可以处理不同的数据。这些线程通常被分组成更小的线程块,每个线程块中的线程可以协调执行相同的指令。

warp是硬件级别上的调度单位,一个 warp 包含32个并行 thread,这些 thread 以不同数据资源执行相同的指令。

Block

  • 块是线程的集合,它们被组织成一个工作单元,可以共享内存和同步。
  • 在GPU编程中,通常会将线程划分成若干个块,以便更有效地管理和协调线程的执行。
  • 块内的线程可以进行协作和通信,通常通过共享内存来提高性能。
  • Block 间并行执行,并且无法通信,也没有执行顺序。

Grid

多个block则会再构成grid。kernel 在 device 上执行时,实际上是启动很多线程,一个kernel 所启动的所有线程称为一个网格(grid)。同一个网格上的线程共享相同的全局内存空间。

当一个 kernel 被执行时,grid 中的线程块被分配到 SM (多核处理器) 上,一个线程块的 thread 只能在一个SM 上调度,SM 一般可以调度多个线程块,大量的 thread 可能被分到不同的 SM 上。每个 thread 拥有它自己的程序计数器和状态寄存器,并且用该线程自己的数据执行指令,这就是所谓的 Single Instruction Multiple Thread (SIMT)。

每个 thread 都有自己的一份 register 和 local memory 的空间。同一个 block 中的每个 thread 则有共享的一份 share memory。此外,所有的 thread (包括不同 block 的 thread) 都共享一份 global memory。不同的 grid 则有各自的 global memory。

从软件的角度来讲:

1.线程处理器 (SP) 对应线程 (thread)。

2.多核处理器 (SM) 对应线程块 (thread block)。

3.设备端 (device) 对应线程块组合体 (grid)。

缓存机制

硬件

GPU 内存硬件的分类,按照是否在芯片上面可以分为片上 (on chip) 内存和片下 (off chip) 内存。

片上内存 主要用于 缓存 (cache) 以及少量特殊存储单元(如texture)。特点是速度快,存储空间小;

片下内存 主要用于全局存储 (global memory) 即常说的显存 ,特点是速度相对慢,存储空间大。不同于 CPU 系统内存可扩展的设计,GPU 的整体内存大小是固定的,在选择好显卡型号后就会确定好,包括缓存和全局存储。

在磁盘/硬盘(Disk/SSD)上面的数据传入到 GPU 的内存要经过: 硬盘 -> 系统内存 -> GPU 内存 的过程。这个速度非常慢,要极力避免这种传输。

系统存储:

  • L1/L2/L3:多级缓存,位置在 CPU 芯片内部;
  • System DRAM:动态 RAM,CPU 芯片外部内存,如内存条
  • Disk/Buffer:外部存储,如磁盘或者固态硬盘。

GPU 设备存储:

  • L1/L2 cache:多级缓存,位置在 GPU 芯片内部;
  • GPU DRAM:通常所指的显存;

传输通道:

  • PCIE BUS:PCIE标准的数据通道,数据就是通过该通道从显卡到达主机;
  • BUS: 总线。计算机内部各个存储之间交互数据的通道;
  • PCIE-to-PCIE:显卡之间通过PCIE直接传输数据;
  • NVLINK:NVIDIA 公司推出的、用于显卡之间的专用的超高速数据传输通道。

PCIE 6.0 的最大理论传输带宽为 128GB/s ;而第四代 NVLINK 的能够提供 GPU 之间 900GB/s 的带宽。二者有接近一个数量级的差距。

架构

在 NVIDIA A100 中,HBM Memory(80GB) 就是我们通常说的显存,在这里我们把一些寄存器文件(Register File)也当作缓存。实际执行单元(SM)希望能够快速获取数据,于是实际执行单元会从寄存器中读取 L2 Cache 的内容。另外一方面呢,希望 L2 Cache 与 显存更近,当 L2 Cache 未命中时,GPU 会从显存中寻找数据。如果在显存还没有找到,就需要通过 PCIe 总线在内存中寻找,但由于 PCIe 的内存带宽非常低(比显存带宽低20倍),因此这会导致时延的大大增加。

多级缓存的形式降低了由于内存传输导致的时延,提高了GPU的算力利用率和总体计算时间。同时,缓存的带宽和容量也在不断增大,以满足GPU对高速数据访问的需求。

按照存储功能进行细分,GPU 内存可以分为:局部内存(local memory)、全局内存(global memory)、常量内存(constant memory)、共享内存(shared memory)、寄存器(register)、L1/L2 缓存等。

其中全局内存、局部内存、常量内存都是片下内存,储存在 HBM 上。所以我们说 HBM 的 大部分 作为全局内存。

关于 SRAM 与 DRAM :

RAM 分为静态 RAM(SRAM)和动态 RAM(DRAM)。SRAM 只要存入数据后,即使不刷新也不会丢失记忆;而 DRAM 的电容需要周期性地充电,否则无法确保记忆长存。

DRAM 密度高、成本低、访问速度较慢、耗电量大。SRAM 则刚好相反。因此 SRAM 首选用于带宽要求高,或者功耗要求低的情境。如:CPU Cache、GPU On-Chip Buffer。DRAM 则一般用于系统内存、显存。

片上 (on chip) 内存-SRAM
L1/L2缓存

L2 缓存可以被所有 SM 访问,速度比全局内存快;L1 缓存用于存储 SM 内的数据,被 SM 内的 CUDA cores 共享,但是跨 SM 之间的 L1 不能相互访问。

合理运用 L2 缓存能够提速运算。A100 的 L2 缓存能够设置至多 40MB 的持续化数据 (persistent data),能够拉升算子 kernel 的带宽和性能。Flash attention 的思路就是尽可能地利用 L2 缓存,减少 HBM 的数据读写时间。

寄存器

寄存器(register)是线程能独立访问的资源,它是片上(on chip)存储,用来存储一些线程的暂存数据。寄存器的速度是访问中 最快 的,但是它的容量较小,只有几百甚至几十 KB,而且要被许多线程均分。

共享内存

共享内存(shared memory) 是一种在线程块内能访问的内存,是片上(on chip)存储,访问速度较快。

共享内存主要是缓存一些需要反复读写的数据。

注:共享内存与 L1 缓存的位置、速度极其类似,区别在于共享内存的控制与生命周期管理与 L1 不同:共享内存受用户控制,L1 受系统控制。共享内存更利于线程块之间数据交互。

片下 (off chip) 内存
全局内存

全局内存(global memory)能被 GPU 的所有线程访问,全局共享。它是片下(off chip)内存,前面提到的硬件 HBM 中的大部分都是用作全局内存。跟 CPU 架构一样,运算单元不能直接使用全局内存的数据,需要经过缓存。

局部内存

局部内存 (local memory) 是线程独享的内存资源,线程之间不可以相互访问。局部内存属于片下内存,所以访问速度跟全局内存一样。它主要是用来应对 寄存器不足 时的场景,即在线程申请的变量超过可用的寄存器大小时,nvcc 会自动将一部数据放置到片下内存里。

常量内存

常量内存(constant memory)是片下(off chip)存储,但是通过特殊的常量内存缓存(constant cache)进行缓存读取,它是只读内存。

常量内存主要是解决一个 warp scheduler 内多个线程 访问相同数据 时速度太慢的问题。假设所有线程都需要访问一个 constant_A 的常量,在存储介质上 constant_A 的数据只保存了一份,而内存的物理读取方式决定了多个线程不能在同一时刻读取到该变量,所以会出现先后访问的问题,这样使得并行计算的线程出现了运算时差。常量内存正是解决这样的问题而设置的,它有对应的 cache 位置产生多个副本,让线程访问时不存在冲突,从而保证并行度。

SM(SP)

从 G80 提出的概念,中文称流式多处理器,核心组件包括CUDA核心、共享内存、寄存器等。SM包含许多为线程执行数学运算的Core,是 NVIDA 的核心。

在CUDA中, 可以并发地执行数百个线程。一个 block 上线程是放在同一个 SM,一个 SM 的有限 Cache 制约了每个 block 的线程数量。

主要包括:

  1. CUDA Core:向量运行单元 (FP32-FPU、FP64-DPU、INT32-ALU); 最开始叫SP(Streaming Processor) , 是GPU最基本的处理单元,在fermi架构开始被叫做 CUDA core 。Volta 架构时期取消了CUDA Core,变成了单独的FPU 和ALU。
  2. Tensor Core:张量运算单元(FP16、BF16、INT8、INT4);
  3. Special Function Units:特殊函数单元 SFU(超越函数和数学函数,e.g. 反平方根、正余弦等);
  4. Warp Scheduler:线程束调度器(XX Thread / clock);
  5. Dispatch Unit:指令分发单元(XX Thread / clock);
  6. Multi level Cache:多级缓存(L0/L1 Instruction Cache、L1
  7. Data Cache & Shared Memory);
  8. Register File:寄存器堆;
  9. Load/Store:访问存储单元LD/ST(负责数据处理);
    GPC ------ 图形处理簇,Graphics Processing Clusters
    TPC ------ 纹理处理簇,Texture Processing Clusters
    SM ------ 流多处理器,Stream Multiprocessors
    HBM ------ 高带宽存储器,High Bandwidth Memory
    包含关系为:GPC > TPC > SM > CORE

大模型的分布式训练

参考链接:https://is-cloud.blog.csdn.net/article/details/155017798

大模型的分布式训练是指利用多台机器(通常每台包含多个计算设备,如GPU)协同工作,来完成传统单机单卡无法处理的超大规模深度学习模型的训练任务。其核心目的是解决模型参数量巨大(如千亿参数以上)或训练数据量极其庞大时,单台设备的显存容量和计算算力不足的瓶颈问题。‌

分布式训练通过将模型、数据或计算任务进行拆分和并行化处理,显著提升了训练效率和可扩展性。在当前的大模型时代,它已成为主流的训练范式。‌

假设神经网络中某一层是做矩阵乘法, 其中的输入 x xx 的形状为 4 × 5 4 \times 54×5, 模型参数 w ww 的形状为 5 × 8 5 \times 85×8, 那么, 矩阵乘法输出形状为 4 × 8 4 \times 84×8 。如下:

1. 数据并行

  • 数据并行:切分数据x到不同设备上,在反向传播过程中,需要对各个设备上的梯度进行 AllReduce,以确保各个设备上的模型始终保持一致。适合数据较大,模型较小的情况,如resnet50
  • allReduce: 先将所有device上的模型的梯度reduce归约,然后将结果广播到所有设备上。

梯度 AllReduce 是深度学习分布式训练中用于‌梯度同步‌的一种核心集体通信操作。‌它本质上是 ‌AllReduce‌ 算法在梯度聚合场景下的具体应用,其主要目的是确保在多机多卡训练中,所有参与计算的设备(如 GPU)都能获得一致的全局梯度,从而同步更新模型参数。‌

其工作流程通常包含两个关键步骤:

  • 归约 (Reduction)‌:所有设备将其计算出的局部梯度进行聚合,最常见的操作是‌求和‌,有时也可能是求平均值。‌
  • 广播 (Broadcast)‌:将归约后得到的全局梯度,分发回给‌所有‌参与计算的设备。‌

通过这一过程,每个设备都能获得由所有设备梯度信息汇总而成的全局梯度,进而使用这个一致的梯度来更新本地的模型参数,保证了分布式训练的一致性和收敛性。‌在实际框架中,如 PyTorch 的分布式数据并行(DDP),loss.backward() 之后会自动调用 AllReduce 操作来完成梯度同步。‌

但是这种数据并行有哪些限制呢?

第一个是可以使用的GPU数量受限于batch的大小,假设batch是64,那你最多也只能用64张卡了。

另一个就是GPU利用率可能拉不满,如果batch的数量固定式512,你GPU太多,分发给每一个GPU的数据量太小,那GPU 更新块,但是通信频率也就增加了,可能会限制训练速度。

所有 GPU 必须一起等到梯度计算完成,才能开始 all-reduce,否则会造成阻塞卡顿。 数据并行最适合模型较小、计算量大、batch size 较大时使用。模型较大或 GPU 太多时要考虑混合并行或 ZeRO 分布式技术。 All-Reduce 是一种通信操作,由 NCCL / MPI / Gloo 等通信库负责实现。它不是库,而是库提供的功能。

2. 模型并行

模型并行:省去了多个设备之间的梯度 AllReduce;但是, 由于每个设备都需要完整的数据输入,因此,数据会在多个设备之间进行广播,产生通信代价。比如,上图中的最终得到的 out ( 4 × 8 ) (4 \times 8)(4×8) ,如果它作为下一层网络的输入,那么它就需要被广播发送到两个设备上。如bert

3. 流水并行

流水并行:模型网络过大,除了用模型并行,还能用流水并行(如下的网络有4层,T1-T4)

4. 混合并行

混合并行:综合上面的多种策略一起训练,如GPT3:

  • 首先被分为 64 个阶段,进行流水并行。每个阶段都运行在 6 台 DGX-A100 主机上。

  • 在6台主机之间,进行的是数据并行训练;

  • 每台主机有 8 张 GPU 显卡,同一台机器上的8张 GPU 显卡之间是进行模型并行训练。

5.张量并行

张量并行 (Tensor Parallelism, TP) / Intra-layer Model Parallelism:将模型中单个大矩阵运算(如Transformer中的QKV投影、MLP层)切分到多个GPU上并行计算。

实现方式:

  • 列并行 (Column Parallelism): 对于 Y = XA 这样的运算,如果A按列切分 (A = [A1, A2]),则 Y = [XA1, XA2]。XA1 和 XA2 可以在不同GPU上计算。这常用于Transformer的QKV投影和MLP的第一层。输入X在各GPU上是相同的(通过广播或All-Gather得到),输出Y是切分的。
  • 行并行 (Row Parallelism): 对于 Y = XA 这样的运算,如果X按行切分 (X = [X1; X2]),则 Y = [X1A; X2A]。但通常在Megatron中,行并行用于 Z = YB,其中Y是前一个列并行操作的输出(已切分),B按行切分。每个GPU计算 YiBi,然后通过 All-Reduce 将结果汇总。这常用于Transformer的Attention输出投影和MLP的第二层。输入Y是切分的,输出Z是聚合后的完整张量(或准备好给下一个列并行层)。
  • 通信: TP内部的GPU通常在同一节点内,通过高速NVLink连接,通信开销相对较低但仍不可忽视。主要通信操作是 All-Reduce, All-Gather, Reduce-Scatter。

优点:

直接减少了每个GPU上的模型参数和激活值存储。

对于非常宽的层(hidden size很大)非常有效。

缺点:

通信频繁,对通信带宽要求高。

扩展性受限于单节点内的GPU数量和NVLink带宽。

6.序列并行

序列并行 (Sequence Parallelism, SP) (Megatron-LM后续加入的优化):在张量并行组内,对于Transformer中那些在序列维度上独立的运算(如LayerNorm、Dropout、逐元素操作),可以将输入张量的序列维度也进行切分。

实现方式: 将输入激活的序列维度 (sequence_length, batch_size, hidden_size) 中的 sequence_length 分割到TP组内的不同GPU上。

通信: 在需要完整序列信息的层(如Attention中的QKV矩阵乘法之后)之前进行All-Gather,在可以重新分片的地方进行Reduce-Scatter。

优点:

进一步减少了TP组内每个GPU上的激活值存储,尤其对于长序列非常有效。

与TP结合,可以更好地利用TP组内的计算资源。

缺点: 增加了额外的通信和同步点。

Flash Attention 技术

参考链接:https://zhuanlan.zhihu.com/p/1985797637404914691

Flash Attention 是一种用于加速 Transformer 模型中注意力机制计算的算法,它通过优化内存访问和计算流程,在不牺牲精度的前提下显著提升了训练和推理速度,并降低了显存占用。

其核心思想是将注意力计算分解为小块(Tile),并利用 GPU 的高速缓存(SRAM)来减少与慢速显存(HBM)之间的数据传输。

主要特点和优势:‌

  • 无损加速‌:与一些近似方法不同,Flash Attention 能够提供与标准注意力完全相同的计算结果,但速度更快。 ‌
  • 内存高效‌:通过避免将大型中间注意力矩阵写入显存,它将内存使用从序列长度的二次方降低到线性,这对于处理长文本序列至关重要。 ‌
  • IO 感知‌:优化重点放在减少数据移动开销上,而非仅仅降低浮点运算量。 ‌

‌核心机制:‌

  • 分块计算(Tiling)‌:将长序列的查询(Q)、键(K)、值(V)和输出(O)矩阵分割成较小的块。 ‌‌
  • SRAM 利用‌:将这些小块加载到速度更快的 GPU 共享内存(SRAM)中进行计算。 ‌‌
  • 重计算‌:在计算过程中,不存储完整的中间结果,而是根据需要重新计算部分值,从而节省显存。 ‌‌
  • 在线算法优化‌:在计算 softmax 函数时,采用特殊技巧(如维持最大值和累加和)来避免数值不稳定,并将除法操作延迟到最后一步。 ‌

为什么需要Flash Attention

Flash Attention 的出现主要是为了解决 Transformer 模型在处理长序列时面临的‌计算效率低‌和‌显存消耗大‌两大核心瓶颈。 ‌

具体来说,传统 Transformer 的自注意力机制存在以下问题:

  • 计算复杂度高‌:自注意力的计算量和空间复杂度均为 (O(N^2)),其中 (N) 为序列长度。这意味着当序列长度增加时,计算量和所需的显存会以平方级增长,严重限制了模型处理长文本的能力。 ‌
  • 显存访问成为瓶颈‌:在计算注意力时,需要生成并存储一个大小为 (N \times N) 的注意力矩阵。在反向传播时,又需要重新读取该矩阵。这不仅消耗大量显存,更导致了频繁且昂贵的高带宽内存(HBM)访问,使得计算过程常常受限于数据读写速度而非 GPU 的计算能力。

Flash Attention 通过一系列创新设计,针对性地解决了这些问题:

  • 加速计算(Fast)‌:它并未减少总的浮点运算次数(FLOPs),而是采用 ‌Tiling(分块计算)‌ 和 ‌算子融合(Kernel Fusion)‌ 技术,将计算过程集中在 GPU 的高速缓存(SRAM)中进行,从而‌大幅减少了与慢速显存(HBM)之间的数据交换次数‌,显著降低了"墙钟时间"。 ‌
  • 节省显存(Memory-efficient)‌:通过一种巧妙的‌在线算法‌,Flash Attention 在计算过程中‌避免了显式地实例化和存储整个 (N \times N) 的注意力矩阵‌。它只保存必要的统计量,将显存复杂度从 (O(N^2)) 降低到 (O(N)),使得训练更长的序列成为可能。 ‌
  • 保持精确性(Exact Attention)‌:与许多通过稀疏化或近似来加速注意力的方法不同,Flash Attention 的计算结果与原始的、精确的注意力机制‌完全等价‌,不会引入近似误差,保证了模型的性能。 ‌
相关推荐
奥特曼_ it6 小时前
【数据分析+机器学习】基于机器学习的招聘数据分析可视化预测推荐系统(完整系统源码+数据库+开发笔记+详细部署教程)✅
笔记·数据挖掘·数据分析
高工智能汽车6 小时前
爱芯元智通过港交所聆讯,智能汽车芯片市场格局加速重构
人工智能·重构·汽车
大力财经6 小时前
悬架、底盘、制动被同时重构,星空计划想把“驾驶”变成一种系统能力
人工智能
A9better6 小时前
嵌入式开发学习日志50——任务调度与状态
stm32·嵌入式硬件·学习
四维碎片7 小时前
QSettings + INI 笔记
笔记·qt·算法
非凡ghost7 小时前
ESET NupDown Tools 数据库下载工具
学习·软件需求
梁下轻语的秋缘7 小时前
Prompt工程核心指南:从入门到精通,让AI精准响应你的需求
大数据·人工智能·prompt
FreeBuf_7 小时前
ChatGPT引用马斯克AI生成的Grokipedia是否陷入“内容陷阱“?
人工智能·chatgpt
福客AI智能客服7 小时前
工单智转:电商智能客服与客服AI系统重构售后服务效率
大数据·人工智能
柳鲲鹏7 小时前
OpenCV:超分辨率、超采样及测试性能
人工智能·opencv·计算机视觉