DeepSpeed 使用 LoRA 训练后文件结构详解

在大语言模型（LLM）的训练过程中，DeepSpeed 提供了强大的分布式训练能力，而 LoRA（Low-Rank Adaptation）通过参数高效微调技术显著减少了资源占用。完成训练后，DeepSpeed 会输出一系列文件，这些文件包括模型权重、训练状态、优化器状态和其他相关配置文件。理解这些文件的用途对于后续模型加载、微调和推理非常重要。

本文将基于训练后生成的文件结构，逐一解析这些文件的作用及使用场景。

下图是我训练之后得到的内容，基于此阅读本篇博文。

在这里插入图片描述

一、文件结构概览

训练完成后生成的文件大致可以分为以下几类：

模型权重文件：
- bf16_zero_pp_rank_*_mp_rank_*.pt
- zero_pp_rank_*_mp_rank_*.pt
训练状态文件：
- random_states_*.pkl
- scheduler.bin
- latest
LoRA 微调适配器文件：
- adapter_config.json
- adapter_model.safetensors
分词器及其他配置文件：
- config.json
- special_tokens_map.json
- tokenizer_config.json
- tokenizer.json
- README.md
转换脚本：
- zero_to_fp32.py

二、详细文件解析

1. 模型权重文件

**`bf16_zero_pp_rank__mp_rank_.pt`**

这些文件是 优化器状态（optimizer states），用于保存训练过程中优化器的中间状态，通常在恢复中断训练或进行多阶段微调时会用到。文件名的结构说明了分布式训练的配置：

bf16 表示训练中使用了半精度（bfloat16）计算。
zero_pp_rank_* 是 ZeRO 优化器的分片编号，pp 代表模型的分布式数据并行策略。
mp_rank_* 是模型并行分片编号。

**`zero_pp_rank__mp_rank_.pt`**

这些文件是 模型状态文件（model states），保存了经过分布式切分的模型权重。它们是 DeepSpeed ZeRO 优化器在分布式训练下生成的模型权重分片。

需要注意的是，这些分片权重在推理时不能直接使用，通常需要通过脚本将它们合并成一个完整的模型权重文件。

2. 训练状态文件

**`random_states_*.pkl`**

这些文件记录了训练时的随机数状态，包括生成器的种子等信息，用于保证训练的可重复性。每个文件对应一个分布式训练节点（rank）。

`scheduler.bin`

这是 学习率调度器（scheduler） 的状态文件，记录了训练过程中学习率的变化曲线。如果你需要恢复训练，这个文件可以帮助调度器从中断点继续。

`latest`

这是一个指针文件，通常用来标记当前最新的 checkpoint，便于训练恢复时直接加载最近的模型状态。

3. LoRA 微调适配器文件

`adapter_config.json`

这是 LoRA 微调的配置文件，记录了 LoRA 的超参数和目标模块，例如：

r：低秩矩阵的秩。
alpha：缩放因子。
target_modules：指定 LoRA 适配的模型模块（如 query 和 value）。
lora_dropout：dropout 的概率。

`adapter_model.safetensors`

这是 LoRA 微调后的适配器权重，使用 safetensors 格式存储以提高安全性和加载速度。它只包含微调过程中训练出的 LoRA 模块权重，相比完整模型权重文件更加轻量。

适配器权重的加载通常需要配合基础模型和 adapter_config.json。

4. 分词器及其他配置文件

`config.json`

基础模型的架构配置文件，包含模型的隐藏层大小、注意力头数、层数等信息。这是模型加载时必不可少的文件。

`special_tokens_map.json`

记录了特殊符号（如 [CLS], [SEP]）的映射关系，用于确保模型正确处理输入文本中的特殊符号。

`tokenizer_config.json`

分词器的配置文件，记录了分词器的超参数，例如是否区分大小写、模型最大序列长度等。

`tokenizer.json`

分词器的完整定义文件，通常包含词汇表、分词规则等信息。

`README.md`

文件的说明文档，通常包含训练设置或模型相关的简要信息。

5. 转换脚本

`zero_to_fp32.py`

这是一个非常重要的脚本，用于将分布式训练生成的 ZeRO 优化器分片权重（zero_pp_rank_* 文件）转换成单个完整的模型权重文件（pytorch_model.bin），便于推理或进一步的模型部署。

使用方式示例：

python zero_to_fp32.py \
  --input_dir /path/to/checkpoint \
  --output_file /path/to/pytorch_model.bin

三、文件间的关系

