FSDP是PyTorch的分布式训练策略，训练大模型时常常需要保存断点以便中断后续训。这篇讲清楚如何完整保存模型分片参数、优化器状态、学习率调度器、随机数种子、Dataloader进度，以及恢复时如何正确重建分片避免精度损失。适合做大模型分布式训练、长周期训练任务、容错训练系统的算法工程师和平台开发人员。

岗位关联

需要学习PyTorch FSDP save/load 状态恢复：模型/优化器/scheduler/随机数状态，这是你的核心专业课，数据治理是数据工程的延伸。本篇涵盖数据清洗/ETL/数据仓库/特征工程，帮助你构建数据管道与洞察。

目录
  1. 可视化演示
  2. FSDP断点续训的核心概念与完整状态组成
  3. 模型分片参数的分片感知保存与加载机制
  4. 优化器状态与学习率调度器的保存恢复
  5. 随机数状态与数据加载进度的持久化
  6. 分片重建与精度无损恢复的工程要点
  7. 容错训练系统的断点续训设计模式
  8. AI 应用场景
  9. Python 演示
  10. 常见问答
  11. 常见误区
  12. 参考来源

可视化演示

展示FSDP训练中模型分片/优化器/调度器/RNG/DataLoader的完整状态保存、恢复校验和分片重建流程，动态演示组件序列化、all-gather收集和unshard参数重建。

本节配有交互动效（进入视口后自动加载，约 3KB）。

保存状态 模型 优化器 FSDP

FSDP断点续训的核心概念与完整状态组成

FSDP（Fully Sharded Data Parallel）在训练过程中将模型参数、梯度以及优化器状态按卡（rank）进行切分存储，每张 GPU 只持有全局状态的一个分片。这种设计大幅降低了单卡显存占用，但也使得断点续训比传统 DataParallel 更加复杂——必须将所有分片的状态完整收集或分别保存，才能在故障恢复时重建出训练中间态。

断点续训需要保存五类核心状态：（1）模型参数——网络权重与偏置；（2）优化器状态——如 Adam 的一阶动量和二阶方差；（3）学习率调度器状态——当前 step 数与预热进度；（4）随机数生成器状态——各 rank 的 CPU 与 CUDA 随机种子；（5）数据加载进度——当前 epoch 与迭代索引。只有五类状态全部持久化，恢复后才能精确复现原始训练轨迹。

如果仅保存模型权重而忽略优化器状态，恢复后优化器将处于冷启动状态：Adam 的动量和方差从零开始累积，导致训练初期梯度方向严重失准，损失曲线出现剧烈尖峰甚至发散。实验表明，在大规模 Transformer 训练中，优化器状态丢失可使恢复后的 loss 需要数百步才能重新收敛到中断前水平。

五类状态缺一不可。模型参数决定前向计算结果，优化器状态决定参数更新方向，调度器状态决定学习率大小，随机数状态决定 dropout 与数据增强的精确行为，数据进度决定训练样本的读取位置。任何一个维度的缺失都会导致恢复后的训练轨迹与原始轨迹产生偏差，在严格复现实验或长时间训练场景下不可接受。

θt+1=θt−ηt⋅mtvt+ϵ\theta_{t+1} = \theta_t - \eta_t \cdot \frac{m_t}{\sqrt{v_t} + \epsilon}

模型分片参数的分片感知保存与加载机制

在 FSDP 模式下，如果直接使用普通的 torch.save 保存模型 state_dict，只会保存当前 rank 所持有的参数分片，而非完整的全局参数。这是因为 FSDP 在运行时将参数切分到各卡，每张卡只能看到自己负责的那部分参数。这样保存得到的文件是不完整的，无法用于恢复或迁移。

分片感知保存（shard-aware saving）的核心思路是让每个 rank 各自落盘自己持有的参数分片。PyTorch 提供了 DistributedCheckpoint（DCP）API，调用时每个 rank 将自身的 local_state_dict 写入共享存储的对应分片文件中。恢复时每个 rank 读取自己的分片，即可精确重建本地参数。这种方式内存开销最小，因为不需要在保存前聚合全量参数。

聚合保存（full-state saving）则是先将所有 rank 的参数分片通过 all_gather 收集到 rank 0（或所有卡），再保存完整的 state_dict。这种方式生成的检查点是自包含的，便于跨集群迁移和不同分片拓扑间的加载，但内存峰值极高——rank 0 需要容纳整个模型的全部参数，对于数十亿参数的模型可能需要数百 GB 内存。

