DeepSpeed Universal Checkpointing 完整实战指南：跨并行拓扑无缝迁移模型与恢复训练

DeepSpeed Universal Checkpointing 完整实战指南：跨并行拓扑无缝迁移模型与恢复训练

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

Universal Checkpointing 是 DeepSpeed 提供的一套通用检查点（Checkpoint）机制，它把 ZeRO 训练中分散在模型、优化器与调度器各分片上的状态统一转换为与并行度无关的“Universal”格式，从而让同一份权重可以在不同的数据并行（DP）规模、张量并行（AutoTP）与专家并行（AutoEP）配置之间自由迁移、无缝续训。本文将以此教程为主体，结合仓库内转换脚本、加载链路与单元测试的源码级细节，带你走完"创建 ZeRO 检查点 → 转换 Universal 格式 → 跨拓扑恢复训练"的完整三步流程。

Universal Checkpointing 要解决的问题

在常规 DeepSpeed ZeRO 训练中，save_checkpoint落盘的模型状态、优化器状态与学习率调度器状态都和当时运行的并行拓扑强绑定：

• ZeRO 会把参数、梯度和优化器状态按数据并行 rank 切分；
• 张量并行（TP）下每个 rank 只保存自己那一份列/行切分；
• 专家并行（AutoEP）下路由专家权重又叠加了专家维度的切分。

这带来的直接后果是：换一个 GPU 数量、换一种并行策略、把 AutoTP 权重挪到纯 ZeRO 拓扑、或者在推理与训练之间切换，旧检查点往往无法直接加载。Universal Checkpointing 的核心思路正是抽象并消除这种拓扑耦合——它把分散在各 rank 分片中的状态重新聚合、重排成统一布局，使其“开箱即用”地支持从小模型到大模型、从单机到多卡多加速器、在不同并行拓扑之间自由穿梭。

从源码常量可见这一设计定位：checkpoint/constants.py 中定义了UNIVERSAL_CHECKPOINT_INFO、UNIVERSAL_CHECKPOINT_VERSION_KEY，当前版本号UNIVERSAL_CHECKPOINT_VERSION_VALUE = 0.4；转换与加载阶段都依赖这份随检查点保存的元数据来决定每个参数该按何种模式重组。

使用前置条件

开始前请确认：

• 已安装 DeepSpeed（pip install deepspeed）；
• 拥有一份基于 DeepSpeed 的分布式训练脚本；
• 训练时启用了 ZeRO（Zero Redundancy Optimizer）优化，即能产出 ZeRO 格式检查点。

三步核心工作流

无论面向 Dense（稠密）模型、AutoTP（自动张量并行）还是 AutoEP（自动专家并行）训练，Universal Checkpointing 的高层流程完全一致，共三步：

1. 保存一份常规 DeepSpeed ZeRO 检查点；
2. 用转换脚本把该检查点转换为 Universal 格式；
3. 在目标运行中开启checkpoint.load_universal，通过常规 API加载这份 Universal 检查点。

Step 1：保存常规 ZeRO 检查点

训练脚本中直接使用 DeepSpeed 常规检查点 API：

engine.save_checkpoint(save_dir, tag=tag)

这段调用对普通 Dense 训练、AutoTP 训练与 AutoEP 训练完全相同。不同之处在于落盘的附加信息：

• AutoTP 检查点会额外携带 Universal Checkpoint 元数据，用于描述张量并行参数的布局（行并行、列并行、复制参数、融合参数、子参数等），使后续转换器能够精确还原完整权重；
• AutoEP 检查点同样使用该 API，但其特有的专家布局约束详见下文“AutoEP 需求与限制”一节。

保存目录中的文件名遵循固定约定（见 checkpoint/constants.py）：模型态文件以mp_rank_/zero_pp_rank_开头、以_model_states.pt结尾，优化器态文件以_optim_states.pt结尾，MoE/AutoEP 专家文件则以layer_开头。

