Enable training cudagraphs for RL - #2452
Conversation
| return type(value)(_clone_nested_tensors(v) for v in value) | ||
| if isinstance(value, dict): | ||
| return {k: _clone_nested_tensors(v) for k, v in value.items()} | ||
| return value |
There was a problem hiding this comment.
Seems a bit dangerous to have Any and not throw an exception for unsupported data types. For instance, if I send a set of tensors, this will not clone them but return a set of tensors which a user of the function would expect to have cloned.
There was a problem hiding this comment.
We actually don't want to throw an exception, but to pass-through for unsupported datatypes. This helper function is scoped very tightly within the cudagraphs code such that the "send set of tensors" case would *not be realistic, in my opinion.
There was a problem hiding this comment.
I did not think of that! Got it now. The only q I have is why if I send a tuple/list of tensors we want to clone them, but if I send a set of tensors we don't clone them?
There was a problem hiding this comment.
Sets of tensors are never used as cudagraph input/output containers — it wouldn't make sense because inputs are ordered, duplicates are meaningful, etc. I would be more concerned if we were sending in a set of tensors at all so I will just throw an error.
This function does see other types (e.g. generators) which need to pass through so I'm not going to modify the typing hints.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
| out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs) | ||
|
|
||
| elif self.training: | ||
| elif self.training and torch.is_grad_enabled(): |
There was a problem hiding this comment.
Not related specifically to this line or meant to be address in this MR, but do we ever need to be able to distinguish between no_grad and inference_mode in the code? If two forwards have the same signature but one is no grad and another is inference mode I wonder if it will error at all in inlace ops
There was a problem hiding this comment.
I've made this unnecessary by flipping to model.eval() and then back to model.train() in logprobs_forward_step so it will never trigger cudagraph capture.
There was a problem hiding this comment.
But yes you're right — right now we rely on cudagraph runner matching guardrails like correct shapes and arguments (e.g. inference passes some arguments which training does not). If we ever have graphs with the exact same shapes and arguments but differing only between inference mode and no_grad mode we may have issues.
… match for cudagraphs
7 participants
What does this PR do ?
Cudagraphs will be enabled for training and inference when
--rl-training-cuda-graphsis passed.Contribution process
flowchart LR A[Pre-checks] --> B[PR Tests] subgraph Code Review/Approval C1[Expert Review] --> C2[Final Review] end B --> C1 C2 --> D[Merge]Pre-checks
Core 0.8)Code review
The following process is enforced via the CODEOWNERS file for changes into
megatron/core. For changes outside ofmegatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.For MRs into `main` branch
(Step 1): Add PR label
Expert Review(Step 2): Collect the expert reviewers reviews
Expert Reviewlabel when your PR is ready for review.Final Review might get declined if these requirements are not fulfilled.
(Step 3): Final Review
Final Reviewlabel(Optional Step 4): Cherry-pick into release branch
If this PR also needs to be merged into
core_r*release branches, after this PR has been merged, selectCherry-pickto open a new PR into the release branch.For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.Merging your PR
Any member of core-adlr and
core-nemowill be able to merge your PR.