[mcore] Cuda Graph support in training and logprob
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Is your feature request related to a problem? Please describe.**
Megatron core natively supports cuda graph for transformer layers, but in RL the complication is that we can switching between train, logprob, and generation, and we may need to offload/onload data between stages which can cause IMA for cuda graph. In addition, sequence padding might be needed for static shape.
This issue tracks the support of cuda graph in training phases (train and logprob) of GRPO.
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.
**Additional context**
Add any other context or screenshots about the feature request here.
Contributor guide
Assessment
This issue has not been assessed yet.