[Fix] Load SwapAdamW state in place to keep pinned host residency by PengchengShi00 · Pull Request #2225 · InternLM/xtuner

PengchengShi00 · 2026-10-27T06:22:43Z

SwapAdamW 的优化器状态设计为常驻 pinned host 内存:self.state 中的 moment 与 _param_to_cpu_states_map 持有的别名指向同一批 buffer。但 resume 时走的是官方 torch.optim.AdamW.load_state_dict,它会:

把所有 moment .to(device=param.device) 搬上 GPU —— 瞬间占用约 2 倍参数量的显存,而这正是 SwapAdamW 要避免的;
用新对象替换 self.state 条目 —— _param_to_cpu_states_map 从此指向旧 buffer, 训练更新的是一套、save_dcp 保存的是另一套,resume 之后的 checkpoint 中 moment 静默失真,且替换后的 tensor 不再 pinned,step 的 non_blocking 拷贝退化为同步。

Fix

覆写 SwapAdamW.load_state_dict:不委托官方 loader,逐 key copy_ 进既有 pinned CPU buffer(对象、storage、别名、pin 属性全部不变);同时重放官方loader 协议(pre/post hooks、id map 位置配对、param_groups 原地 update 以保住_param_to_group_map 引用)。key 集合双向校验(忽略 amsgrad=False 的 None 占位), 空 state 显式原地清零(等价于官方"重置历史"语义)。

Test Plan

新增引擎级测试 tests/engine/test_swap_adamw_engine_resume.py:
两个引擎同一随机初始化,engine1 训练一步保存 DCP 后再训练一步,engine2 以不同 lr 加载(校验超参恢复)后训练一步,对比:

权重与 moment 数值与连续训练对齐;
moment 的 dtype/shape 与所属参数一致(防止丢失 DTensor 全局元数据退化成局部 shard);
moment 保持在 CPU 且 pinned;
self.state 与 _param_to_cpu_states_map 指向同一块 storage(load/step 两个检查点,负向验证过:无本修复时该断言精确报 storage diverged);
step 计数、lr 恢复;单卡运行并打印关键时点显存。

PengchengShi00 · 2026-10-29T09:18:27Z

Summary

修复 swap_optimizer=True 下 DCP resume 的两个静默错误:

保存侧:优化器 state 是普通 CPU 张量(local shard 形状),不带 DTensor 元数据。DCP 只对 DTensor 按
mesh/placements 计算全局布局,普通张量被当成一份完整数据,多 rank 下其余 rank 的 shard 不会被保存。
加载侧:走官方 torch.optim.Optimizer.load_state_dict,它把 moment .to(device=param.device) 并在
__setstate__ 里整体替换 self.state / self.param_groups。后果是:resume 瞬间要一份全量优化器状态的显存
(与 swap 目标相反)、_param_to_cpu_states_map 与 _param_to_group_map 两处别名悬空
(step 更新一套、save_dcp 保存另一套 → checkpoint 静默失真;新 lr 永远不生效)。

修复:state 改为“DTensor 感知的 pinned host 分配”;在 SwapAdamW 内覆写 load_state_dict 做原地 copy_ +
param_group 原地 update,并补齐协议 hooks、参数集合/ key / shape 校验。TrainEngine 零改动。

结论:APPROVE,静态 review 未发现阻断性缺陷。 本轮只应用了注释/文档类建议(见“其他 issues” N2/N6),
其余为测试健壮性建议。

Main Flowchart before this PR

graph TD
  subgraph TE["TrainEngine — xtuner/v1/engine/train_engine.py"]
    TE_save["save_dcp(weights_dir, save_optimizer=True)"]
    TE_load["load_dcp(weights_dir, load_states, load_args)"]
    TE_collect["_get_dcp_state_dict(cpu_offload=True)"]
  end
  subgraph DCP["torch.distributed.checkpoint.state_dict"]
    D_get["get_optimizer_state_dict(model, optimizer, options)"]
    D_set["set_optimizer_state_dict(model, optimizer, optim_state_dict, options)"]
    D_load["_load_optim_state_dict()"]
  end
  subgraph OPT["torch.optim.Optimizer"]
    O_sd["state_dict()"]
    O_lsd["load_state_dict(state_dict)"]
    O_cast["_process_value_according_to_param_policy()<br/>value.to(dtype=param.dtype, device=param.device)"]
    O_setstate["__setstate__()<br/>整体替换 self.state / self.param_groups"]
  end
  subgraph SW["SwapAdamW — xtuner/v1/optim/swap_adamw.py(修改前)"]
    SW_init["__init__() → _init_swap_states()"]
    SW_alloc["state[p][k] = zeros_like(to_local(p)).cpu().pin_memory()<br/>plain Tensor,无 DTensor 元数据"]
    SW_map["_param_to_cpu_states_map[p] = 与 state 同一批 buffer"]
    SW_step["step(closure)"]
    SW_put["put_state_to_device(device) → False"]
  end
  TE_save --> TE_collect --> D_get --> O_sd --> SW_init --> SW_alloc --> SW_map
  TE_load --> TE_collect
  TE_load --> D_set --> D_load --> O_lsd --> O_cast --> O_setstate
  O_setstate -. "state/group 对象被替换 → 别名断裂、moment 上 GPU" .-> SW_step
  SW_map --> SW_step --> SW_put

