refactor: extract and split common logic between GDN & GDN2 - #5843
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
/ok to test c2cbf04 |
xuantengh
changed the title
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
LGTM — clean refactor extracting shared GDN logic into gdn/common.py (_GDNBase) with a back-compat shim. Parameter registration order, split tables, CP all-to-all plumbing, and sharded checkpointing are preserved, and new tests (test_module_construction, test_gpu_forward_backward) plus the updated test_jit_compiled_helpers cover the new API. One trivial docstring typo noted inline.
|
/ok to test ba7e925 |
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
|
/ok to test 44bc4d3 |
svcnvidia-nemo-ci
added
Approved
|
/ok to test 112fc4a |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30286617586 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30293553266 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30329827861 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30337954969 |
What does this PR do?
This is the first of several PRs to refactor the GDN module in order to support different variants of gated delta networks (GDN / GDN2) and different backend implementations (FLA, FlashInfer, FlashQLA, etc.). This PR extracts the parameter initialization logic into
_GDNBaseincommon.pyso that any GDN variant can inherit this base class to avoid code duplication. Also, it also keeps the originalgated_delta_net.pyfile to hold back compatibility.This wandb link shows the refactor produces exactly the same output (lm loss).
#5392 will be the next step to abstract away the CP communication logic.