FSDP 多卡训练时，参数不是完整保存在每个 rank 上，而是按层或按参数被切分到各个 shard。加载预训练权重或恢复 checkpoint 后，直接打印某个 rank 上的参数只能看到本地分片，无法确认所有 shard 是否都加载正确。这里给出一种不依赖外部工具的校验方式：以保存时的完整状态为基准，在加载后逐 rank 对比本地分片参数与完整参数对应切片的值。
校验 FSDP 每个 shard 的参数加载，核心动作是“拿到完整基准”和“拿到本地分片”，然后按维度切片比对。建议在加载后、训练前做一次全参数一致性检查；严格模式只能保证键匹配，不保证值正确。对超大规模模型，可以先比对形状、dtype、非零值比例或抽样张量，再逐步放大。
先确认 checkpoint 保存方式
FSDP 的 state_dict 有两种表现形态：一种是 full state dict（完整参数，一个 rank 上就能看到全部数值），另一种是 sharded state dict（每个 rank 只保存自己的分片）。如果之前保存时用的是 FullStateDictConfig，那加载时可以直接拿到完整参数；如果保存的是原生 sharded state dict，就需要先把完整参数从各 rank 收集起来再做比对。建议先检查保存时的 state_dict_type，再决定校验流程。
校验步骤：从完整参数到分片比对
下面以 PyTorch FSDP 为例，给出一个可复制的校验骨架。整个过程在每个 rank 上运行，最后归约比对结果。
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import FullStateDictConfig, StateDictType
# 假设 model 已经是 FSDP 包装后的模型
torch.distributed.init_process_group(backend="nccl")
rank = dist.get_rank()
# 加载前先设置 full state dict 模式
full_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=False)
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_config):
full_sd = model.state_dict() # 每个 rank 都会获得完整 state_dict（CPU）
# 载入预训练 checkpoint，此处为示例
torch.load("pretrained.pt") # 实际按你的方案加载
# 用 load_state_dict 完成加载，注意 strict 参数按需处理
# 获取加载后的本地分片参数
sharded_sd = model.state_dict() # 默认 sharded state dict，每个 rank 只有本地切片
# 对每个参数名做比对
errors = []
for name, full_tensor in full_sd.items():
if name not in sharded_sd:
errors.append(f"{name} 缺失")
continue
local_tensor = sharded_sd[name]
# full_tensor 是完整参数，local_tensor 是当前 rank 的分片
# 关键：需要按 FSDP 的分片方式切出当前 rank 的对应部分
# 这里以均匀分片为例，实际应依据 FSDP 的 flatten 索引规则
num_shards = dist.get_world_size()
shard_size = full_tensor.numel() // num_shards
start = rank * shard_size
end = (rank + 1) * shard_size if rank != num_shards - 1 else full_tensor.numel()
full_shard = full_tensor.flatten()[start:end].reshape(local_tensor.shape)
if not torch.allclose(full_shard, local_tensor.detach().cpu(), atol=1e-6):
errors.append(f"{name} 数值不一致")
# 汇总到 rank0 并打印
error_tensor = torch.tensor(len(errors), device=torch.device("cuda"))
dist.all_reduce(error_tensor, op=dist.ReduceOp.SUM)
if rank == 0:
print(f"共 {error_tensor.item()} 处错误")
for err in errors[:10]:
print(err)上面例子里的均匀分片只是示意，FSDP 的分片粒度取决于 sharding_strategy 和 flatten 处理，直接切片不一定对应。更稳妥的办法是：先不加载任何权重，用同一随机种子初始化一份完整模型副本，得到完整参数后加载 checkpoint，再在每个 rank 上比较本地分片和完整参数对应位置。也可以先调用一次 model.state_dict()（sharded）获取本地分片，再通过 FSDP.fully_sharded_module_to_device 或其他聚合接口获取完整参数，两相验证。
更省事的校验清单
- 确认加载前后
state_dict_type一致，避免 full/sharded 混用导致比对错位。 - 比较每个参数张量的 shape、dtype、requires_grad 标志是否一致。
- 对于数值比较，建议用
torch.allclose并设置合理的atol和rtol，因为浮点计算可能带来微小误差。 - 如果模型太大无法全量比对，可以按参数名抽样，或只比对每个参数的第一/最后一个元素，但最坏情况可能漏检。
- 如果校验通过但训练指标异常，还需要检查数据加载顺序、loss 计算方式和优化器状态是否按 shard 恢复。
常见问题
1. 为什么加载时没有报错，但输出结果不对？
很可能是 checkpoint 保存的是完整参数，而你用 sharded state dict 加载，或 reverse。FSDP 的 load_state_dict 在 strict=True 时只检查 key 是否齐全，不检查分片位置。请先统一 state_dict 类型，再按上面的方法逐值比对。
2. 如何快速确认本地分片是否真的加载了？
在每个 rank 上打印本地 model.state_dict() 中的参数统计（均值、标准差、norm），与加载前的值对比。如果出现明显的 NaN 或全零，基本可以判断加载失败。但这只能发现大问题，细粒度校验仍需完整比对。
3. 有没有更简单的方式适用于 HF Accelerate？
HuggingFace Accelerate 的 FSDP 接口通常会处理 state_dict 转换，但其内部也是调用 PyTorch FSDP。可以参考 Accelerate 的 save_state 和 load_state 的默认行为，使用时注意 load_state 时是否自动调用 unwrap_model。通用原则相同：先获取完整权重，再和本地分片比对。
校验加载是否正确不是一次性的工作，建议在训练脚本中写成独立函数，每次恢复训练时执行一次。若模型参数极大，可以先只校验关键层（如 embedding、第一层和最后一层）来平衡速度和覆盖范围。