[Dev] Enable selective recompute for norm_out in GDN layers

[wplf]

Summary

Port the gdn_norm_out selective recompute (#4715, merged to main) onto the dev branch, adapting it to dev's GatedDeltaNet structure.

Recompute only the gated output norm + HP-to-CP all-to-all block as a discard-output checkpoint, selected via:

model.recompute_granularity=selective
model.recompute_modules=[gdn_norm_out]

The block output (norm_out) is discarded in the forward once out_proj consumes it, and regenerated from a grad hook on the layer output in the backward, freeing the gated-norm activations. Per #4715 this saves ~2–3% per-layer activation memory at ≤2% iteration-time cost.

Adaptations vs the main-branch PR (#4715)

• The gated norm + all-to-all span lives in dev's _forward_compute helper and uses the dynamic cp_size/cp_group plus the _a2a_hp_to_cp() helper and the thd_cp_a2a_inv permutation, so the extracted _gated_norm_and_a2a closure captures those locals.
• The discard-output checkpoint is gated on self.training (matching dev's whole-module gdn recompute convention) so eval/inference runs the block normally.
• Added a validation guard that gdn and gdn_norm_out cannot be combined: the whole-module gdn recompute (dev-only) already wraps this block, so combining them would nest checkpoints.

Test plan

Adds tests/unit_tests/ssm/test_gated_delta_net.py::TestGatedDeltaNet::test_selective_recompute_norm_out, which builds a no-recompute baseline GDN and one with recompute_modules=["gdn_norm_out"], runs forward + backward, and asserts the output, all parameter grads and the input grad match bit-for-bit.

pytest tests/unit_tests/ssm/test_gated_delta_net.py -k test_selective_recompute_norm_out

🤖 Generated with Claude Code

[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.

[wplf]

/ok to test de151f5