从以上解析可以看出，DeepSpeed 使用 LoRA 训练后的文件大致可以分为三类：

模型相关文件：包括分布式权重、LoRA 模块权重和配置文件。
训练状态文件：用于恢复训练进度，包括优化器状态、随机数状态和学习率调度器状态。
分词器和辅助文件：确保分词器配置与模型训练过程一致，便于模型部署。

这些文件共同构成了完整的训练结果，既可以用于后续训练，也可以通过转换脚本生成推理所需的完整权重文件。

四、总结

通过对 DeepSpeed 使用 LoRA 训练后生成文件的解析，我们可以清晰地了解每个文件的作用以及它们在模型微调和部署中的重要性。在实际使用中，根据需求选择性地加载权重文件、恢复训练状态或进行分词器配置，可以显著提高工作效率。

pp_rank在变而mp_rank不变的原因分析

bf16_zero_pp_rank_0_mp_rank_00_optim_states，bf16_zero_pp_rank_1_mp_rank_00_optim_states等文件，为什么pp后面的rank在变，mp rank却都是00没变啊

这是一个非常好的问题！我们来分析一下 DeepSpeed 的分布式训练中 pp 和 mp 的含义，以及为什么在文件名中 pp 后面的 rank 会变化，而 mp_rank 始终是 00。

1. `pp` 和 `mp` 的含义

在 DeepSpeed 的分布式训练中，pp 和 mp 分别代表：

pp (Pipeline Parallelism)：表示流水线并行的 rank（分片编号）。
- 流水线并行将模型的不同层切分到不同的设备上（GPU/节点）。每个 rank 对应一部分模型的层，例如 Rank 0 处理前几层，Rank 1 处理中间层，Rank 2 处理后几层。
- 如果模型很大（例如数十亿参数），流水线并行可以帮助减小单个设备的内存占用。
mp (Model Parallelism)：表示模型并行的 rank。
- 模型并行将单个模型的权重进一步拆分到多个设备中。例如，某一层的权重可以拆分为多个部分，每个设备只保存部分权重。
- mp_rank_00 表示模型并行的 rank 号为 0，即没有额外的模型切分。

2. 为什么 `pp_rank` 在变化，`mp_rank` 始终是 `00`？

文件名中 pp_rank 是流水线并行的 rank 编号，mp_rank 是模型并行的 rank 编号。
以下是两者变化的关键原因：

2.1 关于 `pp_rank` 的变化

在流水线并行中，每个 GPU 或设备负责模型的一部分层（例如层 0-10 分配到 pp_rank_0，层 11-20 分配到 pp_rank_1）。
每个 pp_rank 对应一段模型的计算任务，因此它会随着分布式训练的设备编号增加而变化。

在你的文件中，pp_rank_0、pp_rank_1、pp_rank_2 和 pp_rank_3 说明模型被分割为了 4 段流水线，每段由不同的 GPU 或设备负责。

2.2 关于 `mp_rank` 没有变化

你的训练设置中，模型并行的配置为 mp=1，也就是说，没有进一步对每个层的权重做切分。
由于模型并行的 rank 总是 00，说明每个层的权重都完整地存储在一个设备上，权重没有再被分割到其他设备。

这与模型的大小和训练资源有关：

如果模型体积足够小，可以直接在每个 GPU 内存中保存整个层的权重，无需模型并行。
如果模型非常大（例如 70B 参数），可能需要启用模型并行（mp>1），这时会看到多个 mp_rank_* 文件。

3. 总结

你看到的现象说明训练使用了 流水线并行（Pipeline Parallelism，pp>1），但没有启用 模型并行（Model Parallelism，mp=1）：

pp_rank_* 的变化表示模型的不同部分被分配到不同的设备进行流水线计算。
mp_rank_00 始终不变，说明模型的每一部分权重没有进一步切分。

4. 深入思考

如果你有更多 GPU 或更大的模型，可以调整 mp 配置。例如设置 mp=2，则会看到文件名中出现 mp_rank_00 和 mp_rank_01，说明模型并行在每层权重内部做了进一步切分。

在配置文件中，pipeline_parallel_size 和 tensor_parallel_size 是控制这两种并行方式的关键参数：

{
  "pipeline_parallel_size": 4,  // 流水线并行大小，对应 pp_rank
  "tensor_parallel_size": 1    // 模型并行大小，对应 mp_rank
}

后记

2024年12月28日22点52分于上海，在GPT4o大模型辅助下完成。