已合并
swap_adam_optimizer #982
swap_adam_optimizer #982
已合并
songjiaqi创建于 6月29日
songjiaqi
songjiaqi成员
6月29日

What type of PR is this?

/kind feature


What does this PR do / why do we need it:

新增 Adam/AdamW 优化器状态换入换出能力,降低训练过程中的设备显存占用。

  • 新增 SwapOptimizerConfigswap_optimizeris_swap_optimizer 公共 API。
  • 通过 Platform.get_swap_optimizer() 统一不同后端的优化器包装器分派。
  • 新增后端无关的 PipelineSwapRuntime,支持按状态大小均衡切分、单批次预取及异步 offload。
  • 支持可选的双 staging buffer packed pipeline,复用 A/B 缓冲区并重叠设备间数据传输与优化器更新。
  • 新增 PyTorch 适配器,覆盖原生 Adam/AdamW、融合 AdamW 和 Hyper AdamW。
  • 新增 MindSpore 适配器,覆盖 nn.Adamnn.AdamWeightDecay 和 MindFormers AdamW。
  • 完善 CPU mirror、checkpoint state_dict 保存/加载,以及 fully-shard/DTensor 参数场景。
  • 本次变更共涉及 22 个文件,新增 8948 行、删除 4 行,其中包含 12 个生产代码文件和 10 个测试文件。

Which issue(s) this PR fixes

Fixes #https://gitcode.com/mindspore/hyper-parallel/issues/313


Test Plan and Test result:What scenarios were tested, and what were the verification results(Function, performance, reliability, etc.)

  • tests/ut/core/optimizer/test_swap_optimizer.py:39 个用例、9 个子测试通过。
  • tests/ut/platform/torch/swap_optimizer/test_adapters.py:19 个用例通过。
  • tests/ut/platform/mindspore/swap_optimizer/test_adapters.py:45 个用例通过。
  • 新增 Torch/MindSpore ST,覆盖 Adam/AdamW 多种配置、AMSGrad、多参数组、预初始化状态、checkpoint、fully-shard 和 MindFormers AdamW 等场景。

Self-checklist:(请自检,在[ ]内打上x,我们将检视你的完成情况,否则会导致pr无法合入

likedislike
Pull Request已成功合入, 合并人@MindSpore-Bot
(感谢 songjiaqi 的贡献)
songjiaqisongjiaqi成员
6月29日 创建了 pull request,commit 138ca792
MindSpore-BotMindSpore-Bot成员
6月29日 添加了label:mindspore-cla/yes
司小南(机器人)
司小南(机器人)成员
6月29日 评论:

@silkage_jiajia, 当前/check-pr未通过,原因如下:

以下Pull Request描述检查项未通过:
选项未通过检查: 设计:PR对应的方案是否已经经过Maintainer评审,方案检视意见是否均已答复并完成方案修改
选项未通过检查: 测试:PR中的代码是否已有UT/ST测试用例进行充分的覆盖,新增测试用例是否随本PR一并上库或已经上库
选项未通过检查: 验证:PR描述信息中是否已包含对该PR对应的Feature、Refactor、Bugfix的预期目标达成情况的详细验证结果描述

以下issue检查项未通过:
Pull Request未关联issue

请修改好上述检查错误后,重新使用/check-pr触发检查。

likedislike
司小南(机器人)司小南(机器人)成员
6月29日 添加了label:pr-check-fail
MindSpore-BotMindSpore-Bot成员
6月29日 添加了label:no-pass-all-review
此处折叠了439条消息 查看更多
司小南(机器人)司小南(机器人)成员
8月10日 删除了label:ci-pipeline-passed
司小南(机器人)司小南(机器人)成员
8月10日 添加了label:ci-pipeline-running
司小南(机器人)
司小南(机器人)成员
8月10日 评论:
Project Name Build_Stage Build Result Details
Hyper-parallel_Atomgit_Gate - FAILURE 6680
likedislike
司小南(机器人)司小南(机器人)成员
8月10日 删除了label:ci-pipeline-running
司小南(机器人)司小南(机器人)成员
8月10日 添加了label:ci-pipeline-failed