Preserve gradient checkpointing kwargs during GKD generation

[DimensionSTP]

What does this PR do?

GKD temporarily disables gradient checkpointing while unwrapping the student or teacher model for generation. The previous restoration path called gradient_checkpointing_enable() without the configured gradient_checkpointing_kwargs.

With reentrant checkpointing under DeepSpeed ZeRO-3, this changes checkpointing behavior between the original forward and backward recomputation. The recomputation can then observe partitioned zero-sized parameters and fail with:

torch.utils.checkpoint.CheckpointError:
Recomputed values have different metadata than during the forward pass.
saved shape: [8192, 2560]
recomputed shape: [0]

This change:

• lets the generation unwrap context receive the original checkpointing kwargs;
• restores checkpointing through the existing disable_gradient_checkpointing context manager, including exception exits;
• passes self.args.gradient_checkpointing_kwargs from both GKD generation paths;
• adds focused tests for normal and exceptional restoration.

The GKD objective, generation behavior, and DeepSpeed gathering policy are unchanged.

Verification

• Two-rank DeepSpeed ZeRO-3 reproducer before the fix: deterministic checkpoint metadata mismatch
• Same reproducer after the fix: one optimizer step completed
• pytest -q tests/test_model_utils.py: 7 passed
• Focused restoration tests: 2 passed
• Ruff check for the changed tests: passed
• Ruff format check for all three changed files: passed
• compileall: passed
• git diff --check: passed

Before submitting

• Did you read the contributor guidelines?
• Was this discussed/approved via a GitHub issue? Not applicable; this is a self-contained bug fix with a minimal reproducer.
• Did you make sure to update the documentation? No user-facing documentation change is required.
• Did you write the necessary tests?

AI writing disclosure

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

---

Note

^{Cursor Bugbot is generating a summary for commit 2369cb4. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 2369cb4. Configure here.}

[cursor]

Checkpoint restore breaks kwargs-only wrappers

High Severity

Generation unwrap now restores checkpointing through disable_gradient_checkpointing, which calls gradient_checkpointing_enable with a positional dict. Wrappers whose method is keyword-only, including PPO's PolicyAndValueWrapper, raise TypeError on exit whenever checkpointing was enabled. Previously unwrap restored with a no-arg enable() call, so PPO generation with checkpointing worked.

Additional Locations (1)

• trl/models/utils.py#L402-L404

 

^{Reviewed by Cursor Bugbot for commit 2369cb4. Configure here.}