Step 2：将 ZeRO 检查点转换为 Universal 格式

得到 ZeRO 检查点后，使用 DeepSpeed 自带的转换脚本ds_to_universal.py：

python deepspeed/checkpoint/ds_to_universal.py \ --input_folder /path/to/ds_checkpoint \ --output_folder /path/to/universal_checkpoint

仓库根目录下对应源码为 deepspeed/checkpoint/ds_to_universal.py，该脚本会读取 ZeRO 检查点、重组权重并写入输出目录；运行时传入--help可查看全部选项。

命令行参数一览

从脚本parse_arguments()的定义（ds_to_universal.py#L62-L89）可确认以下参数及其默认值：

参数	默认值	含义
--input_folder	必填	输入 DeepSpeed 检查点目录
--output_folder	必填	输出 Universal 检查点目录
--num_extract_workers	4	并行抽取 ZeRO 分片的进程数
--num_merge_workers	2	并行合并 TP 切片的进程数（内存占用更高，建议远小于抽取进程数）
--keep_temp_folder	关闭	保留转换过程中的中间切片临时目录（默认会清理，调试用）
--no_strict	关闭	关闭转换结果的合法性校验（默认开启严格校验）
--inject_missing_state	关闭	若检查点缺失必需的元数据状态，则注入后再转换

转换器的内部工作流

根据main()的实现（ds_to_universal.py#L1017-L1200），转换分两条路径执行：

• ZeRO Stage 1/2 路径：先探测zero_stage（从优化器文件读取），用DeepSpeedCheckpoint解析输入的mp_rank_*分片；随后① 按(pp, tp, dp)三维网格并行抽取 ZeRO 参数与优化器片段到临时目录；② 依据UNIVERSAL_CHECKPOINT_INFO中记录的 TP 合并模式把各 TP 切片合并回完整参数；③ 若存在 AutoEP 元数据则额外合并专家状态；④ 保存公共优化器状态zero/optimizer_state.pt；⑤ 拷贝mp_rank_*模型文件并在必要时注入 AutoEP 元数据。全部完成后还会在输出目录的上级写入latest_universal标记文件（引擎加载时用它定位最新 tag）。
• ZeRO Stage 3 路径：检测到zero_pp_rank_*文件后，先从文件名解析出(tp_rank, dp_rank)网格并校验网格完整性（_build_zero3_rank_grid），再按 (tp, dp) 工作项抽取分片，合并时复用 TP-aware 合并逻辑还原数据并行与张量并行两个维度。

输出目录结构

转换完成后，输出目录中每个参数在zero/布局下拥有一个子目录：

<output_folder>/ ├── latest_universal ├── zero/ │ ├── optimizer_state.pt │ └── <参数名>/ │ ├── fp32.pt │ ├── exp_avg.pt │ ├── exp_avg_sq.pt │ └── step.pt └── mp_rank_*_model_states.pt ...

fp32.pt/exp_avg.pt/exp_avg_sq.pt即该参数的完整 fp32 主权重与 Adam 一阶、二阶动量（文件名对应 constants.py 的FP32_WEIGHT_KEY = "fp32"等常量），每个文件同时携带param、cat_dim等描述信息；优化器组级公共状态则汇总进zero/optimizer_state.pt。这也是加载端按名取值的唯一数据源。

合并参数时的布局分类

对 AutoTP 检查点，merge_tp_slices()（ds_to_universal.py#L289-L460）依赖保存侧的UNIVERSAL_CHECKPOINT_INFO对每个参数执行不同合并策略。这些策略正是本教程文档中所说"row-parallel、column-parallel、replicated、fused、sub-parameter"布局的底层实现：

