fix bug where Gemma4 is not working with recompute_granularity = "full"

[warpuv]

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

What does this PR do ?

fix bug where Gemma4 is not working with recompute_granularity = "full"

Issue tracking

Linked issue: #5323

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

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

[github-actions]

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:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[guihong-nv]

@warpuv Thanks for the quick fix. I think the checkpoint arg-splitting is the right direction, but I found one correctness issue that should be fixed before merge.

The reconstruction in checkpointed_forward() currently treats rotary_pos_emb_local is None as meaning “this was a non-dual RoPE input”:

rotary_pos_emb = (
    rotary_pos_emb_global
    if rotary_pos_emb_local is None
    else (rotary_pos_emb_local, rotary_pos_emb_global)
)

That changes semantics for a valid asymmetric dual-RoPE tuple like (None, k_pos_emb): it gets collapsed to k_pos_emb, and then SelfAttention will duplicate non-tuples into (k_pos_emb, k_pos_emb), applying query RoPE even though the caller explicitly disabled it.

Could we preserve whether the original input was dual RoPE and rebuild the tuple whenever it was originally tuple/list, even if the local element is None?

I’d also strongly suggest adding a focused regression test for recompute_granularity="full" with tuple rotary_pos_emb, including an asymmetric case like (None, k_pos_emb).

[warpuv]

Dear @guihong-nv, thank you for your review and valuable suggestions. I fixed the code and added test.

[guihong-nv]

/ok to test 2473b86

[guihong-nv]

LGTM

[Connor-XY]

/claude review

[claude]

LGTM

[warpuv]

@mcore-oncall

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29596792106