【复现DeepSeek-R1之Open R1实战】系列8：混合精度训练、DeepSpeed、vLLM和LightEval介绍

这里写目录标题

1 混合精度训练
- 1.1 FP16和FP32
- 1.2 优点
- 1.3 存在的问题
- 1.4 解决办法
2 DeepSpeed
3 vLLM
- 3.1 存在的问题
- 3.2 解决方法
- - 3.2.1 PagedAttention
  - 3.2.2 KV Cache Manager
  - 3.2.3 其他解码场景
- 3.3 结论
4 LightEval
- 4.1 主要功能
- 4.2 使用方法
- 4.3 应用场景

本文继续深入了解Open R1项目中用到的相关技术，包括训练模型用到的混合精度训练（bfloat16）、DeepSpeed、vLLM，以及评估模型用到的LightEval。

1 混合精度训练

1.1 FP16和FP32

单精度float（32bit，4个字节）如下图所示：

float32

半精度float（16bit，2个字节）如下图所示：

float16

半精度float的存储空间是float32的一半，float16的组成分为了三个部分：

最高位表示符号位sign；
有5位表示exponent位;
有10位表示fraction位。

从其bitmap计算出表示的数字的方式如下：

如果 Exponent 位全部为0：
- 如果 fraction 位全部为0，则表示数字 0
- 如果 fraction 位不为0，则表示一个非常小的数字（subnormal numbers）：(-1)^signbit×2^-14×(0+fraction/1024)
如果 Exponent 位全部位1:
- 如果 fraction 位全部为0，则表示 ±inf
- 如果 fraction 位不为0，则表示 NAN
Exponent 位的其他情况：(-1)^signbit×2^{(exponent-15)}×(1+fraction/1024)

1.2 优点

混合精度训练，指的是单精度 float和半精度 float16 混合，这样的优点是：

内存占用更少：fp16模型占用的内存只需原来的一半，占用的内存越小，训练的时候可以用越大的batchsize，同时通信量（特别是多卡或多机多卡）和等待时间大幅减少。
计算更快：目前不少GPU都有针对 fp16 的计算优化，半精度的计算吞吐量可以是单精度的 2-8 倍。

1.3 存在的问题

但是，fp16也存在明显的问题：

溢出错误（Overflow / Underflow）：fp16 的动态范围 ( 6×10^-8 ~ 65504) 比FP32的 ( 1.4×10^-45 ~ 1.7×10³⁸) 要小很多，因此在计算过程中很容易出现上溢出（Overflow）和下溢出（Underflow）的错误，溢出之后就会出现“NAN”的问题。对于深度学习而言，最大的问题在于 Underflow（下溢出），在训练后期，例如激活函数的梯度会非常小，在深度学习中，由于激活函数的的梯度往往要比权重梯度小，更易出现下溢出的情况。
舍入误差（Rounding Error）：舍入误差指的是当梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败，如下图所示：

舍入误差

1.4 解决办法

混合精度训练（Mixed Precision）：混合精度训练的精髓在于“在内存中用FP16做储存和乘法从而加速计算，用FP32做累加避免舍入误差”。混合精度训练的策略有效地缓解了舍入误差的问题。
损失放大（Loss Scaling）：即使用了混合精度训练，还是会存在无法收敛的情况，原因是激活梯度的值太小，造成了下溢出（Underflow）。损失放大的思路是：
- 反向传播前，将损失变化（dLoss）手动增大 2^k 倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；
- 反向传播后，将权重梯度缩 2^k 倍，恢复正常值。

2 DeepSpeed

DeepSpeed的核心就在于，GPU显存不够，CPU内存来凑。

它将当前时刻，训练模型用不到的参数，缓存到CPU中，等到要用到了，再从CPU挪到GPU。这里的“参数”，不仅指的是模型参数，还指optimizer、梯度等。越多的参数挪到CPU上，GPU的负担就越小；但随之的代价就是，更为频繁的CPU，GPU交互，极大增加了训练推理的时间开销。因此，DeepSpeed使用的一个核心要义是：时间开销和显存占用的权衡。

DeepSpeed 是由微软开发的一个用于深度学习模型训练的优化库，提供了多种技术来优化深度学习模型的训练过程，其中最著名的技术之一是 ZeRO，它是一种减少分布式训练中冗余的方法。

ZeRO（Zero Redundancy Optimizer）分为0、1、2、3四个不同的阶段，对应DeepSpeed 参数是 zero_stage，每个阶段都建立在前一个阶段的基础上，并提供额外的内存节省能力。

Stage 0：这是不采用任何内存优化方案的状态，相当于普通的分布式数据并行（Data Parallelism）。在这种情况下，整个模型会被复制到每个 GPU 上，并且每个 GPU 都会保存一份完整的优化器状态、梯度和参数。
Stage 1：在这个阶段，仅对优化器状态进行分割（Partitioning optimizer state），这意味着每个进程只保存优化器状态的一部分。这减少了每个 GPU 所需的内存，因为不再需要存储整个模型的优化器状态。
Stage 2：除了 Stage 1 的优化器状态分割外，Stage 2 还包括梯度的分割（Partitioning gradients）。这样做的结果是，不仅优化器状态被分割，梯度也被分割成更小的部分，进一步降低了内存需求。
Stage 3：这是最为复杂的阶段，它不仅包含了 Stage 1 和 Stage 2 的所有优化措施，还将模型参数本身进行了分割（Partitioning model parameters）。这意味着模型参数也会分布在多个 GPU 上，从而允许训练比单个 GPU 内存更大的模型。