分片保存与聚合保存各有取舍。分片保存节省内存但要求恢复时的分片数量与拓扑与保存时一致；聚合保存灵活但内存开销大。在实际工程中，通常采用分片保存作为默认策略，仅在需要跨集群迁移或改变并行度时才使用聚合保存。PyTorch 2.x 的 DCP 还支持 resharding 功能，允许在加载时自动适配不同的分片拓扑。

优化器状态与学习率调度器的保存恢复

Adam 优化器在训练过程中为每个参数维护两个内部状态：一阶动量 m_t（梯度的指数移动平均）和二阶方差 v_t（梯度平方的指数移动平均）。这两个状态是历史梯度的累积结果，一旦丢失就等同于冷启动——优化器需要重新积累数百甚至数千步的梯度统计信息，在此期间参数更新方向严重失准。

优化器状态丢失的直接表现是恢复后的训练初期出现损失尖峰（loss spike）。这是因为冷启动的 Adam 缺乏对梯度分布的准确估计，导致学习率自适应机制失效。在大模型训练中，这种尖峰可能触发梯度爆炸，进而导致训练发散。因此，保存优化器状态是断点续训中最重要的环节之一。

学习率调度器（LR Scheduler）记录了当前训练步数、预热（warmup）进度以及学习率衰减阶段等信息。如果不同步恢复调度器状态，恢复后的学习率可能跳变到一个错误的值——例如在 warmup 阶段突然使用最终学习率，或在余弦衰减末期使用过大的学习率，都会严重破坏训练稳定性。

正确的恢复顺序应当是：先构建模型并用 FSDP 包装，再构建优化器（此时优化器状态被初始化为零），然后加载检查点中的模型参数和优化器状态，最后恢复调度器状态并设置到正确的 step。这个顺序至关重要——如果在加载优化器状态之前就执行了训练步骤，优化器的内部计数器就会偏移，导致后续所有学习率计算都产生偏差。

mt=β1mt−1+(1−β1)gt,vt=β2vt−1+(1−β2)gt2m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t, \quad v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2

随机数状态与数据加载进度的持久化

在分布式训练中，每个 rank 拥有独立的随机数生成器（RNG），包括 CPU 端的 Python random 和 torch 的 CPU 生成器，以及每个 CUDA 设备的生成器。这些 RNG 控制着 dropout 的掩码生成、参数初始化、数据增强的随机变换等关键操作。只有全部保存所有 rank 的 RNG 状态，恢复后这些随机操作才能与中断前完全一致，保证训练行为的精确复现。

Dropout 是 Transformer 模型中常用的正则化技术，其本质是根据随机掩码按比例丢弃神经元。如果恢复后 RNG 状态不连续，dropout 掩码序列就会与原始训练不一致，导致梯度信号产生突变。类似地，数据增强（如随机裁剪、颜色抖动）也依赖 RNG，状态不连续会改变增强模式，间接影响模型的收敛轨迹。

数据加载进度如果不持久化，恢复后 DataLoader 将从头开始遍历数据集，导致大量数据被重复训练，而后续数据被跳过。这不仅浪费计算资源，还会破坏数据采样的统计有效性——模型在某些样本上过度拟合，而在另一些样本上训练不足。对于按学习率调度严格安排的数据集（如课程学习），数据进度丢失的影响更加严重。

解决方案是在检查点中记录当前 epoch 编号和迭代索引（iteration index），恢复时将 DataLoader 的 sampler 快进到对应位置。对于 PyTorch 的 DataLoader，可以通过设置 generator 的随机种子和手动跳过前 N 个 batch 来实现。对于更复杂的采样策略（如分布式采样器），还需要保存采样器的内部状态以确保各 rank 的数据分配与中断前一致。

分片重建与精度无损恢复的工程要点

恢复 FSDP 训练时，必须先将模型用 FSDP 包装完成分片，然后再加载分片状态。如果顺序颠倒——先加载完整状态再用 FSDP 包装——FSDP 的初始化过程会覆盖或重新切分已加载的参数，导致状态丢失或错位。正确的流程是：构建原始模型、用 FSDP 包装、构建优化器、加载检查点。

分片拓扑必须与保存时保持一致。所谓分片拓扑包括 rank 数量、每张卡负责的参数分片范围以及通信组（process group）的配置。如果恢复时的拓扑与保存时不同，各 rank 读取到的参数分片可能对应错误的参数区间，出现'张冠李戴'的严重问题——某些参数被重复加载，而另一些参数完全丢失。PyTorch DCP 的 resharding 功能可以部分缓解这一限制。

