数据并行
数据并行是跨多个 GPU 扩展模型训练最常用的方法之一。在这种方法中,每个GPU设备都需要存储全尺寸模型并处理不同批次的训练数据。除了存储模型和训练数据之外,每个GPU设备还保存临时结果,其中包括前向传播输出和梯度下降结果。
每个副本完成前向传递阶段后, NCCL All Reduce<白话ZeRO>会跨GPU节点同步模型参数,确保模型更新就像单个节点处理过所有批次数据。
ZeRO数据并行<白话ZeRO>是对传统数据并行的增强,旨在减少 GPU内存使用。与每个GPU存储完整模型不同,ZeRO跨GPU分割模型参数。这种方法仍然是数据并行的一种形式,因为每一层的张量计算不会跨GPU分割。GPU负责持久化模型分片,并可以在完成必要的计算后丢弃它们。
这种方法显着减少了内存消耗,使得可以在相同的硬件上训练更大的模型。PyTorch 提供完全分片数据并行 ( FSDP ),它将模型参数、优化器状态和跨分布式数据并行 ( DDP ) 等级的梯度进行分片,因此进一步减少内存消耗。
数据并行简单直接,同时可以将训练分布在多GPU。特别大的模型而言,即便仅仅加载大型模型的单层也会对GPU的内存带来很大的压力。
在单GPU上训练模型时,若模型超过GPU内存,则训练过程中需要频繁地通过HostToDevice操作在主机内存和GPU内存之间不断地交换模型参数。这种频繁的I/O操作会导致GPU空闲,从而拖累了整个训练过程。
张量并行是模型并行的一种,模型参数被分割到多个GPU,从而实现并行计算。模型参数具备矩阵运算特点,可以切分按行或按列划分,允许每个 GPU独立执行乘法操作。最终通过将不同的GPU计算的子结果合并就可以得到最终的运算结果,也可以充分的利用所有的GPU设备。下图很直观的展示将矩阵分割,两路并行计算。
管道并行
除了数据并行和张量并行之外,管道并行可以进一步提高训练期间的GPU利用率。在管道并行性中,模型根据其层被分为块,允许前向传播和后向传播分布在不同的GPU上。这种方法可能会导致GPU出现空闲等待时间。
若模型足够小,单个GPU能够搞定,则可以使用数据并行将其扩展到多个节点。随着模型大小的增加,可能需要张量并行才能将模型分布到单个节点内的多个 GPU 上。如果模型变得更大,可以在同一节点内应用张量并行,而在不同节点之间使用管道并行。
当然需要确保参与管道并行的节点位于同一网络等级内以实现最佳 I/O 性能至关重要。对于读者而言,小编建议先从数据并行入手,熟悉和感受下分布式训练方法。目前随着分布式计算越发的成熟,涌现出许多流行的训练框架,例如PyTorch Distributed Data Parallel ( DDP )、 DeepSpeed和Megatron-LM ,这些框架都提供并行方法的实现。