两个断点:

保存:Optimizer.state_dict() 输出普通 CPU tensor,DCP 不会按 rank 切分 → 多 rank checkpoint 语义被破坏;
加载:_cast 把 moment 搬到 param.device,__setstate__ 换掉 self.state/self.param_groups →
显存尖峰 + 两处别名悬空(_param_to_cpu_states_map 供 step(), _param_to_group_map 供 lr/betas)。

Main Flowchart after this PR

graph TD
  classDef changed fill:#ffe3e3,stroke:#d40000,stroke-width:2px
  subgraph TE["TrainEngine(未改动)"]
    TE_save["save_dcp(weights_dir, save_optimizer=True)"]
    TE_load["load_dcp(weights_dir, load_states, load_args)"]
    TE_collect["_get_dcp_state_dict(cpu_offload=True)"]
  end
  subgraph DCP["torch.distributed.checkpoint.state_dict"]
    D_get["get_optimizer_state_dict(model, optimizer, options)"]
    D_set["set_optimizer_state_dict(model, optimizer, optim_state_dict, options)"]
    D_load["_load_optim_state_dict() → optim.load_state_dict(state_dict=...)"]
  end
  subgraph OPT["torch.optim.Optimizer"]
    O_sd["state_dict()"]
    O_setstate["__setstate__() — 不再被走到"]
  end
  subgraph SW["SwapAdamW — xtuner/v1/optim/swap_adamw.py(修改后)"]
    SW_init["__init__() → _init_swap_states()"]
    SW_alloc["state[p][k] = empty_like(p, device='cpu', pin_memory=True).to_local().zero_()<br/>DTensor 保留 global shape / placements"]
    SW_map["_param_to_cpu_states_map[p] = state[p] 的 local view(同一 storage,未改动)"]
    SW_lsd["load_state_dict(state_dict) ← 新增覆写<br/>pre-hooks → 校验 group 布局 → 原地 copy_ → 原地更新 group → post-hooks"]
    SW_step["step(closure)"]
    SW_put["put_state_to_device(device) → False"]
  end
  TE_save --> TE_collect --> D_get --> O_sd
  O_sd -. "DTensor:每个 rank 只写自己的 shard" .-> TE_save
  TE_load --> TE_collect
  TE_load --> D_set --> D_load --> SW_lsd
  SW_lsd -. "只改值不改对象 → 别名与 pin 属性保持" .-> SW_step
  SW_init --> SW_alloc --> SW_map --> SW_step --> SW_put
  class SW_alloc,SW_lsd changed

说明(mermaid 8.8 支持的强调方式:classDef + class,红色节点为本次改动):

流程节点	修改点与代码位置
初始化 swap 状态	`_init_swap_states()`(`swap_adamw.py:169-208`)改用 `torch.empty_like(param, device="cpu", pin_memory=True)`:DTensor 参数得到“CPU + pinned 的 DTensor”,其 `to_local()` 即本 rank 的 pinned shard;不再先在 GPU `zeros_like` 再拷回 CPU
保存 checkpoint	`Optimizer.state_dict()` 未改;输出已是 DTensor,`InterleavedShardSavePlanner` 走 `RuntimeLayout.from_dtensor` 分支,每 rank 写自己的 shard
加载 checkpoint	新增 `SwapAdamW.load_state_dict()`(`swap_adamw.py:67-161`);`set_optimizer_state_dict` → `_load_optim_state_dict` 直接调用覆写方法,不再经过 `_cast` / `__setstate__`
训练步	`step()` 未改;storage 未变,step 更新的 buffer 与 `state_dict()` 导出的 buffer 仍是同一块

核心原理

