如何从多个 pth 文件加载 pytorch 分片模型

Question

我有一个模型，已经使用 FSDP 进行了训练；因此，例如，我有两个 .pth 文件，consolidated.00.pth 和 consolidated.01.pth，其中包含模型的分片权重。我希望使用未分片的模型，但是，我无法使用像

这样的典型代码来加载权重

checkpoint = torch.load("consolidated.00.pth")
model.load_state_dict(checkpoint)

这给了我关于分裂权重形状的明显错误：

layers.23.attention.wk.weight 的大小不匹配：从检查点复制形状为 torch.Size([2560, 5120]) 的参数，当前模型中的形状为 torch.Size([5120, 5120])。

如何从这两个（或更多）分片向模型加载权重？或者，有没有办法将分片权重合并到一个文件中，我可以加载？