microsoft / microsoft/onnxruntime
[Feature Request] [Training] Add Scan and other control flow operators
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the feature request
I want training support for mainly Scan, and possibly other control flow (If and Loop). Because Scan hold a subgraph to compute, we cannot use existing way to implement gradient for ops described in [here](https://github.com/microsoft/onnxruntime/wiki/Per-Op-Gradient-Building).
Implementing gradient support for Scan enables on-device training for RNN-based model (like RWKV, or Mamba) without custom kernel, so I think it will be great for less popular devices/accelerators (if ORT training support varieties of EP).
I think GetExternalGradient does gradient calculation by using custom ops and is considered as one of the solutions, but still documentations are not enough. (Sorry if it's wrong)
### Describe scenario use case
As mentioned above, on-device training without custom kernel, or on-device training with complex operations that cannot be expressed with CumSum/CumProd and other operators.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the linked Per-Op Gradient Building wiki page and investigating GetExternalGradient, which the issue identifies as a possible approach. Clarify the intended training support for Scan and whether If and Loop are also in scope; done should include a defined implementation path and documented behavior for the requested control-flow operators.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100