元数据模式（constants 键）	合并行为
tp_replicated_parameter_patterns	各 TP 切片应完全一致，取第一个切片并断言其余相同
parameter_to_average_patterns	取所有 TP 切片的均值（如部分归一化层/Bias 参数）
parameter_with_row_parallelism_patterns	沿cat_dim=1拼接（行并行）
默认（其余列并行）	沿cat_dim=0拼接
parameter_with_2_sub_params_cat_dim_0	把每个切片按 dim 0 等分为两半，先分别拼接两半、再合并（如 Megatron 的 fused 参数）
parameter_with_sub_params+sub_param_shard_widths	带子参数物理宽度的精密切分（版本 0.4 起记录；0.4 之前的旧元数据仅能假设均匀切分，遇到不均匀子参数会直接拒绝）
vocabulary_parameter_patterns	拼接后按original_vocab_size裁掉 padding 并打上vocab_tensor标记

例如 MoE 场景中常见的mlp.h_to_4h这类"单个连续分配但按多个子参数使用"的参数，转换器会先按子参数重组、再对每个子参数横向拼接（见 constants.py#L73-L98 关于SUB_PARAM_SHAPE、SUB_PARAM_SHARD_WIDTHS的注释）。严格模式下，任何未被任何参数命中的元数据模式都会被断言捕获（Unused patterns ... while merging tp slices），用于发现保存侧与模型定义不一致的问题。

Step 3：加载 Universal 检查点恢复训练

目标运行（新拓扑下的训练）只需在 DeepSpeed 配置中加入：

{ "checkpoint": { "load_universal": true } }

然后沿用常规 DeepSpeed 检查点 API 加载：

engine.load_checkpoint("/path/to/universal_checkpoint", tag=tag)

配置项的解析链路清晰可见：默认值定义在 runtime/constants.py（load_universal=false，注释中还给出tag_validation、use_node_local_storage等同级配置），解析位于 runtime/config.py#L892-L900。引擎侧 runtime/engine.py 的load_universal_checkpoint()会把该开关传给加载逻辑，并由此决定：tag 文件读取latest_universal而非latest、以"通用零分片加载"路径替换常规load_module_state_dict，以及跳过与 MoE 常规加载相关的分支等（见 engine.py#L4399-L4607 一带的多个判断点）。

加载时注意：目标训练仍需配置与你期望恢复所用模型与拓扑匹配的 DeepSpeed 并行配置——Universal 检查点承担的是"状态重排"，而并行组本身由新一轮运行的配置建立。

load_checkpoint还支持三个常用控制开关（与文档中 AutoEP 部分提到的语义一致）：

• load_optimizer_states：是否连优化器状态一起加载；
• load_module_only：仅恢复模块权重（module-only），不触碰优化器；
• load_module_strict：模型加载是否严格匹配（对应strict语义）。

AutoTP：自动张量并行检查点的转换

AutoTP 检查点保存时已写入 Universal Checkpoint 元数据（UNIVERSAL_CHECKPOINT_INFO），因此转换器能精确还原张量并行参数布局。需要强调的是：转换后的 AutoTP 检查点可以脱离原 TP 拓扑运行——加载端会结合目标 rank 上每个参数携带的ds_autotp_universal_checkpoint_meta（常量DS_AUTOTP_UC_META，见 constants.py#L62-L64），依据partition_dim与logical_shape重新切出当前 TP rank 的局部视图，从而支持把同一份权重加载到与原 TP 度数不同的运行中。这部分参数裁剪/拼接（含融合注意力权重按 key/value 头边界切分、旧版本均匀切分假设的护栏断言等）在 checkpoint/universal_checkpoint.py 中有完整实现。

AutoEP：需求与限制

AutoEP 检查点虽然也以常规 DeepSpeed 检查点保存，但路由专家权重存在依赖 ZeRO 阶段的额外布局。

两种布局格式

