大模型-计算优化-DeepSpeed

这是一个用来理解DeepSpeed是什么的概念的笔记，关于具体的细节没有特别详细。

Introduction

大模型训练都会用到，用来对多 GPU 中的计算、存储、通信做优化，从而在有限的硬件上训练更大的模型。
由微软开发。

原理

ZeRO (Zero Redundancy Optimizer)

Introduction

深度学习模型通常在多个 GPU 上进行分布式训练。在常规的数据并行策略中，每个 GPU 都存储模型的完整副本、梯度和优化器的状态，这会导致大量的冗余存储。

ZeRO 通过三个阶段的优化来减少这种冗余。它在多个 GPU 之间划分和共享模型参数、梯度和优化器状态，从而显著减少每个 GPU 的内存使用量。

这允许更大的模型被拟合进 GPU，并加速了模型的训练，因为在多个 GPU 之间的数据通信更少了。

ZeRO 的三个阶段

ZeRO-0

最基础的配置，它与传统的数据并行训练非常相似。
每个GPU都存储完整的模型参数、梯度和优化器状态。这个阶段并没有减少任何冗余。

ZeRO-1 (ZeRO-DP)

模型的梯度被分片并存储在各个GPU中，这样每个GPU只需要存储其分片的梯度。

模型参数和优化器状态仍然在每个GPU上完整地存储。

梯度需要同步：虽然每个GPU只存储了梯度的一个分片，但在进行模型更新之前，这些分片梯度仍然需要在所有GPU之间同步以得到完整的梯度。在ZeRO-1中，这是通过一个高效的reduce-scatter和all-gather通信操作来完成的。

ZeRO-2 (ZeRO-SP)

除了梯度分片，优化器状态也被分片并存储在各个GPU中。

模型参数仍然在每个GPU上完整地存储。

ZeRO-3 (ZeRO-MP)

模型参数、梯度和优化器状态都被分片并分布在各个GPU中，是最高级的优化。

也就是说，为了训练更大的模型，可以直接配置 ZeRO-3

参数、梯度和优化器状态的动态移动，以及通信开销

由于所有这些都是分片存储的，因此在执行前向和后向传播时，需要跨 GPU 动态地移动这些片段以确保正确的计算。这会引入一定的通信开销，但由于减少了内存需求，所以在大多数情况下，这种开销是值得的。

激活检查点和重计算

Introduction

为了进一步减少内存使用，DeepSpeed 提供了激活检查点策略，它保存了前向传播中的某些中间激活值，并在反向传播中重计算它们。

用更多的计算换取了在GPU内存中可以存储更大的模型，对大模型来说非常有用。

activation 的内存开销

在模型的前向传播和反向传播过程中，所有层的激活都必须存储在内存中，因为反向传播需要它们来计算梯度。对于大型和/或深层模型，这些激活可以占用大量的GPU内存。

activation 数量会远超模型参数本身

前向传播时需要存储 2 种数值

1. 模型参数
2. 每条样本输入到模型后在每一层的计算结果
   也就是这里的 activation

每条样本在每一层都要计算 activation。
如果模型层数多、参数多，样本多，那么 activation 的数量是巨大的，activation 数量会远超模型参数本身的数量

Activation Checkpoint 激活检查点

与传统的训练方法相比，其中所有激活都被保存下来，激活检查点只保存某些层的激活。这些特定的激活被称为“检查点（checkpoint）”

重计算

由于我们在前向传播中没有保存所有的激活，所以在反向传播时，我们需要重新计算这些丢失的激活。

当我们在反向传播中到达一个检查点时，我们从该检查点开始，再次进行前向传播，重新计算丢失的激活。

这样就可以用额外的计算时间来换取内存的节省。

计算过程

从 loss 开始逐层反向计算每个参数的梯度，用链式法则逐层求导。

当求导求到一个具体的参数值时，从反向传播过来的部分是清楚的，也就是用链式法则求导时上层落下来的，参数的系数；

参数自己这一层的数值是不知道的，需要从参数的前面保存了activation 的 checkpoint 重新计算，并前向传递到参数这里。

混合精度训练

Introduction

使用 16 位浮点数（而不是传统的 32 位）来存储模型参数和计算梯度，从而加速训练和减少内存使用。DeepSpeed 还提供了动态损失缩放，以防止在 16 位训练中发生数值不稳定。

计算梯度、正向反向传播时使用 FP16

可以提高计算速度，因为较低的精度意味着更简单的算术运算。

始终保持一份主权重

为了防止精度损失，通常会在FP32精度中保留一份模型权重的复制（称为“主权重”）

梯度累积与权重更新

梯度可以在FP16中计算，但权重更新通常在FP32中进行，以保持精度。
计算得到的FP16梯度首先转换为FP32，然后应用到主权重上，之后再将更新的权重转回FP16。

动态损失缩放

使用FP16可能会导致数值不稳定性，尤其是当梯度值很小并导致下溢时。

此时可以将损失乘以一个较大的数（缩放因子），从而将梯度值提升到一个较安全的范围。

如果在训练中检测到梯度值溢出，则缩放因子会动态调整。

这里应该还有更多细节，但是暂时没有必要了解。

某些特定的操作可能仍然在FP32中执行，因为它们在FP16中可能不稳定。

一个常见的例子是批量归一化（batch normalization）。

通信优化

