PyTorch 模型加载 - DDP 与 torch.compile 权重前缀

flyfish

在 PyTorch 分布式训练与性能优化场景中，DistributedDataParallel (DDP) 和 torch.compile 二者都会对模型的 state_dict（权重字典）键名进行包装修改，导致训练保存的权重无法直接加载到原生模型中，出现 Missing key(s) / Unexpected key(s) 报错。

DDP 对模型权重的影响

1. 包装原理

DDP 是 PyTorch 官方分布式训练方案，它会将原始模型作为子模块，封装到 DistributedDataParallel 类中，实现多卡梯度同步。
原始模型：model（原生 nn.Module）
DDP 包装后：ddp_model = DDP(model)，其中 ddp_model.module 才是原始模型

2. 对 state_dict 的影响

调用 ddp_model.state_dict() 时，PyTorch 会递归包含所有子模块的参数，因此所有权重键名都会自动加上 module. 前缀：
原始键名：head.weight、features.0.conv.weight
DDP 后键名：module.head.weight、module.features.0.conv.weight

DDP 仅修改键名前缀，不改变权重的数值本身。

torch.compile 对模型权重的影响

1. 包装原理

PyTorch 2.0 引入的 torch.compile 通过 JIT 编译优化计算图，显著提升训练/推理速度。它会将原始模型包装为 OptimizedModule 对象。
原始模型：model（原生 nn.Module）
编译后：compiled_model = torch.compile(model)，其中 compiled_model._orig_mod 才是原始模型

2. 对 state_dict 的影响

调用 compiled_model.state_dict() 时，所有权重键名会加上 _orig_mod. 前缀：
原始键名：head.weight
编译后键名：_orig_mod.head.weight

torch.compile 同样只修改键名，不改变权重数值；且编译产物与 PyTorch 版本强绑定，跨版本兼容性差。

DDP + compile 同时使用

这里顺序为：先编译模型，再用 DDP 包装（谁先谁后继续看下面）。

# 每个进程内执行
model = MyModel().to(device)
model = torch.compile(model)       # 1. 先编译模型
model = DDP(model, device_ids=[rank])  # 2. 再用 DDP 包装

此时 state_dict 的键名会叠加两层前缀：module._orig_mod. + 原始键名，例如 module._orig_mod.head.weight。

DDP works with TorchDynamo. When used with TorchDynamo, apply the DDP
model wrapper before compiling the model, such that torchdynamo can
apply DDPOptimizer (graph-break optimizations) based on DDP bucket
sizes.

参考网址 https://docs.pytorch.org/docs/2.12/notes/ddp.html

DDP 可与 TorchDynamo 搭配使用。结合 TorchDynamo 时，应先为模型添加 DDP 包装，再对模型执行编译，这样 TorchDynamo 才能基于 DDP 的梯度桶大小应用 DDPOptimizer（图中断优化）

训练时正确保存权重

同时使用 DDP + compile

# 取出最内层原始模型再保存
torch.save(model.module._orig_mod.state_dict(), "checkpoint.pth")

1. 只在主进程（rank 0）中执行保存，避免多进程写入冲突；
2. 永远保存 state_dict，不要直接保存整个模型对象（torch.save(model)），后者兼容性极差。

推理加载：已有权重如何清洗前缀

如果已经有了带前缀的权重文件（如第三方预训练权重、历史 checkpoint），就需要在加载时清洗前缀

1. 清洗逻辑

遍历权重字典的所有键，依次移除 module. 和 _orig_mod. 前缀，兼容单层、双层前缀的场景。

state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
new_state_dict = {}
for k, v in state_dict.items():
    # 移除 DDP 带来的 module. 前缀
    if k.startswith("module."):
        k = k[7:]
    # 移除 compile 带来的 _orig_mod. 前缀
    if k.startswith("_orig_mod."):
        k = k[10:]
    new_state_dict[k] = v

对于 module._orig_mod.head.weight 这种双层前缀，先去掉 module. 得到 _orig_mod.head.weight，再去掉 _orig_mod. 得到原始键名。

2. strict 参数的使用

strict=True（默认）：要求权重键与模型结构完全匹配，多键、少键都会报错。适合正式推理环境，能及时发现结构不匹配问题。
strict=False：忽略不匹配的键，仅加载匹配部分。仅在微调、加载部分权重时使用，纯推理场景不建议开启。