feat(gdn): Pass GDN kernel backend to Mcore - #3995
vasunvidia wants to merge 1 commit into
Conversation
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
9f18ce0 to
c6daa54
Compare
vasunvidia
changed the title
jepio
added
the
CI:Lfast
jepio
left a comment
jepio
left a comment
There was a problem hiding this comment.
Reviewed by a team of specialized agents (rl-expert, bug-finder, test-agent, design-reviewer, devil's-advocate).
One finding below, refined with the author's context: this PR intentionally lands the gdn_kernel_backend passthrough ahead of a scheduled Megatron-core pin bump, to avoid coupling this change to that bump's stricter testing/schedule.
Generated by Claude Code
| # so the canonical expanded-sequence contract must use backend dispatch. | ||
| attention_backend = config["megatron_cfg"].get("attention_backend") | ||
| if "gdn_kernel_backend" in config["megatron_cfg"]: | ||
| model_cfg.gdn_kernel_backend = config["megatron_cfg"]["gdn_kernel_backend"] |
There was a problem hiding this comment.
nemo_rl/models/megatron/setup.py:1124-1125
1 action item.
model_cfg.gdn_kernel_backend = ... sets an attribute that doesn't exist on TransformerConfig at the currently pinned Megatron-LM SHA (731b7914) — confirmed via grep across both submodules (zero other hits) and reading the GDN dispatch logic (gdn.py#L57-L60, gdn2.py#L98-L101), which is driven solely by deterministic_mode, not a kernel-backend field. Since TransformerConfig is a plain, non-slotted dataclass, this succeeds silently rather than erroring.
This is presumably intentional prep for a separate, upcoming Megatron-core pin bump that introduces the real field — landing the plumbing early avoids coupling this PR to that bump's stricter testing/schedule, which is reasonable. But nothing in the PR currently signals that sequencing to a reviewer.
Action: Update the PR description to state this is staging for the upcoming Megatron-core pin bump (link the tracking issue/PR if one exists) and that gdn_kernel_backend is inert until then. Once the pin bump lands, follow up with: gdn_kernel_backend: NotRequired[str] on MegatronConfig (nemo_rl/models/policy/__init__.py:329-380), exemplar YAML docs, and a present/omitted unit test pair mirroring the existing cuda_graph_modules tests in tests/unit/models/megatron/test_megatron_setup.py.
@jepio, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
|
/ok to test c6daa54 |
|
This is not required for MLPerf because #3535 is merged and can be used to enable new features. |
3 participants
What does this PR do ?
Add a one line overview of what this PR aims to accomplish.
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information