提供了优化的通信算法，例如 1-bit Adam，它可以显著减少在多个 GPU 之间同步数据所需的时间。

异步 I/O

通过异步 I/O 来确保 GPU 在训练过程中始终被充分利用，从而减少由于等待数据导致的延迟。

Offload to CPU

允许将某些数据，如优化器状态，卸载到 CPU 内存或甚至 NVMe 存储，从而进一步节省 GPU 内存。

可以通过参数配置来决定 offload both optimizer states and params to NVMe, or just one of them or none.

NVMe（Non-Volatile Memory express）：一个协议和接口标准，专为固态驱动器（SSD）设计，旨在充分发挥其高速性能。

自动地并行化模型

这部分没有看太懂，不理解它和其他的模块的联系和异同，需要了再细看。

Introduction

模型并行性意味着模型的不同部分（例如，神经网络的不同层）需要被分配到不同的 GPU 上。

2 modules

要实现模型，就需要先分割模型，然后再让不同部分之间交叉通信。

分割模型

将模型的不同部分做分割，然后分配给不同的设备。一下是几种常见的分割策略

1. 纵向分割
   将模型的连续层分配给不同的设备。例如，在一个五层模型中，前两层可以放在一个GPU上，而后三层可以放在另一个GPU上。
2. 横向分割
   对于有大量参数的模型组件（例如，一个大型的全连接层或一个转换器的自注意力头），可以将参数和计算分割到多个设备上。例如，一个全连接层的权重可以被分割成两半，每半放在一个GPU上。
3. 流水线并行性
   这是一个更复杂的策略，其中不同设备负责不同的前向和反向传播阶段。当一个设备完成其前向传播的部分并将结果传递给下一个设备时，它可以开始进行反向传播，从而在多个设备上同时进行前向和反向传播。

交叉通信

不同 GPU 之间交换数据

ZeRO 和模型并行化的关系

两者是独立的。

不同点

1. 目标
   ZeRO旨在通过减少冗余存储来提高数据并行性的效率。它不涉及跨GPU的模型计算分解。而模型并行性主要关注的是如何将一个大型模型的计算分布到多个GPU上。
2. 存储vs计算
   ZeRO的主要目标是优化存储（参数、梯度、优化器状态），而模型并行性主要关注的是如何将模型的计算分布到多个GPU上。
3. 通信
   在ZeRO中，跨GPU的通信主要是为了同步分片的参数和梯度。在模型并行性中，跨GPU的通信是为了满足模型各部分之间的计算依赖。

是哪个层面实现的

与CUDA的关系

DeepSpeed运行时依赖CUDA，因为它是为在NVIDIA GPU上运行的深度学习模型而设计的。
DeepSpeed底层会直接或间接地与CUDA API交互，以实现对GPU的高效操作。

与PyTorch的关系

1. DeepSpeed为PyTorch提供了一个易于使用的接口，允许研究者和开发者轻松地将DeepSpeed的优化集成到他们的PyTorch代码中。

2. DeepSpeed提供了许多PyTorch原生不包括的扩展性和优化功能，如ZeRO、模型并行性和内存优化。

   通过DeepSpeed和PyTorch的集成，用户可以用很少的代码更改来训练非常大的模型，这些模型在没有DeepSpeed优化的情况下可能很难或不可能训练。

DeepSpeed处理模型的梯度、优化器、参数等时，是用 pytorch 来处理的吗？

虽然DeepSpeed与PyTorch紧密集成，但在处理这些内容时，DeepSpeed会使用其自己的方法和实现，而不仅仅依赖PyTorch的默认机制。

1. 梯度
   尽管梯度的计算可能仍然使用PyTorch的自动微分功能，但存储和同步策略是DeepSpeed特有的。
2. 优化器
   DeepSpeed提供了对常见优化器（如Adam和Lamb）的优化版本。这些优化器考虑了ZeRO分片策略，并进行了针对性的调整以支持大规模并行性。
   这些实现可能与PyTorch原生的优化器有所不同，但目的是为了更好地与DeepSpeed的其他功能集成。
3. 参数
   与梯度类似，DeepSpeed的ZeRO策略还对模型参数进行了特殊处理，将其分片到多个GPU上。
   这种参数管理方式与PyTorch的默认策略不同，但可以显著减少单个GPU的存储需求。

用什么语言实现

使用多种编程语言，主要有 2 种：

1. Python
   DeepSpeed的上层API和许多高级功能都是使用Python实现的，以便与其他深度学习框架（如PyTorch）进行集成。
2. C++/CUDA
   为了实现高效的计算和并行处理，许多底层操作和核心功能都是使用C++和CUDA来编写的。

如何使用

DeepSpeed的接口使用很容易，但是需要指定配置文件

pytorch

使用时只需要在 nn.module 和 optimizer 外面包一层 deepspeed即可，其他的就像使用原生 pytorch 一样，示例如下：

import deepspeed
model, optimizer, _, _ = deepspeed.initialize(model, optimizer, deepspeed_config="ds_config.json")

transformers

huggingface 的 transformers 库集成了 DeepSpeed，可以在命令行启动时增加 DeepSpeed 参数，或者在创建 Trainer 时添加 DeepSpeed 参数。详细见transformers文档

参数配置

配置参数首先需要对DeepSpeed的实现原理有大概的理解，ZeRO、激活检查点重计算等，然后参考其他的公开出来的配置即可。（纯属个人意见）