已合并
trainer_recompute #1132
已合并
songjiaqi创建于 8月7日
songjiaqi
songjiaqi成员
8月7日

What type of PR is this?

/kind feature


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

为 Trainer 增加统一的 Activation Checkpointing 配置和执行流程,支持 off、full、selective 三种模式,降低大模型训练的激活显存占用。

主要变更:

  • 新增分布式 Activation Checkpointing 实现,支持:
    • Full activation checkpointing
    • Selective activation checkpointing
      计算密集型算子、通信算子及可选扩展算子的策略配置
  • 在模型配置、from_pretrained、模型构建和 Trainer 流程中透传 activation_checkpoint 配置。
  • 按照 sharding plan -> activation checkpointing -> torch.compile -> FSDP2 的顺序应用相关处理,保证 checkpoint wrapper 位于 TP/EP 转换之后及 FSDP 边界之内。
  • 支持已知 Transformer/VLM 模型的语言层和视觉层容器识别,并对未知模型提供兼容性回退。
  • 对 Hugging Face 原生 gradient checkpointing、KV sharing 场景及 compile 场景进行兼容处理。
  • 为 MindSpore 和 PyTorch SAC 增加运行时忽略算子注册能力。
  • 更新训练示例配置,将旧的 gradient_checkpointing 配置迁移为 activation_checkpoint.mode。

Which issue(s) this PR fixes:

Fixes #


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

新增并更新相关单元测试,覆盖:

  • off、full、selective 配置解析及禁用行为。
  • Full/Selective checkpoint wrapper 的创建和反向重计算。
  • Selective checkpoint 对 TopK 原地修改输出的处理。
  • SAC profiler/FSDP runtime operator 注册。
  • KV sharing 场景下的 checkpoint 回退逻辑。
  • 已知 VLM 模型的语言层/视觉层容器发现及去重。
  • 未知模型的兼容性容器发现。
  • Hugging Face 原生 gradient checkpointing 及 compile 场景。
  • Activation Checkpointing 与 sharding、FSDP2 的执行顺序。
  • MindSpore/PyTorch SAC 忽略算子逻辑。
  • Trainer 模型阶段对 Activation Checkpointing 配置的透传。

本次提交包含测试用例更新,但当前未在本地执行完整测试套件,建议在 CI 中完成 PyTorch、MindSpore 及分布式场景验证。


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

likedislike
Pull Request已成功合入, 合并人@MindSpore-Bot
(感谢 songjiaqi 的贡献)
songjiaqisongjiaqi成员
8月7日 创建了 pull request,commit 5af47489
songjiaqisongjiaqi成员
8月7日 virtual merging failed, update merge request[project_id: 8714117, iid: 1132, target_commit_sha: 5af47489ecd310db201b64f63021b46b4db17df2], message: Conflict detected
songjiaqisongjiaqi成员
8月7日 审查状态已重置,审查人: changzherui1,yao_yf
songjiaqisongjiaqi成员
8月7日 强制推送  1 个提交:4ddaf7ae-trainer_recompute
songjiaqisongjiaqi成员
8月7日 virtual merging failed, update merge request[project_id: 8714117, iid: 1132, target_commit_sha: 5af47489ecd310db201b64f63021b46b4db17df2], message: Conflict detected
此处折叠了92条消息 查看更多
司小南(机器人)司小南(机器人)成员
8月14日 添加了label:ci-pipeline-running
司小南(机器人)
司小南(机器人)成员
8月14日 评论:
Project Name Build_Stage Build Result Details
Hyper-parallel_Atomgit_Gate - ✅ SUCCESS 6959
likedislike
司小南(机器人)司小南(机器人)成员
8月14日 删除了label:ci-pipeline-running
司小南(机器人)司小南(机器人)成员
8月14日 添加了label:ci-pipeline-passed
MindSpore-BotMindSpore-Bot成员
8月14日 合入了pull request,合并节点 SHA:be4c89aec9b5f06a47a9cc91b31dbd8651c42016