• ZeRO Stage 1 / Stage 2：DeepSpeed 把路由专家权重（w1、w2、w3）单独写入每个专家的文件中，命名形如layer_<moe_layer_id>_expert_<global_expert_id>_mp_rank_<NN>_model_states.pt； 而 router、gate、共享专家以及其他非路由专家参数仍存放在常规mp_rank_*_model_states.pt文件中，走标准的 Universal Checkpointing 路径。
• ZeRO Stage 3：AutoEP 检查点是partition-native（分区原生）的——不再产生独立专家文件，专家权重以 ZeRO 分片形式存放在zero_pp_rank_*_model_states.pt及优化器分片文件中，并记录分片组元数据。

两种布局下，常规模型检查点都会把 AutoEP 元数据记录在ds_autoep_layers键中；更早的检查点则可能使用旧键名autoep_layers（二者在 constants.py#L100-L118 中定义，脚本读取时二者兼容）。

转换与加载行为

ZeRO Stage 1/2 与 ZeRO Stage 3 的 AutoEP 检查点均可转换为 Universal 格式。对 ZeRO Stage 3，ds_to_universal.py会从zero_pp_rank_*_model_states.pt中探测 partition-native 的 AutoEP 元数据，在专家副本组内把每个专家参数从各分片中汇聚出来，并写出与其他 Stage 相同的zero/参数布局（实现见_consolidate_zero3_autoep_expert_states，ds_to_universal.py#L719-L824，其中包含 EP rank 完整性、分区 rank 集合与形状一致性的多重校验）。

ZeRO Stage 3 AutoEP 还额外支持load_module_only=True的模块级加载与load_optimizer_states=False的无优化器状态加载，两种方式都能同时作用于分区分片与 Universal 格式。转换为 Universal 格式后，ZeRO Stage 3 AutoEP 允许在不同数据并行 world size、不同autoep_size或二者同时变化的情况下加载含优化器或仅权重/模块的检查点——前提是目标 AutoEP 拓扑对相同的模型参数名与专家数量有效。需要特别指出的是：分区原生的 AutoEP 检查点不支持用zero_to_fp32.py合并（该脚本会抛出NotImplementedError），这类检查点请改用ds_to_universal.py。

转换期间，ds_to_universal.py读取ds_autoep_layers（或旧键autoep_layers），汇聚每个 AutoEP 层的路由专家状态（Stage 1/2 来自逐专家文件，Stage 3 来自分区分片），并把完整专家张量写入形如zero/<expert_key_prefix>.w1/fp32.pt的路径。这些文件以is_expert_param与ep_num_experts打标——二者正是加载期执行 AutoEP 专家重分片（resharding）的触发信号（常量定义见 constants.py#L121-L125）。当存在匹配的专家优化器分片时，转换器还会把exp_avg.pt、exp_avg_sq.pt等优化器状态文件写到对应参数旁。

在 Universal 加载路径中，AutoEP 路由专家从zero/参数布局恢复，而非常规的layer_*_expert_*_model_states.pt文件。目标运行的 AutoEP 进程组提供加载侧的专家并行 rank 与 size；对每个带专家标记的张量，加载器先按ep_rank与ep_size切分保存的专家维度，再套用目标 ZeRO 分区组与 padding。

对目标模型的约束

目标模型仍需暴露名字匹配、形状兼容的 AutoEP 参数，例如：

<module_path>.experts.w1 <module_path>.experts.w2 <module_path>.experts.w3

需要明确边界：Universal Checkpointing 只会改变匹配张量的专家并行切分方式，它不会在不同模型族、不同模块路径或任意专家参数名之间做翻译。同时，目标 AutoEP 配置在加载前必须本身合法：autoep_size必须能整除目标流水线 stage size（world_size / pp_size），且能整除检测到的每个目标层专家数。

拓扑变化边界与数值约束

AutoEP 的拓扑变化被限定在两类：数据并行 world size 变化与匹配的 AutoEP 管理专家参数的autoep_size重分片。对检查点中每个 AutoEP 层，保存的ep_num_experts必须能被目标autoep_size整除。例如一个 8 专家的检查点可以加载到autoep_size为 1、2、4 或 8 的目标运行，但不能是 3。当autoep_size=1时专家张量不会被切分，但目标参数仍须具备兼容的完整专家形状。

