RuntimeError: don’t know how to restore data location of torch.storage.UntypedStorage (tagged with cpu:0)

快速结论：该报错通常出现在 CPU + 多进程（如 torchrun + gloo）场景下用 Trainer 从 checkpoint 恢复训练时，加载 optimizer state 的过程中失败。优先排查 map_location 是否被传成了带索引的 CPU 设备（如 cpu:0），而不是纯 "cpu"。

适用环境：Issue 中确认的环境为 transformers 5.17.0、macOS-27.0-arm64、Python 3.12.8、PyTorch 2.14.0、huggingface_hub 1.32.0、safetensors 0.8.0、accelerate 1.15.0、未安装 DeepSpeed，分布式方式为 torchrun --nproc_per_node=2 加 CPU / gloo 后端。

最快修复方案：暂无确认的一步修复方案。Issue 中维护者已合入修复 PR #49123，其思路是在多进程场景下，当设备为 CPU 时，让 _load_optimizer_and_scheduler 使用纯 "cpu" 作为 map_location，而不是带索引的 cpu:0；可优先尝试升级到包含该修复的 transformers 版本。

注意事项：该修复仅针对 CPU 设备的 map_location 回退逻辑，多 GPU 行为保持不变。Issue 中提到新增的回归测试是在单进程内 mock world_size=2 与 cpu:0，并未真正跑两 worker 的 resume 路径；真实两 worker 的恢复由 @gitedmond 和 @neevmodh 分别手动验证通过，但自动化覆盖仍属于后续可补充项。

问题场景

用户在 CPU 环境下用 torchrun --nproc_per_node=2 和 gloo 后端跑自定义 Trainer 训练任务。从零开始训练可以正常保存 checkpoint，但新建 Trainer 并设置 resume_from_checkpoint="out/checkpoint-2" 时，两个进程都在加载 optimizer state 阶段崩溃，无法继续训练。

报错原文

training from scratch: OK
  File ".../transformers/trainer.py", line 1745, in _prepare_for_training
    self._load_optimizer_and_scheduler(resume_from_checkpoint)
  File ".../transformers/trainer.py", line 3835, in _load_optimizer_and_scheduler
    torch.load(
  File ".../torch/serialization.py", line 1605, in load
RuntimeError: don't know how to restore data location of torch.storage.UntypedStorage (tagged with cpu:0)

原因分析

根据 Issue 中的定位，根因在 _load_optimizer_and_scheduler：当 world_size > 1 时，Trainer 会把 self.args.device 作为 map_location 传给 torch.load，以便把 optimizer state 直接加载到各设备上。在 CPU 多进程场景下，self.args.device 是一个带索引的设备对象（如 cpu:0），而 PyTorch 的 CPU 反序列化器只认纯 "cpu" 这种 location tag，因此每个 rank 都会在 resume 时崩溃。单进程路径之所以不受影响，是因为它已经回退到使用 "cpu" 作为 map location。

环境排查

确认 transformers 版本是否仍包含未修复的 _load_optimizer_and_scheduler 逻辑，Issue 中为 5.17.0。
确认 PyTorch 版本，Issue 中为 2.14.0。
确认是否在 CPU 上以多进程方式运行，Issue 中为 torchrun --nproc_per_node=2 加 gloo 后端。
确认触发场景是否为从 checkpoint 恢复，而不是从零训练。
确认 args.device 是否可能被解析为带索引的 CPU 设备（如 cpu:0）。

解决步骤

先确认当前是否命中同一条件：多进程、CPU 设备、resume_from_checkpoint，且报错发生在 _load_optimizer_and_scheduler 的 torch.load 调用处。
可优先尝试升级到包含 PR #49123 修复的 transformers 版本；该 PR 已将 CPU 设备的 map_location 回退为纯 "cpu"，并保留非 CPU 设备的原有行为。
如果暂时无法升级，可参考该修复思路，在本地检查或修改 _load_optimizer_and_scheduler 中传入 torch.load 的 map_location：当设备是 CPU 时使用 "cpu"，仅在确为 GPU 设备时使用带索引的设备对象。修改前请自行评估对其它分布式场景的影响。
修改或升级后，用相同的两进程 CPU resume 流程重新运行复现脚本，确认 checkpoint-2 能恢复并继续训练。

验证方法

用 Issue 中的复现方式重新执行：先跑一次从零训练生成 out/checkpoint-2，再用 resume_from_checkpoint="out/checkpoint-2" 启动新的 Trainer。若两个进程都不再报 RuntimeError: don't know how to restore data location of torch.storage.UntypedStorage (tagged with cpu:0)，并最终打印出 resume: OK，则说明问题已解决。Issue 中 @gitedmond 与 @neevmodh 均基于修复 PR 手动验证过真实两 worker 的 CPU resume 路径。