state 必须是 DTensor。 DCP 保存计划对每个 value 分两类:isinstance(obj, DTensor) 时用 device_mesh /
placements 计算全局形状与 offset;否则按普通张量整份保存(多 rank 去重后只落一份)。
empty_like(param, device="cpu") 让 state 的 spec 与参数一致,才能按 rank 分片读写。
该 op 在 DTensor 上由 aten.empty_like.default 的 propagate_single_input_strategy 处理:
输出 spec 继承输入,非张量 kwargs(device / pin_memory / memory_format)透传给 local op,
因此 pin 是真正作用在 local tensor 上的,不会被静默丢弃。
必须原地加载。 官方 loader 的 _cast 对非 step key 一律 value.to(dtype=param.dtype, device=param.device),
随后 __setstate__ 替换 self.state、self.param_groups。原地 copy_ 同时保住:对象与 storage(别名有效)、
is_pinned(non_blocking 仍有效)、dtype 转换(copy_ 自带 cast)。
比较 local shape 而非 global shape。 加载时 dcp.load 已按当前 rank 的布局完成再分片,
target / source 都是 local shard;target.shape != source.shape 只会在“旧格式 checkpoint / 布局变化”时触发,
而这正是 copy_ 会静默广播、写进错 moment 的场合。
meta 参数下的分配。 TrainEngine.build_model() 在 torch.device("meta") 下建模,SwapAdamW 在 meta 参数上构造;
empty_like(meta_dtensor, device="cpu") 按 meta 的 shape/spec 直接分配真实 pinned 内存,
初始化阶段不再产生 GPU 临时张量。该路径由引擎级用例真实覆盖。

核心实现与对应的单测

实现	位置	覆盖它的断言
DTensor + pinned 状态分配	`swap_adamw.py:189-201`	`TestSwapAdamWResume`:`isinstance(moment, DTensor)`、`moment.shape == param.shape`、`moment.placements == param.placements`、`local2.is_pinned()`
原地加载(值 / step / groups)	`swap_adamw.py:118-157`	`test_load_state_dict_applies_moments_on_pinned_host`:moment `rtol=0, atol=0`、`step == 7`、`lr == 5e-4`、仍在 CPU + pinned
空 state 重置历史	`swap_adamw.py:124-128`	同一用例后半段:`step == 0`、moment 全 0(对应 engine 的 `load_states=False`)
参数集合 / key / shape 校验	`swap_adamw.py:97-116`、`130-147`	`test_load_state_dict_rejects_incompatible_state`:key 多、key 少、shape 错三种 `ValueError`
端到端等价	全流程	`TestSwapAdamWResume`:`lr == 1e-4`、`step == 2`、moment 与权重逐参数 `assert_close`(world_size=2)

测试结构(198 行 / 2 个 TestClass / 3 个用例):

TestSwapAdamWResume(DeterministicDDPTestCase,world_size = 2):真实 save_dcp → load_dcp → 再训一步;
TestSwapAdamWLoadStateDict(普通 TestCase + nn.Linear(8, 4),毫秒级):只打 load_state_dict 公开协议。

这两层刚好对应 bug 的两个面:修复前 resume 用例会在 DTensor/pinned/moment 断言处失败;
把 _init_swap_states 回退成普通张量时,isinstance(moment, DTensor) 会失败。

抽象与信息隐藏评估

好的部分

修复完全收敛在 SwapAdamW:TrainEngine / Optimizer 基类 / DCP 适配层零改动、零新字段、零新类。
_init_swap_states 由“isinstance(DTensor) 分叉 + 两套分配”收敛成一条 empty_like + _to_local_tensor,
减少分支,把“state 布局随参数”的规则藏在一处。
_to_local_tensor 不是纯转发(隐藏 DTensor / 普通张量差异),保留合理。
load_state_dict 用 5 个带注释的功能块表达 pre-hooks → 校验 → 拷贝 → 更新 group → post-hooks,
读者不用跳转即可理解顺序。

与基类协议的 3 处分歧,均已在代码中说明

空 state → 显式清零(基类不触碰):docstring 已写“An empty state explicitly resets history”;
参数 id 集合 / state key 集合 / shape 严格校验(基类容忍):对应代码块上方均注释了“为什么选择早失败”;
group 更新为 merge 而非“以 saved group 为准重建”(checkpoint 缺失 key 保留本地值):已在块头注明
“Unlike torch, this merges into the existing groups, so a key absent from the checkpoint keeps its local value”。

依赖 torch 私有 attrs(可接受取舍,已在 docstring 注明)

load_state_dict 直接遍历 self._optimizer_load_state_dict_pre_hooks / _optimizer_load_state_dict_post_hooks,
没有公开替代(调用 super() 就会触发 _cast)。docstring 已注明 “replayed from the base class' private hook
dicts (torch>=2.6)”,便于以后判断升级风险。

_param_to_cpu_states_map 是为 step() 预先持有 local view 的既有设计,本 PR 不动,也不在此提出调整建议。

新增功能守则(既有类零污染)

本 PR 是 bug fix,不新增功能类:

