fix: Disable cudnn sdpa backend when using activation checkpointing

[yfw]

What does this PR do ?

The previous fix (#1676) did not cover all the cases. Specifically, it was failing for certain configurations where the cudnn backend still seemed to be selected during activation recomputation. This is because the context manager for setting the sdpa backend was only activated in the forward call, not during activation recomputation. This PR disables cudnn backend globally when activation checkpointing is enabled for the dtensor path to avoid this.

Related: #1663

This is likely related to this change in the cudnn backend: https://github.com/pytorch/pytorch/pull/155958/files#diff-0af86060a6f34f46e562971d76a9ad8ddaeb945c8fbd6693186f1d60304de438L263

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

Summary by CodeRabbit

• Bug Fixes

  • Improved activation checkpointing reliability by ensuring consistent CUDA backend behavior during model training.

• Refactor

  • Simplified internal configuration logic for activation checkpointing and CUDA optimization coordination.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[github-actions]

⚠️ File Consistency Check

Check based on commit: 9b386a4 (PR #1717 from yifu/fix_act2)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Modified dtensor_policy_worker_v2.py to refine activation checkpointing and SDPA backend handling. Changed context-parallel activation checkpointing condition, removed SDPA exclusion logic, and added explicit runtime disablement of CUDNN SDPA backend when activation checkpointing is enabled.

Changes

Cohort / File(s)	Summary
SDPA and Activation Checkpointing Logic nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py	Added CUDA backend utilities import. Narrowed CP-related activation checkpointing condition to trigger only when cp_size > 1 (excluding OR clause). Removed conditional block excluding activation_checkpointing in SDPA init setup. Introduced explicit runtime disablement of CUDNN SDPA backend when activation checkpointing is enabled post-model-load.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• fix: Fix crash when using activation_checkpointing #1676: Directly modifies the same dtensor_policy_worker_v2.py file with related CUDNN SDPA/attention handling when activation_checkpointing is enabled.
• cp: fix: Fix crash when using activation_checkpointing (1676) into r0.5.0 #1682: Modifies identical activation_checkpointing and SDPA/CUDNN attention logic in the same file with overlapping backend treatment changes.

Suggested labels

CI:L1, Run CICD

Suggested reviewers

• terrykong
• joyang-nv

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR implements cuDNN SDPA backend disablement for activation checkpointing but lacks validation showing the fix resolves the recomputation error and doesn't cause convergence regressions.	Add test results confirming the fix resolves the recomputation metadata error, validates no convergence regressions occur, and documents performance impact of disabling cuDNN SDPA backend.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check	✅ Passed	The title accurately summarizes the main change: disabling the cudnn SDPA backend when activation checkpointing is enabled, which is the core fix described in the PR objectives.

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

⚠️ File Consistency Check

Check based on commit: c91fb39 (PR #1717 from yifu/fix_act2)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}