在混合精度训练（AMP）中，模型同时维护低精度（FP16/BF16）的计算副本和高精度（FP32）的主权重。优化器状态也是基于高精度主权重维护的。检查点必须保存高精度主权重，而非低精度的计算副本。如果只保存了低精度副本，恢复后不仅精度损失不可逆，优化器状态与主权重之间的对应关系也会错乱。

保存与恢复之间不能修改模型结构。任何层的增删、参数形状的改变、甚至 FSDP wrapping 策略的调整，都会导致 state_dict 的 key 或 tensor shape 与检查点不匹配，进而引发加载失败。如果确实需要修改模型结构，必须使用聚合保存获取完整 state_dict，然后手动进行状态映射和转换。

容错训练系统的断点续训设计模式

定期异步保存检查点是容错训练的基础设计模式。同步保存会阻塞训练主循环——在保存期间所有 GPU 停止计算，等待 I/O 完成。对于大模型，一次检查点保存可能需要数十 GB 的磁盘写入，耗时数十秒甚至数分钟。异步保存将序列化与写入操作放到独立线程或进程中，训练主循环可以继续执行，仅在'快照'瞬间有短暂同步开销。

保留多份历史检查点是应对检查点静默损坏（Silent Data Corruption, SDC）的关键策略。SDC 是指存储介质或传输通道在没有明显错误的情况下发生比特翻转，导致检查点文件内容损坏但校验和可能仍然通过。如果只保留最新一份检查点且该检查点已损坏，训练将被迫从更早的随机初始化重新开始。保留最近 3-5 份检查点并提供回滚机制，可以显著降低 SDC 带来的风险。

心跳检测与自动重启机制实现故障的自动感知与恢复。训练系统通过定期心跳监控各 rank 的健康状态，一旦检测到某个 rank 无响应（GPU 掉卡、节点宕机、网络分区），自动触发故障恢复流程：终止当前训练、从最近的检查点重新初始化所有 rank、恢复完整状态并继续训练。这要求编排系统（如 Kubernetes）具备自动重新调度 Pod 的能力。

分层存储策略将新检查点放在高速存储（如 NVMe SSD 或分布式文件系统的热数据层）以保证快速恢复，将旧检查点定期归档到对象存储（如 S3 或 OSS）以降低成本。这种策略在成本和恢复速度之间取得平衡：最近几次的检查点可以快速加载，而更早的检查点虽然加载较慢但作为兜底备份仍然存在。部分系统还支持增量检查点，只保存与上一份检查点的差异部分，进一步减少存储和 I/O 开销。

AI 应用场景

千亿参数大模型预训练的故障自动恢复

解决的问题：在一个使用 128 张 A100 GPU 训练千亿参数大语言模型的集群中，训练周期长达数月。在此期间，GPU 硬件故障、网络抖动、存储异常等问题不可避免。系统采用 FSDP 进行模型分片训练，配合分布式检查点（DCP）实现异步保存。每隔 500 步自动触发一次检查点保存，保留最近 3 份。当监控系统检测到某个节点 GPU ECC 错误率超过阈值时，自动触发故障切换：终止当前训练任务，在备用节点上重新调度 Pod，从最近的检查点恢复模型参数、优化器状态、调度器状态、RNG 状态和数据进度，整个恢复过程在 10 分钟内完成，训练损失曲线几乎无感知。

具体案例：在一个使用 128 张 A100 GPU 训练千亿参数大语言模型的集群中，训练周期长达数月。在此期间，GPU 硬件故障、网络抖动、存储异常等问题不可避免。系统采用 FSDP 进行模型分片训练，配合分布式检查点（DCP）实现异步保存。每隔 500 步自动触发一次检查点保存，保留最近 3 份。当监控系统检测到某个节点 GPU ECC 错误率超过阈值时，自动触发故障切换：终止当前训练任务，在备用节点上重新调度 Pod，从最近的检查点恢复模型参数、优化器状态、调度器状态、RNG 状态和数据进度，整个恢复过程在 10 分钟内完成，训练损失曲线几乎无感知。

跨数据中心的模型微调容灾迁移

解决的问题：某企业在 A 数据中心使用 FSDP 对预训练大模型进行领域微调，训练进行到一半时因机房电力维护需要整体迁移到 B 数据中心。由于采用了分片感知的检查点保存策略，每个 rank 的参数分片已持久化到共享对象存储。迁移时只需将检查点文件同步到新集群，在新集群上以相同或不同的 GPU 拓扑重新部署训练任务。利用 DCP 的 resharding 功能，即使新集群的 GPU 数量与原来不同（如从 64 卡变为 96 卡），也能自动重新分配参数分片并正确加载，确保微调过程不中断、模型精度不损失。

