深度学习笔记——模型训练工具(DeepSpeed、Accelerate)

454 阅读6分钟

大家好,这里是好评笔记,本文为试读,查看全文请移步公主号:Goodnote。本文详细介绍深度学习模型训练中常用的工具及其特点,包括DeepSpeed和Accelerate。

15.png

@[toc]

DeepSpeed

DeepSpeed 是由微软开源的一个深度学习大规模分布式训练,特别适合大规模预训练模型(如 GPT、BERT、T5 等)的高效训练。DeepSpeed 提供了一系列技术和工具,以提升模型的计算效率、内存使用效率和训练速度。它不仅支持模型的分布式训练,还包含一些用于训练和推理的高效技术。

DeepSpeed 的核心功能

ZeRO 优化器(Zero Redundancy Optimizer)

  • 作用:通过消除冗余的数据存储和计算,大幅减少训练大规模模型时的内存需求,使得超大规模模型的训练成为可能。
  • 原理:将优化器状态梯度模型参数分布到多个 GPU 中,每个 GPU 只存储一部分数据,通过分布式计算实现更高的内存利用率
  • ZeRO 阶段性优化:ZeRO 优化器分为多个阶段,每个阶段逐步减少内存冗余,实现更高的内存效率。
    • ZeRO Stage 1:分布式存储优化器状态。每个 GPU 只存储模型的部分优化器状态,这样多个设备可以共享优化器的负担。
    • ZeRO Stage 2:分布式存储梯度信息。模型的梯度分布到不同设备中,而不是每个设备都存储完整的梯度。
    • ZeRO Stage 3:分布式存储模型参数。每个 GPU 仅保留部分模型参数,使得整个模型可以在多 GPU 上进行分布式存储,从而大幅减少显存占用。
  • 显存高效性:ZeRO 的三级优化可以将显存需求减少到原本的很小比例,使得在有限显存的硬件上训练数百亿到数千亿参数的大规模模型成为可能。

混合精度训练(Mixed Precision Training)

  • 作用:使用半精度(FP16) 进行训练,在保证模型精度的同时大幅减少显存占用和计算成本。
  • 特点:DeepSpeed 在混合精度训练中采用 NVIDIA 的 Apex 技术,支持 FP16 和 FP32 的混合精度计算,有效加速模型训练。

1. 混合精度实现 传统上,神经网络的训练使用 32位浮点数(FP32)来表示模型参数计算结果。这种高精度表示在某些任务中可能并不总是必要的,而且会带来显存和计算上的高成本。DeepSpeed混合精度训练通过在模型中引入较低精度的数值表示(如16 位浮点数,FP16)来减小显存消耗和加速计算。

  • DeepSpeed 主要使用 FP16(半精度浮点数) 和 FP32(单精度浮点数) 进行混合精度训练,降低了显存需求和计算量。
    • FP16(半精度浮点数):具有更低的存储需求(16 位),用于模型的前向和反向传播计算,以及大部分参数和激活函数的存储。用于加速计算
    • FP32(单精度浮点数):用于存储模型的关键权重累计梯度等更敏感的数据。保证精度和稳定性
  • 它通过动态损失缩放、优化器的低精度支持,深度集成硬件加速等技术,确保在混合精度下的数值稳定性。

其中的 大部分权重和激活函数除了使用FP16 也可以使用 BF16

FP16 和 BF16 的对比如下:

特性FP16(半精度浮点数)BF16(BFloat16)
位数结构16 位(1 符号位,5 指数位,10 尾数位)16 位(1 符号位,8 指数位,7 尾数位)
指数范围较小,容易出现下溢或上溢与 FP32 相同,数值范围较大
精度的尾数精度,精度优于 BF16精度略于 FP16
数值稳定性,特别是在处理大动态范围数据时,更适合大动态范围的任务
硬件支持NVIDIA V100、A100、RTX 系列等大部分 GPUNVIDIA A100、H100 和 Google TPU
损失缩放需求需要动态损失缩放(Dynamic Loss Scaling)通常不需要损失缩放
  • FP16:适合广泛的硬件环境和常规任务,特别适用于在大多数 GPU 上进行的计算密集型任务。需要结合损失缩放来解决数值不稳定的问题。
  • BF16:适合需要高数值稳定性的任务,如在 A100 或 TPU 上训练大规模模型,可以替代 FP16,且通常无需损失缩放。BF16 在支持硬件上提供更高的稳定性,更适合于大动态范围的模型训练。

2. 自动混合精度 AMP

DeepSpeed 的混合精度训练不是简单地降低模型参数精度,而是会进行自动选择性的调整精度来实现高效且稳定的混合精度训练:

  • 自动混合精度(Automatic Mixed Precision, AMP):DeepSpeed 会在适当的时候自动将模型的参数和梯度转换为 16 位。例如,在大多数计算(如前向和反向传播)中将参数和梯度转换为 FP16 或 BF16,以提高速度并减少内存消耗;在需要更高精度的步骤(如关键参数的更新)中使用 FP32

3. 动态损失缩放 在混合精度训练中,特别是使用 FP16 时,由于数值表示的范围变小,梯度可能会下溢(变为零)或上溢(变得过大)。DeepSpeed 使用动态损失缩放动态调整损失的缩放因子,以确保数值稳定性。

动态损失缩放核心思路在反向传播中将损失值乘以一个损失缩放因子,从而放大梯度的数值。这个因子在整个训练过程中会动态调整,以确保梯度在 FP16 精度下计算时不会因为数值过小而变为零。

假设损失值为 ( LL ):

  • 通过 缩放因子 ( SS ),在反向传播前,将损失值放大 ( SS ) 倍,即使用放大的损失值 ( L×SL \times S ) 来计算梯度。
  • 计算得到的梯度也会被放大 ( S ) 倍。然后,在更新参数前,将梯度除以 ( SS ) 还原。

这种缩放只影响梯度的计算过程,而不影响参数的更新,确保模型学习过程中梯度的有效性。

详细全文请移步公主号:Goodnote。

参考:欢迎来到好评笔记(Goodnote)!