Repository navigation
Conversation
Consolidate the tested final change; omit superseded attempts and unrelated topic history.
|
@claude review |
|
Claude: Summary本 PR 通过 Main Flowchart after this PRflowchart TD
A["train_engine.load_dcp\nset_optimizer_state_dict"] --> B["torch _load_optim_state_dict"]
B --> C["Optimizer.load_state_dict\n_process_value_according_to_param_policy\nvalue.to(param.dtype, param.device)"]
C --> D["post hook _restore_swap_states\n(本 PR 新增)"]
D --> E["_init_swap_states\n复用 restored moments + pin_memory 回 CPU"]
E --> F["step() 使用 _param_to_cpu_states_map"]
style C fill:#ffd9d9,stroke:#d33,stroke-width:2px
style D fill:#fff3cd,stroke:#e0a800,stroke-width:2px
style E fill:#fff3cd,stroke:#e0a800,stroke-width:2px
核心原理实现与单测核心行为（恢复 moments / step / param-group 别名，而非归零）有真实代码路径覆盖：新增的 未被真实代码路径覆盖的是 DTensor / 多卡恢复分支（ 抽象与信息隐藏评估
单测建议
VerdictCOMMENT |
| self._param_to_group_map: dict[torch.Tensor, dict] = {} | ||
| self._param_to_cpu_states_map: dict[torch.Tensor, dict[str, torch.Tensor | None]] = {} | ||
| self._init_swap_states() | ||
| self.register_load_state_dict_post_hook(self._restore_swap_states) |
There was a problem hiding this comment.
Claude: [正确性] [复杂] post-load hook 的 Seam 位置偏晚。torch.optim.Optimizer.load_state_dict 在调用 hook 之前，就已经通过 _process_value_according_to_param_policy 把 state_dict["state"] 里每个 moment 做 value.to(dtype=param.dtype, device=param.device)，并且是先构造完整的新 state 字典、再 __setstate__。也就是说在 hook 运行之前，全部 exp_avg / exp_avg_sq（amsgrad 时还有 max_exp_avg_sq）会同时落在 accelerator 上，峰值约为 2~3 倍参数量。而 SwapAdamW 存在的前提恰恰是 optimizer state 放不下显存，因此 resume 很可能在 set_optimizer_state_dict（train_engine.py#L580）处 OOM，修复在生产配置下可能走不到。
需要说明的是，这不是本 PR 引入的回归：修复前这些 device 张量还会长期留在 self.state 里（常驻显存泄漏，且 step() 仍读旧 pinned buffer），本 PR 已经消除了常驻部分，只是没有消除峰值。
建议把 Seam 前移到 load_state_dict 本身，先把 moments 摘出来再委托给父类，这样只有 step 经过 torch 的 cast：
def load_state_dict(self, state_dict):
moments = {pid: {k: s.pop(k) for k in self._state_keys if k in s}
for pid, s in state_dict["state"].items()}
super().load_state_dict(state_dict) # 只剩 step 需要 cast
self._init_swap_states() # 重建 group / state 别名
self._assign_cpu_moments(moments) # pin 到 CPU，全程不碰 deviceDCP 路径最终也是调用 optim.load_state_dict，所以覆盖面与当前 hook 相同。
| local_param, device="cpu", memory_format=torch.preserve_format | ||
| ) | ||
| else: | ||
| cpu_tensor = self._to_local_tensor(restored).to(device="cpu") |
There was a problem hiding this comment.
Claude: [设计] [复杂] 这里恢复后写回 self.state[param] 的是 plain local CPU tensor（与 __init__ 时 zeros_like(local_param) 的既有行为一致），对 DTensor 参数并不会还原成 DTensor。存档端 train_engine.py#L383 的 get_optimizer_state_dict 对 FSDP2 并不做 plain→DTensor 转换（info.fsdp_modules 为空），因此 dcp.save 的 planner 会按相同 FQN 跨 rank 去重，存在只持久化单个 rank 分片的风险；那么本 PR 修好的 restore 在 world_size > 1 时会忠实地恢复出错误的状态。
存档端行为是既有问题、不在本 PR diff 内，但它决定了这次修复能否真正达成「resume 正确」的目标，所以希望确认一下多卡 round-trip：
- 变量名
cpu_state_dtensor暗示这里原本期望存 DTensor，是否应在 restore / init 时用DTensor.from_local保持 placement？ - 或者明确记录「SwapAdamW 仅支持单 rank / replicate 存档」这一约束。
另外，新增的 self._to_local_tensor(restored) DTensor 分支目前没有任何测试覆盖（PR 描述中已说明多 rank 未验证），建议补一个 @pytest.mark.gpu 的多 rank save→load→step 用例，否则这条生产主路径的回归无法被 CI 捕获。
AdamW.load_state_dict casts every saved moment to the parameter device before any post hook runs, so resuming briefly held the full optimizer state (2-3x parameter size) in accelerator memory - the exact allocation SwapAdamW exists to avoid. Override load_state_dict to extract the moments first so only the scalar step tensors go through the base cast, then rebuild the swap states from host-side tensors. The DCP resume path funnels through load_state_dict as well, so coverage is unchanged. The load hook is folded into the override: the old post hook ran before the moments could be re-injected and rebuilt zero state twice.
|
Addressed both review findings; new head is 5f1feab (rebased on latest main). W1 — post-load hook seam is too late (2-3x-parameter VRAM peak at resume): accepted. `load_state_dict` is now overridden directly (a7db228): the moments are extracted from `state_dict["state"]` before delegating to `super().load_state_dict`, so only the scalar `step` tensors pass through the base-class device cast; the swap states are then rebuilt from host-side tensors and pinned without any device round-trip. The registered post-load hook is gone — it fired inside `super().load_state_dict` before the moments could be re-injected and rebuilt zero state twice. Coverage is unchanged: the DCP resume path (`set_optimizer_state_dict` → `_load_optim_state_dict`) ends in `optim.load_state_dict(...)`, so both entry points hit the override. Regression test `test_load_state_dict_never_stages_moments_on_device` spies on `Optimizer._process_value_according_to_param_policy` and asserts no tensor with `numel() > 1` ever enters the cast path. W2 — plain-tensor restore loses DTensor placement / multi-rank risk: addressed in 20272cc on the same push. `_init_swap_states` now keeps each moment as a CPU DTensor (`torch.empty_like(param, device="cpu", pin_memory=True)`) so DCP persists every rank's shard instead of deduplicating plain locals as replicas; a shape check rejects moments lacking complete DTensor metadata, and `_validate_optimizer_checkpoint_shapes` (train_engine) refuses to resume from legacy checkpoints whose saved global shape is smaller than the live parameter — with a message naming the lost shards. The variable that motivated the question (`cpu_state_dtensor`) now genuinely holds a DTensor. Covered by `test_sharded_checkpoint_keeps_every_rank_moment`, a real 2-rank save→load→step round-trip (gloo on CPU, NCCL on 2×GPU) asserting per-rank shard equality, pinned-CPU placement, and identical next-step numerics for both sync and async DCP saves, plus `test_incomplete_optimizer_checkpoint_is_rejected_before_loading`. All CPU variants of `tests/optim/test_swap_adamw_resume.py` pass locally after the rebase. |
torch.empty_like on a DTensor parameter keeps the DTensor type at runtime, but mypy models the return as plain Tensor and rejects .to_local(); assert the placement so strict checking passes.
|
有两个问题：
|
| if load_states: | ||
| _validate_optimizer_checkpoint_shapes( | ||
| state_dict["optimizer"], dcp.FileSystemReader(weights_dir).read_metadata() | ||
| ) |
Motivation
AdamW.load_state_dict()replaces optimizer state and parameter-group objects, butSwapAdamWcaches separate references to both. After loading a checkpoint, its swap path can therefore use stale groups or moment buffers. Simply rerunning the existing initializer is also unsafe because it initializes moments and step counters to zero, discarding the restored optimizer history.Modification
Register a post-load hook that rebuilds the parameter-group and CPU-state maps from the loaded state. Make initialization preserve existing moments, AMSGrad maximum moments and step counters, moving the canonical state back to CPU and pinning moment buffers. Allocate zero state only for entries that do not already exist.
The hook covers ordinary
load_state_dict()and the distributed-checkpoint optimizer loading path. The checkpoint format and optimizer API are unchanged.Testing
4 passed, 4 skipped. Compares the next update with uninterrupted SwapAdamW and standard AdamW after direct and single-process DCP loading, with AMSGrad enabled and disabled. The CPU tests replace device synchronization and pinning. CUDA cases were skipped; multi-rank restoration was not exercised.