此外，还有一个相关的概念叫做 ZeRO-offload，它可以将部分或全部的优化器状态、梯度和参数卸载到 CPU 内存上，以进一步降低 GPU 显存的使用量。不过，这样做通常会带来一定的计算性能损失，因为需要通过 PCI-E 总线传输数据。

3 vLLM

vLLM（Virtualized Language Learning Model）是一种用于自然语言处理（NLP）的模型架构或框架，旨在提高大规模语言模型（如GPT等）的性能和效率。论文：《Efficient Memory Management for Large Language Model Serving with PagedAttention》。

在模型推理时，GPU的内存分配如下：

内存分配
其中：

Parameters 保留权重等参数，是静态的，这部分无法优化；
KV Cache 是 Transformer 的 attention 机制引入的中间缓存；
Others 是临时激活函数使用，占用比例较小，优化空间不大。

从上面 GPU 的内存分配来看，KV Cache 是影响推理吞吐量的瓶颈，如果 KV Cache 管理不好，导致一次推理输出的数量太少，就会导致推理速度降低。

3.1 存在的问题

KV Cache 太大：LLM 服务需要为每个请求维护一个键值（KV）缓存，用于存储模型在生成文本时的上下文信息。随着请求数量的增加，KV缓存的大小迅速增长，占用大量 GPU 内存。
复杂的解码算法：LLM 服务通常提供多种解码算法供用户选择，如贪婪解码、采样解码和束搜索（beam search）。这些算法对内存管理的复杂性有不同的影响。
未知的输入和输出长度：LLM 服务的输入和输出长度是变化的，这要求内存管理系统能够适应不同长度的提示。随着请求的输出长度在解码过程中增长，所需的 KV 缓存内存也会增加，可能会耗尽用于新请求或现有的内存。

3.2 解决方法

为了解决这些挑战，vLLM 提出了一种新的注意力算法 PagedAttention，并构建了一个高效的内存管理系统：KV Cache Manager，通过分页技术来管理 KV Cache，从而提高内存的利用效率，减少内存浪费，并支持更复杂的解码算法。这种方法允许在非连续的物理内存中存储连续的键和值，使得内存管理更加灵活，能够更有效地处理 LLM 服务中的内存挑战。vLLM架构如下图所示：

vLLM架构

3.2.1 PagedAttention

PagedAttention 是一种受操作系统中虚拟内存和分页技术启发的注意力算法。它允许将连续的 K 和 V 向量存储在非连续的内存空间中。这一点与传统的注意力算法不同，后者通常要求 K 和 V 向量在内存中连续存储。

PagedAttention

3.2.2 KV Cache Manager

KV Cache Manager 是 vLLM 系统中的一个核心组件，负责以分页的方式高效管理 KV Cache。这一管理器的设计灵感来源于操作系统中的虚拟内存管理技术，特别是分页机制。

Manager

3.2.3 其他解码场景

其他解码场景的优化包括：Parallel sampling、Beam search和共享前缀。

3.3 结论

评估表明，与 FasterTransformer 和 Orca 等最先进的系统相比，vLLM 在相同的延迟水平下将流行 LLM 的吞吐量提高了2-4倍。

4 LightEval

LightEval 是由 Hugging Face 推出的一款轻量级AI评估工具，专门用于评估大型语言模型（LLMs），它支持多任务处理和复杂模型配置，并且能够在多种硬件上运行，包括 CPU、GPU 和 TPU 。

4.1 主要功能

多设备支持：LightEval 支持在多种设备上进行模型评估，这使得它能够适应不同的硬件环境，满足企业的需求。
易于使用：即使是技术水平不高的用户也能轻松上手，可以在多种流行基准上评估模型，甚至可以定义自己的自定义任务。
自定义评估：用户可以根据需求定制化评估过程，包括指定模型评估的配置，如权重、管道并行性等。
与 Hugging Face 生态系统集成：LightEval 可以与 Hugging Face Hub 等工具配合使用，方便模型的管理和共享。
支持复杂配置：通过配置文件加载模型，可以进行复杂的评估配置，例如使用适配器/增量权重或更复杂的配置选项。
流水线并行评估：支持在16位精度下评估大于约40B参数的模型，通过流水线并行技术将模型分片到多个 GPU 以适应 VRAM。

4.2 使用方法

安装 LightEval：首先需要克隆 LightEval 的 GitHub 仓库到本地，然后创建一个虚拟环境并激活之，最后安装 LightEval 及其依赖项。
配置评估环境：可以通过 accelerate config 命令来配置多 GPU 环境。
运行评估：使用 run_evals_accelerate.py 脚本在单个或多个 GPU 上评估模型。可以通过命令行参数指定模型和任务的配置。
指定任务和模型参数：通过 --tasks 参数指定要运行的任务，通过 --model_args 参数指定模型的路径或名称，使用 --override_batch_size 来覆盖默认的批处理大小，使用 --output_dir 指定输出目录。
自定义任务和指标：如果需要添加新的任务或指标，可以修改 tasks_table.jsonl 文件或创建新的 Python 文件来定义它们，确保新任务可以通过 LightEval 运行。
查看和分析结果：评估完成后，结果将保存在指定的输出目录中。可以查看生成的日志文件和结果文件来分析模型的性能。