分布式训练

分布式训练

基本概念

• 单个服务器上一般存在CPU(中央处理单元),GPU(图形处理单元),TPU(张量处理单元)三种设备加速器
• 多设备加速比:多设备加速比通常是通过将训练任务拆分到多个设备上并行执行，然后与在单个设备上执行相同任务的时间进行比较来计算的.衡量了分布式训练中的性能优势.
• 分布式训练存在的挑战:
  • 计算墙:单个计算设备所能提供的计算能力与大语言模型所需的总计算量之间存在巨大差异。
  • 显存墙：单个计算设备无法完整存储一个大语言模型的参数。
  • 通信墙：分布式训练系统中各计算设备之间需要频繁地进行参数传输和同步。由于通信的延迟和带宽限制，这可能成为训练过程的瓶颈。

目标

• 因为目前单个机器上有限的资源已无法满足大语言模型训练的对算力要求.需要将一个模型训练任务拆分成多个子任务,并将子任务分发给多个计算设备,从而解决资源瓶颈.

原理

• 分布式训练并行策略

  • 数据并行:对数据进行切分,在多个设备上使用同一个模型并行执行不同的数据分片;而在基于Transformer架构的大语言模型中,因为基本上每个算子(自注意力机制)都是对单个token进行处理,所以数据并行时不会涉及到同步问题

  

  • 模型并行:对模型(以神经网络层为单位)进行划分,将不同的算子(神经网络层)分发到多个设备分别完成

    • 流水线并行:按模型的层切分到不同设备,即层间并行或算子间并行

      • GPipe方法:通过将小批次**(Mini-batch)数据进一步划分为微批次(Micro-batch)**来缓解朴素流水线策略存在的并行气泡问题.但存在需要等待一个微批次所有前向计算完成才能进行后向计算
        

      • 1F1B流水线策略:主要分为热身阶段,前向-后向阶段,后向阶段.可以保证流水线的满负荷运转.

        

    • 张量并行:将计算图层内的参数切分到不同设备,实现层内并行或算子内并行.Transformer结构主要的算子为Embedding,矩阵乘和交叉熵损失,主要的张量并行策略也着重于这三方面.

      • Embedding:对于Embedding层的参数,对于嵌入表示层的参数，可以按照词维度切分， 每个计算设备只存储部分词向量，然后通过汇总各个设备上的部分词向量，从而得到完整的词向量。主要通过划分word_size到不同的设备进行张量并行.
      • 矩阵乘:通过将参数矩阵分块到不同计算设备上的方式进行张量并行,分块方式主要有按行分块和按列分块.

    

    • 交叉熵损失:将logit矩阵(全连接层结果)按照类别维度切分到不同的计算设备上,同时通过中间结果(每部分数据的交叉熵损失)通信,得到最终的全局的交叉熵损失.

  • 混合并行:同时使用多种并行策略:例如在针对千亿规模的LLM时,通常在每个服务器内部使用张量并行策略,再通过流水线并行,将模型的不同层划分为多个阶段,每个阶段有不同的机器负责计算.最后通过数据并行,将训练数据分发到多组服务器上进行并行处理,每组服务器处理不同的数据批次.

  

相关技术

分布式训练的集群架构

​ 计算集群的硬件架构:在LLM的分布式训练中,因为其参数量过大,为了保证模型参数在各个服务器间的有效通信,舍弃了以往传统的多层树架构(交换机-交换机-机柜),选择采取胖树拓扑结构,试图实现网络带宽的无收敛.

​ 计算集群的架构需要根据分布式系统、大语言模型结构、优化算法等综合因素进行设计.常见的架构一般有:参数服务器架构和去中心化架构

• 参数服务器架构:主要包含训练服务器和,参数服务器.训练方式主要包含同步训练和异步训练.

  • 同步训练:训练服务器在完成一个小批次的训练后，将梯度推送给参数服务器。参数服务器 在接收到所有训练服务器的梯度后，进行梯度聚合和参数更新。整体训练较慢
  • 异步训练:异步训练：训练服务器在完成一个小批次的训练后，将梯度推送给参数服务器。但是参数服 务器不再等待接收所有训练服务器的梯度，而是直接基于已接收到的梯度进行参数更新。异步更新的机制会导致训 练效果有所波动。

  

• 去中心化架构:此架构中没有中央服务器或控制节点,而是由节点之间进行直接通信和协调.节点间并行的进行训练和通信,可以显著降低通信开销减少通信墙的影响.

  • 节点之间需要周期性地交换参数更新和梯度信息。可以通过集合通信技术来实现.主要的通信原语有:Broadcast、Scatter、Reduce、All-Reduce、Gather、All-Gather、Reduce-Scatter、All-to-All 等.他们表示了节点间参数传递的各种方式.

  

  • 通信库:Pytorch等DL框架通常通过通信库对网络硬件进行操作.常用的通信库有MPI,GLOO和NCCL,根据不同的通信库的在CPU,GPU上的对通信原语的支持情况不同,不同的分布式集群所采用的通信库也会不同.

    • MPI:是一种广泛使用的并行计算通信库，常用于 在多个进程之间进行通信和协调。
    • Gloo:FaceBook 推出的一个类似 MPI 的集合通信库,大体遵照 MPI 提供的接口规定,，支持 CPU 和 GPU 上的分布式训练。
    • NCCL:NVIDIA 开发的高性能 GPU 间通信库，专门用于在多个 GPU 之间进行快速通信和同步,在 NVIDIA 硬件 上，NCCL 的效果往往比其它的通信库更好.

  因此,在 CPU 集群上进行训练时，可选择使用 MPI 或 Gloo 作为通信库；而如果是在 GPU 集群上进行训练，则 可以选择 NCCL 作为通信库。