NVIDIA-NeMo / NVIDIA-NeMo/RL

[mcore] Cuda Graph support in training and logprob

Open
#1,246 0 comments 0 reactions 1 assignee Claimed by @seonjinn View on GitHub
Performance research t-mcore
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.