[源码分析] Facebook如何训练超大模型---(4)

Python微信订餐小程序课程视频

https://edu.csdn.net/course/detail/36074

Python实战量化交易理财系统

https://edu.csdn.net/course/detail/35475

[源码分析] Facebook如何训练超大模型 — (4)

目录* [源码分析] Facebook如何训练超大模型 — (4)
+ 0x00 摘要
+ 0x01 背景知识
- 1.1 单精度、双精度和半精度浮点格式的区别
- 1.2 多精度和混合精度计算的区别
- 1.3 混合精度
- 1.4 训练过程
+ 0x02 PyTorch
- 2.1 英伟达算力
- 2.2 Torch.cuda.amp
* 2.2.1 使用
* 2.2.2 多Model，losses和优化器
* 2.2.3 分布式
+ 0x03 FSDP 使用
- 3.1 成员变量
- 3.2 Scaler
- 3.3 初始化
- 3.4 重建
- 3.5 cast操作
- 3.6 _post_reduction_hook
+ 0xFF 参考

0x00 摘要

我们在前文介绍过，微软 ZeRO 可以对一个万亿参数模型可以使用 8 路模型并行、64 路管道并行和 8 路数据并行在 4,096 个 NVIDIA A100 GPU 上进行扩展。而FSDP（Fully Sharded Data Parallel）是Facebook 深度借鉴微软ZeRO之后提出的PyTorch DDP升级版本，可以认为是对标微软 ZeRO，其本质是 parameter sharding。Parameter sharding 就是把模型参数等切分到各个GPU之上。我们会以 Google，微软和 Facebook 的论文，博客以及代码来进行学习分析。

之前文章之中我们谈到了FSDP支持混合精度训练，所以我们再来看看相关知识。

本系列其他文章如下：

[源码解析] PyTorch 分布式之 ZeroRedundancyOptimizer

[论文翻译] 分布式训练 Parameter sharding 之 ZeRO

[论文翻译] 分布式训练 Parameter Sharding 之 Google Weight Sharding

[源码分析] Facebook如何训练超大模型—(1)

[源码分析] Facebook如何训练超大模型 — (2)

[源码分析] Facebook如何训练超大模型 — (3)

0x01 背景知识

1.1 单精度、双精度和半精度浮点格式的区别

我们从NVIDIA官博 What’s the Difference Between Single-, Double-, Multi- and Mixed-Precision Computing?摘录如下：

IEEE 浮点算术标准是在计算机上用二进制表示数字的通用约定。在双精度格式中，每个数字占用 64 位。单精度格式使用 32 位，而半精度只有 16 位。

在传统的科学记数法中，pi 写为 3.14 x 10010^0。但是计算机将这些信息以二进制形式存储为浮点数，即表示数字及其相应指数的一系列 1 和 0，在本例中为 1.1001001 x212^1。

在单精度 32 位格式中，一位用于判断数字是正数还是负数。为指数保留了八位，指数（因为它是二进制的）是 2 的某个幂。剩余的 23 位用于表示组成数字的数字，称为有效数。

相反，双精度为指数保留 11 位，为有效数保留 52 位，大大扩展了它可以表示的数字的范围和大小。半精度占据了更小的部分，只有 5 个位用于指数，10 个位用于有效数。

以下是 pi 在每个精度级别的样子

1.2 多精度和混合精度计算的区别

多精度计算意味着使用能够以不同精度进行计算的处理器——在需要时使用双精度，并依赖于应用程序的其他部分的半精度或单精度算法。

混合精度，也称为超精度，计算改为在单个操作中使用不同的精度级别，以在不牺牲精度的情况下实现计算效率。在混合精度中，计算从快速矩阵数学的半精度值开始。但是随着数字的计算，机器以更高的精度存储结果。例如，如果将两个 16 位矩阵相乘，则答案大小为 32 位。

使用这种方法，当应用程序完成计算时，累积的答案在准确度上可与在双精度算术中运行整个事情相媲美。这种技术可以将传统双精度应用程序的速度提高多达 25 倍，同时减少运行它们所需的内存、运行时间和功耗。它可用于 AI 和模拟 HPC 工作负载。

1.3 混合精度

采用FP16的优势如下：

• 内存占用更少。如果采用FP16，则模型占用是FP32的一半，这样可以训练更大的模型，使用更大的batch size，通信量更少。
• 计算更快。FP16的加速优化可以加快训练和推理的计算。
• 另外，随着NVIDIA Tensor Core 的普及，FP6计算也越来越快。

FP16的问题主要是其表示范围比FP32狭窄，所以会带来两个问题：溢出错误 和 舍入误差。因此，百度和NVIDIA联手在论文之中提出了一些技术。

• 保留一份FP32格式的权重主备份。
• 使用loss scale来避免梯度过小。
• 使用FP16计算但是用FP32进行累加。

比如，对于主备份，论文之中图例如下：

1.4 训练过程

上面介绍的三种技术对于训练过程是一个良好的补充，我们从NVIDIA官方文档 https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html 摘录训练过程具体如下。

1. 维护一份FP32的参数主副本。
2. 对于每次迭代：
   1. 制作权重的FP16副本。
   2. 使用FP16权重和激活进行向前传播。
   3. 将得到的损失乘以比例因子S。
   4. 使用FP16权重，激活和它们的梯度进行后向传播。
   5. 将权重梯度乘以1/S。
   6. 完成权重更新（包括gradient clipping等）。

一个更稳健的方法是动态地选择损失比例因子。其基本思想是以一个大的比例因子开始，然后在每次训练迭代中重新考虑它。如果在选定的迭代次数N中没有发生溢出，则增加比例因子。如果发生溢出，则跳过权重更新，降低比例因子。我们发现，只要不频繁地跳过更新，训练计划就不必调整，就可以达到与FP32训练相同的精度。

请注意，N有效地限制了我们可以溢出和跳过更新的频率。缩放因子的更新率可以通过选择增加/减少的乘数以及N（增加前的非溢出迭代次数）来调整。

动态损失缩放方法对应了了以下训练流程：

1. 在FP32中保持一份权重的主副本。
2. 将S初始化为一个大的数值。
3. 对于每个迭代
   1. 制作一个