具体案例：某企业在 A 数据中心使用 FSDP 对预训练大模型进行领域微调，训练进行到一半时因机房电力维护需要整体迁移到 B 数据中心。由于采用了分片感知的检查点保存策略，每个 rank 的参数分片已持久化到共享对象存储。迁移时只需将检查点文件同步到新集群，在新集群上以相同或不同的 GPU 拓扑重新部署训练任务。利用 DCP 的 resharding 功能，即使新集群的 GPU 数量与原来不同（如从 64 卡变为 96 卡），也能自动重新分配参数分片并正确加载，确保微调过程不中断、模型精度不损失。

Python 演示

python75 行
import osimport torchimport torch.nn as nnfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDPfrom torch.distributed.fsdp import StateDictType, FullStateDictConfigimport torch.distributed.checkpoint as DCPfrom torch.distributed.checkpoint import FileSystemWriter, FileSystemReader def save_fsdp_checkpoint(model, optimizer, scheduler, rng_states, data_progress, step, save_dir):    """保存 FSDP 分布式检查点（分片感知方式）"""    os.makedirs(save_dir, exist_ok=True)    ckpt_path = os.path.join(save_dir, f"checkpoint_{step}")     # 使用 DistributedCheckpoint 分片保存    state = {        "model": model.state_dict(),        "optimizer": optimizer.state_dict(),        "scheduler": scheduler.state_dict(),        "rng_states": rng_states,        "data_progress": data_progress,        "step": step,    }    DCP.save(state, storage_writer=FileSystemWriter(ckpt_path))    print(f"[Rank {torch.distributed.get_rank()}] Checkpoint saved to {ckpt_path}") def load_fsdp_checkpoint(model, optimizer, scheduler, save_dir, step):    """加载 FSDP 分布式检查点并恢复完整训练状态"""    ckpt_path = os.path.join(save_dir, f"checkpoint_{step}")     state = {        "model": model.state_dict(),        "optimizer": optimizer.state_dict(),        "scheduler": scheduler.state_dict(),        "rng_states": None,        "data_progress": None,        "step": 0,    }    DCP.load(state, storage_writer=FileSystemReader(ckpt_path))     model.load_state_dict(state["model"])    optimizer.load_state_dict(state["optimizer"])    scheduler.load_state_dict(state["scheduler"])    return state["rng_states"], state["data_progress"], state["step"] def capture_rng_states():    """捕获所有 RNG 状态"""    return {        "python_rng": __import__('random').getstate(),        "torch_cpu_rng": torch.random.get_rng_state(),        "torch_cuda_rng": torch.cuda.get_rng_state_all(),    } def restore_rng_states(rng_states):    """恢复所有 RNG 状态"""    import random    random.setstate(rng_states["python_rng"])    torch.random.set_rng_state(rng_states["torch_cpu_rng"])    torch.cuda.set_rng_state_all(rng_states["torch_cuda_rng"]) # ---- 使用示例 ----# 1. 构建模型并用 FSDP 包装# model = MyModel()# model = FSDP(model)# 2. 构建优化器和调度器# optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10000)# 3. 训练循环中定期保存# for step in range(max_steps):#     train_step(model, optimizer, scheduler)#     if step % 500 == 0:#         rng = capture_rng_states()#         save_fsdp_checkpoint(model, optimizer, scheduler, rng, data_progress, step, "./ckpts")# 4. 故障恢复# rng, data_progress, resumed_step = load_fsdp_checkpoint(model, optimizer, scheduler, "./ckpts", last_step)# restore_rng_states(rng)
目的：演示 FSDP 分布式检查点的完整保存与加载流程，涵盖模型参数、优化器状态、调度器状态、随机数状态和数据进度的持久化与恢复。
  1. import torch.distributed.checkpoint as DCP导入 PyTorch 分布式检查点模块，提供分片感知的保存/加载 API
  2. "model": model.state_dict(),将模型参数加入 state 字典，与优化器、调度器、RNG 和数据进度一起组装为统一的检查点状态
  3. DCP.save(state, storage_writer=FileSystemWriter(ckpt_path))使用 DCP 分片保存，每个 rank 自动写入自己的参数分片到共享存储
  4. DCP.load(state, storage_writer=FileSystemReader(ckpt_path))使用 DCP 分片加载，每个 rank 读取对应的参数分片并重建本地状态
  5. model.load_state_dict(state["model"])将加载的模型状态应用到 FSDP 包装后的模型中
  6. optimizer.load_state_dict(state["optimizer"])恢复优化器的动量和方差状态，避免冷启动导致的损失尖峰
  7. scheduler.load_state_dict(state["scheduler"])恢复调度器的 step 计数和预热进度，确保学习率连续
  8. def capture_rng_states():捕获 Python random、torch CPU 和所有 CUDA 设备的随机数状态
  9. def restore_rng_states(rng_states):恢复所有 RNG 状态，确保 dropout 和数据增强的精确复现
  10. if step % 500 == 0:每 500 步触发一次检查点保存，平衡保存频率与 I/O 开销

