Skip to content

fix(transformer): AUT-1497 support dual RoPE full recompute - #6519

Merged
ko3n1g merged 1 commit into
NVIDIA:devfrom
svcnemo-autobot:fix/gemma4-dual-rope-recompute-dev-260813
Aug 14, 2026
Merged

fix(transformer): AUT-1497 support dual RoPE full recompute#6519
ko3n1g merged 1 commit into
NVIDIA:devfrom
svcnemo-autobot:fix/gemma4-dual-rope-recompute-dev-260813

Conversation

@svcnemo-autobot

Copy link
Copy Markdown
Collaborator

Background

Bridge Gemma4-VL full recompute fails on the selected MCore dev pin because a dual-RoPE tuple reaches autograd checkpoint storage. MCore main already carries the equivalent repair from #5324.

What changed

  • Split dual RoPE into tensor checkpoint arguments and reconstruct it inside recompute.
  • Preserve the non-checkpointed block path.
  • Add focused tensor, dual-RoPE, local-only, and global-only coverage.

Details

flowchart LR
  R[dual RoPE tuple] --> S[split checkpoint args]
  S --> C[autograd checkpoint]
  C --> J[reconstruct RoPE]
Loading

This PR is based directly on dev pin 12eed0a5ac3f1ced6a3a4f76727db530bdbf7391 for coupled Bridge validation.

Tested

  • H100 full-checkpoint tests: 5 passed
  • Black, isort, Ruff, and Pylint
  • Coupled Bridge Gemma4-VL rerun pending

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
@svcnemo-autobot
svcnemo-autobot requested review from a team as code owners August 13, 2026 18:29
@copy-pr-bot

copy-pr-bot Bot commented Aug 13, 2026

Copy link
Copy Markdown

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.

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

/ok to test b3b1e2d

@ko3n1g
ko3n1g enabled auto-merge August 13, 2026 22:18
@ko3n1g
ko3n1g added this pull request to the merge queue Aug 14, 2026
Merged via the queue into NVIDIA:dev with commit b8b0a41 Aug 14, 2026
236 of 256 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants