[dev] Refactor: extract and split common logic between GDN & GDN2 (cherry-pick #5843)

[yuzhongw-nvidia]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Cherry-pick #5843 and part of #5765 into dev.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuzhongw-nvidia]

/claude strict-review

[claude]

Strict Review Summary

Findings by severity — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1

This PR is a cherry-pick of #5843 that splits the 1453-line megatron/core/ssm/gated_delta_net.py into a package (__init__.py, common.py, gdn.py, gdn2.py), extracting a shared _GDNBase for GDN and the new GDN2 variant.

Verification performed

I diffed the new package against the deleted original (origin/dev) along the correctness-sensitive paths:

• Public import surface preserved. __init__.py re-exports GatedDeltaNet, GatedDeltaNetSubmodules, torch_chunk_gated_delta_rule, _build_head_perm_for_split_sections, _build_thd_cp_a2a_perm, tensor_a2a_cp2hp/hp2cp, _split_tensor_factory, etc. All existing importers (hybrid_layer_specs, experimental_attention_variant_module_specs, modelopt specs, tests) resolve unchanged. get_parameter_local_cp_headwise was renamed to get_parameter_local_cp but it is a private helper with no external callers.
• Compute paths are behavior-preserving. torch_chunk_gated_delta_rule, the conv1d/deterministic branch, _prepare_input_for_gated_delta_rule (formerly _prepare_qkv_* + _compute_g_and_beta), the g/beta math (-A_log.exp() * softplus(alpha.float()+dt_bias), fp32), the L2-norm/GQA repeat_interleave, and the CP a2a + zigzag reshuffle all match the original line-for-line.
• Checkpoint compatibility. sharded_state_dict splits (in_proj.weight → query/key/value/z/beta/alpha; conv1d → query/key/value) and TP axis maps are identical, and gate parameters (dt_bias, A_log) are still registered after conv1d, preserving parameter order for existing checkpoints.
• Process-group usage. No new direct parallel_state.get_*_group() reads in library code — groups flow through pg_collection/ProcessGroupCollection, consistent with the CLAUDE.md guidance.
• No unused new identifiers. feat_dim_split, in_proj_extra_dim, dt_bias_dim, a_log_dim all have real read paths (validated by the attrs_to_check assert loop and consumed in split/reset logic). GDN2 (GatedDeltaNet2, chunk_gdn2, torch_chunk_gdn2) is not yet wired into any model spec, but it is a deliberately-introduced public API for a follow-up, exported from __init__.py.

Only finding

One stale-docstring nit: _setup_variant_attrs's docstring lists in_proj_dim as an attribute the variant must set, but the base class derives it and variants set in_proj_extra_dim instead.

Risk: Low. Mechanically faithful refactor with preserved public API, checkpoint layout, and numerics; the added test_deterministic_mode/test_module_construction coverage is a plus.

[yuzhongw-nvidia]

/claude strict-review

[yuzhongw-nvidia]

/ok to test 28a2106

[claude]

Strict Review Summary — PR #6088 (cherry-pick #5843)

Findings by severity: CRITICAL: 0 · IMPORTANT: 3 · SUGGESTION: 1

This is a large, mostly-mechanical refactor that splits the monolithic megatron/core/ssm/gated_delta_net.py into a package (_GDNBase in common.py, GatedDeltaNet in gdn.py, new GatedDeltaNet2 in gdn2.py) while preserving the public import surface via __init__.py. I traced the GDN v1 forward/CP/checkpoint/sharded-state-dict paths against the deleted original and they are faithfully preserved (in_proj sizing, conv1d TP/CP splits, headwise vs chunkwise A2A plumbing, _build_thd_cp_a2a_perm, _build_head_perm_for_split_sections, and the torch-native deterministic kernel body are all unchanged). Public re-exports and the external importers (experimental_attention_variant_module_specs.py, hybrid_layer_specs.py, post_training/.../model_specs.py, tests) remain valid. Removed test imports were genuinely unused.

The new surface (the GatedDeltaNet2 variant, its torch kernel, and the recompute_norm_out path) is where the findings cluster:

Most impactful

1. recompute_norm_out is unreachable (IMPORTANT). It keys off "gdn_norm_out" in recompute_modules, but that token is not in the allowed_modules set validated in TransformerConfig.__post_init__ — enabling it fails config validation. This PR doesn't touch transformer_config.py, so all the new norm-recompute machinery in both variants is dead until that allowance lands.
2. Missing None guard (IMPORTANT). recompute_norm_out reads "gdn_norm_out" in self.config.recompute_modules without the truthiness guard used one line below for recompute_gdn; recompute_modules defaults to None, risking a TypeError on config paths that bypass normalization.
3. Lost defensive assert (IMPORTANT). The torch-native kernels dropped the explicit assert cp_context is None, now absorbing cp_context via **kwargs. Harmless today (chunkwise-CP + deterministic isn't exercised), but a future misuse would silently produce wrong results instead of erroring.

Note

• GDN2 has no test/spec coverage in this PR. GatedDeltaNet2, chunk_gdn2, and torch_chunk_gdn2 are net-new and not referenced by any module spec or unit test in the repo. The added test_deterministic_mode covers GDN v1 only. Given this is a cherry-pick, that may be intentional (coverage arriving separately), but the numerically intricate torch_chunk_gdn2 (centered-decay WY recurrence) shipping untested is a risk worth flagging.

Overall risk: Low–Moderate

The GDN v1 refactor is behavior-preserving and low-risk. The moderate items are all confined to newly-added, currently-inactive code paths (norm recompute is unreachable; GDN2 is unwired). None blocks the refactor's core intent, but the recompute_norm_out config gap and the untested GDN2 kernel should be resolved before those features are relied upon.

[yuzhongw-nvidia]

/ok to test 1aa8afe