常见问答

FSDP 断点续训需要保存哪五类状态？缺少任何一类会导致什么后果？
五类状态分别是：模型参数、优化器状态、学习率调度器状态、随机数生成器状态、数据加载进度。缺少模型参数则无法恢复前向计算；缺少优化器状态（如 Adam 动量）会导致冷启动和损失尖峰；缺少调度器状态会导致学习率跳变；缺少随机数状态会导致 dropout 和数据增强行为不连续；缺少数据进度会导致数据重复训练或跳过。
为什么不能直接用 torch.save 保存 FSDP 模型的 state_dict？
因为 FSDP 在运行时将参数切分到各个 rank，每个 rank 只持有全局参数的一个分片。直接使用 torch.save 只会保存当前 rank 的分片，得到的是不完整的参数。应使用分片感知保存（如 DistributedCheckpoint），让每个 rank 各自保存自己的分片，或使用 all_gather 聚合后在 rank 0 保存完整状态。
FSDP 检查点加载时为什么要先包装 FSDP 再加载状态？
因为 FSDP 包装过程会重新切分和管理参数。如果先加载状态再包装 FSDP，FSDP 的初始化会覆盖或重新分配已加载的参数，导致状态丢失或错位。正确顺序是：构建模型 → FSDP 包装 → 构建优化器 → 加载检查点状态。
混合精度训练中，检查点应该保存 FP16 计算副本还是 FP32 主权重？为什么？
应该保存 FP32 高精度主权重。因为优化器状态是基于 FP32 主权重维护的，如果只保存 FP16 副本，恢复后精度损失不可逆，且优化器状态与主权重之间的对应关系会错乱。FP32 主权重是参数更新的基准，FP16 副本仅用于前向/反向计算加速。
异步检查点保存相比同步保存有什么优势？如何实现？
异步保存不会阻塞训练主循环，训练可以在检查点写入磁盘的同时继续执行，大幅减少 GPU 空闲时间。实现方式是将检查点的序列化和写入操作放到独立线程或进程中，训练主循环仅在'快照'瞬间短暂同步以获取状态副本，随后立即恢复训练。

常见误区

误区只要保存了模型权重就能完美恢复训练，优化器状态不重要。

正解优化器状态（如 Adam 的动量和方差）是历史梯度的累积统计，丢失后等同于冷启动，会导致恢复初期梯度方向严重失准、损失尖峰甚至训练发散。必须同时保存优化器状态才能实现平滑恢复。

为何常见：初学者往往只关注模型参数而忽视优化器内部状态，因为在小型实验中优化器冷启动的影响不明显，但在大规模训练中影响极为严重。

误区FSDP 模型可以直接用 torch.save 保存完整模型参数。

正解FSDP 将参数分片存储在各 rank 上，普通 torch.save 只能保存当前 rank 的分片。必须使用分片感知保存（如 DistributedCheckpoint）或先 all_gather 聚合再保存。

为何常见：torch.save 是单机训练中最常用的保存方式，开发者容易习惯性地在分布式场景下直接复用，但 FSDP 的分片机制使得这种方式无法获得完整参数。

误区恢复训练时只要加载模型和优化器状态就够了，随机数状态影响可以忽略。

正解随机数状态控制 dropout 掩码、数据增强等随机操作。不恢复 RNG 状态会导致这些操作的行为与中断前不一致，虽然短期影响可能不大，但在严格复现实验或长时间训练中会导致训练轨迹偏差累积。

为何常见：随机数状态的影响是隐式的、渐进的，不像模型参数丢失那样立即可见，因此容易被忽视。

误区恢复 FSDP 训练时可以先加载 state_dict 再用 FSDP 包装模型。

正解FSDP 包装过程会重新切分和管理参数，先加载再包装会导致已加载的状态被覆盖或错位。正确顺序是先 FSDP 包装、再加载状态。

为何常见：单机训练中先加载权重再包装模型是常见做法，开发者容易将这个习惯带入 FSDP 场景，但 FSDP 的参数管理机制与单机完全不同。

参考来源