新增类 0 个;TypedDict / 数据类 0 个。
既有类接触点:SwapAdamW 新增 1 个公开方法 + _init_swap_states 内 1 处分配改写,无新增字段;
TrainEngine / AdamW / Optimizer 接触点 0(R1 满足)。
既有类中未出现新特性的实现细节(切分规则、指标 key 等)(R2 满足)。
没有新增“按特性分叉”的条件,反而删掉一个 isinstance(DTensor) 分叉(R3 满足)。

公开接口的线性业务流程评估

load_state_dict 是 5 个同层级步骤的线性叙述:pre-hooks → 校验 group 布局 → 原地拷贝 state → 原地更新 group → post-hooks。
曾考虑拆成 2~3 个私有方法,最终保留“单一方法 + 块注释”:这些块共享 saved_groups / saved_state / id_map,
拆函数要反复传参,Locality 反而变差;每块都短,过不了 deletion test。

唯一注意点:方法已约 95 行;若将来再增加“按参数类型分派”之类的逻辑,优先把“拷贝单个参数 state”这一步拆出去
(规则明确、边界清晰,可独立测试):

# 仅在未来膨胀时考虑
def _copy_one_param_state(self, target_state, source_state, param_id) -> None:
    ...

单测建议

world_size = 2 是硬性要求,单卡机器会直接失败而非 skip。建议加 guard,让本地单卡环境优雅跳过:

@unittest.skipUnless(torch.cuda.device_count() >= 2, "needs >=2 GPUs for the 2-rank DCP resume case")
def test_resume_after_dcp_matches_continuous_training(self) -> None:
    ...

测试只验证正确性,不内建计时;需要耗时用 pytest --durations=0(resume 用例是唯一重负载,2 rank 并行)。

其他 issues

N1(已应用)补齐 @torch._disable_dynamo。 已加到覆写方法上(外层 @torch._disable_dynamo、内层 @torch.no_grad()),
与 torch.optim.Optimizer.load_state_dict 保持一致:告诉 Dynamo 不要追踪该方法体。该方法带 ParamSpec 重载,
不会像公共的 torch.compiler.disable 那样把签名退化成 Any;eager 路径下无任何运行时开销。

N2(已应用)类 docstring 补充 DTensor 不变量。 已加:对 DTensor 参数,state 必须是同 mesh / 同 placements 的
DTensor,因此每个 rank 只 checkpoint 自己的 shard。

N3(nit)expected_keys 是 set,错误路径可能部分写入。 for key in expected_keys 顺序不确定,shape 校验在循环内,
一旦后面的 key 报错,前面已 copy_ 的 key 已写入 buffer。错误路径直接抛出、进程中止,可接受;
若要“要么全成功要么全不动”,可先统一校验 shape 再统一拷贝。

N4(nit)assert isinstance(source, torch.Tensor) 在 python -O 下会被移除;后面紧跟 shape 校验,风险很小。

N5(流程)本地分支落后 main(merge-base 8e6c1824),提 PR / 跑 CI 前建议先 rebase。

Verdict

APPROVE — 无阻断项。本轮已应用:N1(@torch._disable_dynamo)、N2(类 docstring 的 DTensor 不变量)、
group merge 语义与私有 hook dict 依赖的注释说明。
仍待处理(测试代码改动,本轮未动):单测建议 1(2 卡 guard)、单测建议 2(amsgrad 分支);N3/N4 为可选 nit。

PengchengShi00 added 2 commits October 27, 2026 14:19

[Fix] Load SwapAdamW state in place to keep pinned host residency

02d738b

[Test] Number param id map globally across param groups

6898f01

PengchengShi00 requested a review from jayhenry October 27, 2026 09:05

PengchengShi00 mentioned this pull request Oct 27, 2026

[Fix] Restore SwapAdamW state and parameter aliases after checkpoint loading #2175

Closed

[Fix] Validate SwapAdamW state layout before loading it in place

36e76ce

Navigation Menu

Repository navigation

[Fix] Load SwapAdamW state in place to keep pinned host residency - #2225

[Fix] Load SwapAdamW state in place to keep pinned host residency#2225
PengchengShi00 wants to merge 3 commits into
InternLM:mainfrom
PengchengShi00:fix_swapAdamW

PengchengShi00 commented Oct 27, 2026

Uh oh!

PengchengShi00 commented Oct 29, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

PengchengShi00 commented Oct 27, 2026

Fix

Test Plan

Uh oh!

PengchengShi00 commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Main Flowchart before this PR

Main Flowchart after this PR

核心原理

核心实现与对应的单测

抽象与信息隐藏评估

新增功能守则(既有类零污染)

公开接口的线性业务流程评估

单测建议

其他 issues

Verdict

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

PengchengShi00 commented Oct 29, 2026 •

edited

Loading