【Bug已解决】FSDP2 + SHARDED_STATE_DICT： optimizer checkpoint load fails with missing step key. (RuntimeE

【Bug已解决】FSDP2 SHARDED_STATE_DICT optimizer checkpoint load fails with missing step key. (RuntimeError Missing key in checkpoint state_dict optimizer.state.0.step.) 解决方案一、现象长什么样用 FSDP2 做全分片训练并在保存 / 恢复时启用StateDictType.SHARDED_STATE_DICT以便每个 rank 只存自己那一份分片。保存时一切正常但重启加载时直接抛出RuntimeError Missing key in checkpoint state_dict optimizer.state.0.step.也就是优化器状态字典里的state子字典期望有optimizer.state.0.step第 0 号参数的步数计数器可加载进来的分片里根本没有这个键。最让人困惑的地方保存阶段没有报错torch.save成功落盘优化器的step在训练时明明一直在自增学习率调度依赖它加载代码看起来和官方示例一致optimizer.load_state_dict(ckpt[optimizer])只要改用FULL_STATE_DICT就能正常加载SHARDED_STATE_DICT一上就炸。这说明问题不在step 没被保存而在sharded 格式的键结构与加载端期望的键结构对不上。二、背景FSDP2 通过torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type这个上下文管理器来决定状态字典的形态FULL_STATE_DICT聚合回完整、未分片的字典键是全局 param 索引SHARDED_STATE_DICT每个 rank 只持有自己分片的那部分键是本地param 索引从 0 开始且state里每个参数组只含本 rank 负责的张量LOCAL_STATE_DICT最原始的本地位姿键也是本地索引。关键点优化器状态字典的state子字典其键param 索引是在当前这个 state_dict_type 上下文里生成的。保存时若在SHARDED_STATE_DICT上下文里调用optimizer.state_dict()得到的是按本地索引、step被切到对应 rank 的字典加载时若不在相同上下文里调用optimizer.load_state_dict()PyTorch 会用另一套键结构去对齐于是找不到optimizer.state.0.step。更隐蔽的是step是一个 Python 标量 / 小张量sharded 保存时它属于哪个 rank 负责该参数分片的那一份。若保存端和加载端的 rank 数、分片方式、或者 state_dict_type 不一致step这一项就会落在错误的分片里加载端自然缺失。三、根因把不一致的上下文抽象成代码示意非照抄源码# 保存端正确在 SHARDED 上下文里取 optimizer state with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): optim_sd optimizer.state_dict() # 键是本地索引step 随分片分布 torch.save({optimizer: optim_sd}, path) # 加载端错误脱离了上下文用默认 FULL 的键结构去对齐 ckpt torch.load(path) optimizer.load_state_dict(ckpt[optimizer]) # 键对不齐 - Missing step根因链条optimizer.state_dict()的键结构由调用它时所处的state_dict_type上下文决定保存端在SHARDED_STATE_DICT内键为本地索引step被切到对应 rank加载端没有重新进入相同上下文PyTorch 默认按FULL_STATE_DICT的全局索引去匹配两边键空间不一致optimizer.state.0.step在加载端期望全局索引 0但文件里它是某个本地分片里的键对齐失败抛Missing key——这是典型的保存 / 加载上下文不对称导致的 KeyError。另一个常见变体保存时用了SHARDED但optimizer.state_dict()之前没有model被state_dict_type包裹即 optimizer 状态其实还是 FULL 形态存进去的加载端却按 SHARDED 去读同样对不上。四、最小可运行复现用一段纯 Python 模拟键空间不一致导致step缺失# repro_sharded_step.py class OptState: def __init__(self): self.state {} # param_index - {step: int, exp_avg: ...} def save_sharded(self, rank, world): # 模拟 SHARDED每个 rank 只存自己负责的本地索引 sharded {} for idx in range(rank, len(self.state), world): sharded[idx] self.state[idx] # 本地索引 return sharded def load_full_expect(self, sharded): # 模拟加载端用 FULL 键结构全局索引 0..N-1去对齐 for idx in range(len(self.state)): if idx not in sharded: raise KeyError(foptimizer.state.{idx}.step) def main(): opt OptState() opt.state {0: {step: 7}, 1: {step: 7}, 2: {step: 7}, 3: {step: 7}} world 2 rank0 opt.save_sharded(0, world) # 本地索引 {0:.., 2:..} rank1 opt.save_sharded(1, world) # 本地索引 {1:.., 3:..} merged {**rank0, **rank1} # 合并后键是 0,1,2,3 - 其实能对齐 # 但下面模拟加载端只拿到了 rank0 的分片却按全局 0..3 期望 try: opt.load_full_expect(rank0) # 只给 rank0 分片 except KeyError as e: print(复现成功 -, e) if __name__ __main__: main()运行输出复现成功 - optimizer.state.1.step这正是真实 bug 的抽象step被切进了不同 rank 的分片加载端若没用对称的上下文 / 没合并完整分片就会缺少某些step键。五、解决方案第一层最小直接修复保证保存与加载都在同一个state_dict_type上下文里这是最小且必须的一步# fix_layer1.py import torch from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import StateDictType def save_sharded(model, optimizer, path): with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): sd { model: model.state_dict(), optimizer: optimizer.state_dict(), # 与上下文一致 } torch.save(sd, path) def load_sharded(model, optimizer, path): ckpt torch.load(path) # 关键加载也进入完全相同的 SHARDED 上下文 with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer]) # 键结构对齐 - step 不再缺失只要 save / load 都用SHARDED_STATE_DICToptimizer.state.0.step这类键就会在两侧用同一套本地索引键空间生成与消费缺失问题消失。六、解决方案第二层结构性改进把对称上下文收敛成一个职责单一的 IO 助手避免任何调用方忘记配对上下文并显式确保step被纳入优化器状态# fix_layer2.py from dataclasses import dataclass from enum import Enum from typing import Any, Callable class StateDictKind(str, Enum): SHARDED sharded FULL full LOCAL local dataclass(frozenTrue) class CkptSpec: kind: StateDictKind def make_state_dict_context(model, spec: CkptSpec): 单一来源根据 spec 返回对应的 state_dict_type 上下文。 from torch.distributed.fsdp import FSDP, StateDictType mapping { StateDictKind.SHARDED: StateDictType.SHARDED_STATE_DICT, StateDictKind.FULL: StateDictType.FULL_STATE_DICT, StateDictKind.LOCAL: StateDictType.LOCAL_STATE_DICT, } return FSDP.state_dict_type(model, mapping[spec.kind]) def save_checkpoint(model, optimizer, path, spec: CkptSpec): with make_state_dict_context(model, spec): torch.save({ # type: ignore[name-defined] model: model.state_dict(), optimizer: optimizer.state_dict(), }, path) def load_checkpoint(model, optimizer, path, spec: CkptSpec): ckpt torch.load(path) with make_state_dict_context(model, spec): # 强制对称 model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer]) def assert_step_present(optimizer) - None: 结构性守护确保 step 计数器确实在优化器状态里。 for grp in optimizer.param_groups: for p in grp[params]: st optimizer.state[p] assert step in st, optimizer.state 缺少 step 计数器要点save_checkpoint/load_checkpoint共用make_state_dict_context上下文对称由结构保证无法被某个调用方漏掉assert_step_present在加载后立即校验每个参数的step把缺失提前变成显式异常CkptSpec成为唯一真相来源切换 FULL / SHARDED 只需改一处。七、解决方案第三层断言 / CI 守护写一条 pytest保存 - 加载往返断言step不丢且数值相等把回归锁进 CI# test_sharded_optim_ckpt.py import pytest class FakeOpt: def __init__(self, n): self.state {i: {step: 7} for i in range(n)} self.param_groups [{params: list(range(n))}] def state_dict(self): return {state: {k: dict(v) for k, v in self.state.items()}} def load_state_dict(self, sd): # 模拟 PyTorch按 state 键对齐缺失即报错 for idx in sd[state]: pass for idx in self.state: if idx not in sd[state]: raise KeyError(foptimizer.state.{idx}.step) self.state {int(k): v for k, v in sd[state].items()} def test_step_roundtrip_sharded(): opt FakeOpt(4) sd opt.state_dict() opt2 FakeOpt(4) opt2.load_state_dict(sd) # 对称同一键空间 assert opt2.state[0][step] 7 def test_missing_step_detected(): opt FakeOpt(4) sd opt.state_dict() del sd[state][0] # 模拟 step 被切丢 with pytest.raises(KeyError): opt.load_state_dict(sd) def test_context_symmetry_required(): 加载端必须和保存端用同一 state_dict_type 键空间。 saved {0: 7, 2: 7} # 仅 rank0 分片本地索引 missing {1, 3} # 全局期望 0..3 for idx in missing: assert idx not in saved # 说明不对称会缺键CI 一旦回归到加载脱离上下文test_missing_step_detected与往返测试会立即变红。八、排查清单遇到Missing key ... optimizer.state.0.step时确认保存端optimizer.state_dict()是否包在FSDP.state_dict_type(..., SHARDED_STATE_DICT)内确认加载端optimizer.load_state_dict()是否包在完全相同的上下文内比对两边 rank 数 / 分片方式是否一致world size 变了也会错位若用accelerate确认accelerator.load_state的state_dict_type与保存时一致加载后立即用assert_step_present(optimizer)校验step存在临时改用FULL_STATE_DICT验证是不是键空间不对称——能加载就坐实本 bug把第七节的 pytest 接进 CI作为 sharded checkpoint 的回归护栏。九、小结FSDP2 SHARDED_STATE_DICT下优化器加载报Missing key optimizer.state.0.step根因是保存与加载所处的state_dict_type上下文不对称保存端在 SHARDED 上下文里按本地索引切分了step加载端却脱离上下文、用全局索引键空间去对齐于是键对不上。三层层级第一层保存与加载都进入相同的SHARDED_STATE_DICT上下文键结构对齐第二层用CkptSpecmake_state_dict_context把对称上下文收敛成单一入口并加assert_step_present提前暴露缺失第三层写 pytest 做保存-加载往返与缺失检测锁进 CI。核心教训任何状态字典形态由上下文决定的 API保存与加载必须成对地处在同一上下文里凡是不成对的都是 KeyError / 静默错位的温床。