优雅谈大模型:白话ZeRO 下

文摘   科技   2024-08-31 13:35   上海  

庞德公
编辑郭嘉

机器学习模型的复杂性和规模不断增长,分布式训练变得比以往任何时候都更加重要。训练具有数千亿参数的大型语言模型( LLMs )将是机器学习基础设施面临的挑战。与传统的分布式计算框架不同的地方在于GPU的分布式训练需要将数据传递给GPU芯片等物理硬件层。GPU设备之间会进行频繁、大规模的数据交换以进行高效训练,今天将揭开分布式训练的神秘面纱。

上图为基本的机器学习训练框架,数据准备占据1/3,训练占据1/3,将中间版本进行保存以及调整参数占据1/3。在模型的训练过程中,若读者拥有多片GPU,那么如何让这些GPU同时工作,目前有数据并行,模型并行以及管道并行。

数据并行

数据并行是跨多个 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 ,这些框架都提供并行方法的实现。


--->更多内容,请移步“鲁班秘笈”!!<---

鲁班模锤
基于开源技术生态,跟踪与普及人工智能、隐私计算、区块链以及数据空间的最新技术。着力于数据治理自动化,依托语料库构建数据可信流通方案,为所有企业提供平等的数据增益服务。
 最新文章