Support set recomputed for cell output in jit · Pull Request !91585 · MindSpore/mindspore

1. Support set recomputed for cell output in jit
---
<!--  Thanks for sending a pull request!  Here are some tips for you:

1) If this is your first time, please read our contributor guidelines: https://gitee.com/mindspore/mindspore/blob/master/CONTRIBUTING.md

2) If you want to contribute your code but don't know who will review and merge, please add label `mindspore-assistant` to the pull request, we will find and do it as soon as possible.
-->

**What type of PR is this?**

/kind feature

---

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

重计算是平衡内存占用与计算性能的核心技术。通过在反向传播时重新执行前向算子，可以避免保存大尺寸的中间激活值。
在之前的实现中，静态图重计算功能缺乏对 Cell 输出是否参与重计算的精细控制机制，某些大模型训练场景无法更进一步地减少显存占用。动态图此前已支持配置Cell的输出做重计算，故此 PR 是在静态图下支持该功能。详细设计请查阅issue #IDLK1V

---

**Which issue(s) this PR fixes**:

Fixes #IDLK1V

---

**Test Plan and Test result：What scenarios were tested, and what were the verification results（Function, performance, reliability, etc.）**：

添加了以下三种场景的测试用例，通过检查内存峰值和IR是否符合预期来验证：
1、Jit下执行grad场景的重计算：tests/st/compiler/recompute/test_recompute_with_jit.py
2、Jit下执行grad且带LazyInline场景下重计算：tests/st/compiler/recompute/test_recompute_with_lazy_inline.py
3、GradJit场景下重计算：tests/st/compiler/recompute/test_recompute_under_gradjit.py

---

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

+ - [x] **设计**：PR对应的方案是否已经经过Maintainer评审，方案检视意见是否均已答复并完成方案修改
+ - [x] **测试**：PR中的代码是否已有UT/ST测试用例进行充分的覆盖，新增测试用例是否随本PR一并上库或已经上库
+ - [x] **验证**：PR描述信息中是否已包含对该PR对应的Feature、Refactor、Bugfix的预期目标达成情况的详细验证结果描述
+ - [ ] **接口**：是否涉及对外接口变更，相应变更已得到接口评审组织的通过，API对应的注释信息已经刷新正确
+ - [ ] **文档**：是否涉及官网文档修改，如果涉及请及时提交资料到Doc仓

GVP MindSpore / mindspore
关闭

发现个问题

GVPMindSpore / mindspore关闭

搜索帮助

GVP MindSpore / mindspore
关闭