[Medium] Fix gradient checkpointing incompatibility with intervention hooks (#231)

[pablocs116]

Description

Closes #231.

Training a PyReFT model with gradient checkpointing (use_reentrant=False) fails
in the backward pass with CheckpointError: A different number of tensors was saved during the original forward and recomputation.

IntervenableModel.forward() removed the intervention hooks and cleared their
cached state immediately after the intervened forward — i.e. before the
Trainer calls .backward(). With use_reentrant=False, PyTorch recomputes each
module's forward during backward; by then the hooks are gone, so the recompute
saves a different number of tensors than the original forward and the checkpoint
guard raises.

Fix: defer the hook teardown (removing the handles + _cleanup_states) to a
post-backward callback whenever gradients are being tracked, so the hooks stay
registered during recomputation and are cleaned up once the whole backward graph
has run. The teardown is scheduled via a hook on the grad-requiring output
tensors that queues the real cleanup on the autograd engine (one-shot guarded).
When there is nothing to back-propagate through, teardown happens immediately —
no behavior change off the training path.

Testing Done

Added tests/integration_tests/GradientCheckpointingTestCase.py (3 tests):

• forward+backward under use_reentrant=False no longer raises;
• loss and intervention gradients match a non-checkpointed run exactly
  (checkpointing must not change numerics);
• hooks and cached state are fully cleared after backward (no leak across steps).

Confirmed the new regression test fails on main with the reported
CheckpointError and passes with this change:

# on main (before fix):
tests/integration_tests/GradientCheckpointingTestCase.py .FF   -> CheckpointError
# with fix:
tests/integration_tests/GradientCheckpointingTestCase.py ...   3 passed

Full existing suite (python -m unittest discover -p '*TestCase.py'):
Ran 75 tests — 74 pass. The single unrelated error
(test_customized_intervention_function_get, a KeyError collecting attention
weights) also fails on a clean main with current transformers and is not
touched by this PR.

Checklist:

• My PR title strictly follows the format: [Your Priority] Your Title
• I have attached the testing log above
• I provide enough comments to my code
• I have changed documentations
• I have added tests for my changes

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 071627946c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".