ZeRO Stage 3 下，含优化器的加载会按目标运行时拓扑对路由专家参数及其 Adamfp32、exp_avg、exp_avg_sq状态重新切片；仅权重/模块加载则只从 universalfp32.pt文件对路由专家参数与常规参数重新切片，无需优化器状态。

AutoEP 失败场景清单

文档与本仓库实现共同给出了以下失败或受限情形，便于排查：

• 缺少元数据：ZeRO Stage 1/2 转换时，专家检查点文件若没有ds_autoep_layers或autoep_layers元数据，将抛出RuntimeError；
• 原生 MoE 混淆：既有的 DeepSpeed MoE / Megatron-DeepSpeed 专家检查点文件可能同名（layer_<moe_layer_id>_expert_<global_expert_id>_mp_rank_<NN>_model_states.pt），但它们使用原生deepspeed_moe专家参数名且不带 AutoEP 元数据；把这类检查点加载或转换成 AutoEP 需要单独的、模型特定的迁移步骤（转换脚本检测到这类文件时仅打印提示并跳过 AutoEP 合并）；
• 专家文件缺失：存在 AutoEP 元数据但预期的逐专家模型文件缺失时，转换抛出FileNotFoundError；
• 同一 (layer, expert) 多个mp_rank_*专家文件：抛出NotImplementedError；本路径不覆盖 AutoEP + AutoTP 组合拓扑变更；
• 优化器状态尽力合并：AutoEP 优化器状态合并是"尽力而为"——对通常的、包含匹配专家优化器分片的 ZeRO Stage 1/2 AutoEP 训练检查点会成功；若expp_rank_*_mp_rank_*_optim_states.pt文件或匹配状态条目缺失，转换器仍会写出模型参数的fp32.pt文件并跳过缺失的优化器状态文件；
• ZeRO Stage 3 兼容要求：Stage 3 通用转换需要 partition-native AutoEP 检查点格式，否则抛出NotImplementedError；冻结的专家参数也不能从该格式转换。

典型验证与测试参考

仓库的单元测试从端到端角度印证了"转换 + 跨拓扑加载"的核心卖点。例如 tests/unit/checkpoint/test_universal_checkpoint.py 提供了test_dp_world_size_2to2、test_dp_world_size_4to2、test_dp_world_size_2to4等用例——在 DP world size 2 与 4 之间双向变化并断言加载结果，测试中会以{"checkpoint": {"load_universal": True}}配置驱动加载路径（同文件第 225 行附近）；tests/unit/checkpoint/test_autotp_uc_checkpoint.py 与 tests/unit/v1/moe/test_autoep_checkpoint.py 则分别覆盖 AutoTP 元数据还原与 AutoEP 专家重分片场景，是理解各布局参数合并语义的绝佳阅读材料。

结语

DeepSpeed Universal Checkpointing 简化了模型状态的管理：通过"常规 ZeRO 保存 →ds_to_universal.py转换 →load_universal加载"三步，它让模型、优化器与调度器状态可以跨训练会话、跨数据并行规模、跨 AutoTP/AutoEP 拓扑自由转移，从而支撑起基于预训练模型的多配置实验与大规模分布式训练的平滑演进。针对 AutoEP 的分区原生（Stage 3）与逐专家文件（Stage 1/2）双布局支持，加上转换与加载两侧的严格元数据校验，使得专家模型也能在改变autoep_size与 DP world size 后安全续训。

若想进一步了解设计思路与更大规模集群上的实践细节，可继续阅读仓库内 DeepSpeed Universal Checkpointing 博客（并提供中文版与日文版）；对 Megatron-DeepSpeed 训练框架中的完整通用检查点示例与高级配置，可参考其官方examples_deepspeed/universal_checkpointing目录。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed