DeepSpeed基础及内存优化特性

DeepSpeed

1.基础概念

DeepSpeed是一个由Microsoft 公司开发的开源深度学习优化库，旨在提高大规模模型训练的效率和可扩展性，使研究人员和工程师能够更快地迭代和探索新的深度学习模型和算法。它采用了多种技术手段来加速训练，包括模型并行化、梯度累积、动态精度缩放和本地模式混合精度 等。此外，DeepSpeed 还提供了一些辅助工具，例如分布式训练管理 、内存优化和模型压缩，以帮助开发者更好地管理和优化大规模深度学习训练任务。

1. DeepSpeed 提供了分布式计算框架，首先需要明确几个重要的基础的概念：节点编号、全局进程编号、局部进程编号、全局总进程数和主节点。

• DeepSpeed 主节点（master_ip+master_port）负责协调所有其他节点和进程的工作，由主节点所在服务器的IP 地址和主节点进程的端口号来确定主节点。主节点还负责监控系统状态、处理任务分配和结果汇总等任务，因此是整个系统的关键部分。
• 节点编号（node_rank）是系统中每个节点的唯一标识符，用于区分不同计算机之间的通信。
• 全局进程编号（rank）是整个系统中的每个进程的唯一标识符，用于区分不同进程之间的通信。
• 局部进程编号（local_rank）：是单个节点内的每个进程的唯一标识符，用于区分同一节点内的不同
  进程之间的通信。
• 全局总进程数（word_size）是整个系统中运行的所有进程的总数，用于确定可以并行完成多少工作以及需要完成任务所需的资源数量。

2. 在网络通信策略方面，DeepSpeed 提供了MPI、GLOO 和NCCL 等选项，可以根据具体情况进行选择和配置。DeepSpeed 配置文件中，在optimizer 部分配置通信策略，如下是使用OneBitAdam优化器的配置样例，配置中其中使用了nccl 通讯库：

主要包含三个部分：

• APIs：DeepSpeed 提供了易于使用的API 接口，简化了训练模型和推断的过程。用户只需通过调用几个API 接口即可完成任务。通过"initialize"接口可以初始化引擎，并在参数中配置训练参数和优化技术等。这些配置参数通常保存在名为"ds_config.json/deepspeed.json"的文件中。
• RunTime：DeepSpeed 的核心运行时组件，使用Python 语言实现，负责管理、执行和优化性能。它承担了将训练任务部署到分布式设备的功能，包括数据分区、模型分区、系统优化、微调、故障检测以及检查点的保存和加载等任务。
• Ops：DeepSpeed 的底层内核组件，使用C++ 和CUDA 实现。它优化计算和通信过程，提供了一系列底层操作，包括Ultrafast Transformer Kernels、fuse LAN kernels、Customary Deals等。Ops 的目标是通过高效的计算和通信加速深度学习训练过程。

2.补充概念

1.常见浮点数表示方法：

2.当前大语言模型训练通常采用Adam 优化算法，

除了需要每个参数梯度之外，还需要一阶动量（Momentum）和二阶动量（Variance）。虽然Adam 优化算法相较SGD 算法通常效果更好也更稳定，但是对计算设备内存的占用显著增大。为了降低内存占用，大多数系统已经采用了混合精度训练（Mixed Precision Training）方式，即同时存在FP16（16 位浮点数）或者BF16（Bfloat16）和FP32（32 位浮点数）两种格式的数值。

混合精度优化过程：

Adam 优化器状态包括采用FP32 保存的模型参数备份，一阶动量和二阶动量也都采用FP32 格式存储。

假设模型参数量为Φ，模型参数和梯度都是用FP16格式存储，则共需要2Φ + 2Φ + (4Φ + 4Φ + 4Φ) = 16Φ 字节存储。其中Adam 状态占比75%。动态损失缩放反向传播前，将损失变化（dLoss）手动增大2K 倍，因此反向传播时得到的激活函数梯度则不会溢出；反向传播后，将权重梯度缩小2K 倍，恢复正常值。

举例来说 ，对于包含75 亿个参数模型，如果用FP16 格式，只需要15GB 计算设备内存，但是在训练阶段模型状态实际上需要耗费120GB。计算卡内存占用中除了模型状态之外，还有剩余状态（ResidualStates），包括激活值（Activation）、各种临时缓冲区（Buffer）以及无法使用的显存碎片（Fragmentation）等。由于激活值可以用检查点（Activation Checkpointing）方式使得激活值内存占用大幅度减少，因此如何

减少模型状态尤其是Adam 优化器状态是解决内存占用问题的关键。

3.主要特性

零冗余优化器（Zero Redundancy Data Parallelism，ZeRO）目标就是针对模型状态的存储 进行去除冗余的优化。ZeRO 使用分区的方法，即将模型状态量分割成多个分区，每个计算设备只保存其中的一部分。这样整个训练系统内只需要维护一份模型状态，减少了内存消耗和通信开销。具体来说，如下图所示，ZeRO 包含以下三种方法：

1.对优化器进行分区 ；（对应下图Pos，对应DeepSpeed框架下为ZeRO-1）

模型参数和梯度依然是每个计算设备保存一份。此时，每个计算设备所需内存是4Φ+ 12Φ/N 字节，其中N 是计算设备总数。当N 比较大时，每个计算设备占用内存趋向于4ΦB，也就是原来16ΦB 的1/4。

2.对模型梯度进行分区 ；（对应下图Pos+g，对应DeepSpeed框架下为ZeRO-2）

模型参数依然是每个计算设备保存一份。此时，每个计算设备所需内存是2Φ + (2Φ+12Φ)/N 字节。当N 比较大时，每个计算设备占用内存趋向于2ΦB，也就是原来16ΦB 的1/8。

3.对模型参数也进行分区 ;（对应下图Pos+g+p，对应DeepSpeed框架下为ZeRO-3）

此时，每个计算设备所需内存是16Φ/N*B。当N比较大时，每个计算设备占用内存趋向于0。

如DeepSpeed使用ZeRO-3配置文件：

{
"zero_optimization": {
"stage": 3,
},
"fp16": {
"enabled": true
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 0.001,
"betas": [
0.8,
0.999
],
"eps": 1e-8,
"weight_decay": 3e-7
}
},
...
}

4.ZeRO Offload

将优化器状态和计算转移到CPU上,在配置文件中添加如下代码：

{
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu"
}
},
...
}

也可进一步将模型参数装到CPU内存当中，在配置文件中添加如下代码：

{
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu"
}
"offload_param": {
"device": "cpu"
}
},
...
}

5.重计算

---

日常学习总结，欢迎交流