test(char-causal-lm): cover gradient checkpointing delegation

[AmitMY]

Audit item #7 — but tests first, not the deletion.

gradient_checkpointing_enable / gradient_checkpointing_disable on CharacterCausalLMWrapper had zero test coverage, so there was no way to tell whether toggling on the wrapper actually reached the wrapped model. Four tests now pin that down:

• test_enable_reaches_wrapped_model — flag off, enable, flag on
• test_disable_reaches_wrapped_model — round trip back off
• test_enable_forwards_kwargs — the gradient_checkpointing_kwargs={"use_reentrant": False} form
• test_backward_still_produces_gradients — forward + backward with checkpointing on, asserting gradients actually arrive once activations are being recomputed

The tiny_wrapper fixture builds a 2-layer LlamaConfig in code, so these stay in the fast unit path and download nothing — unlike the existing integration classes, which pull sign/utf8-lm-tiny from the hub behind @pytest.mark.slow.

245 passed on transformers 4.57.6 and 5.14.1 both. ruff clean. Test-only change, no library code touched.

What this now tells us about the overrides

With coverage in place I could answer the original audit question, and the answer is that both overrides are redundant. Deleting them and letting PreTrainedModel's own implementation walk submodules produces identical results on both versions:

	with override	base implementation
transformers 5.14.1	enable→True, disable→False	enable→True, disable→False
transformers 4.57.6	enable→True, disable→False	enable→True, disable→False

The kwargs form works through the base implementation too.

I have not deleted them here — this PR is the safety net, and the deletion belongs in its own commit now that it can be verified rather than assumed. Happy to send it as a follow-up.

🤖 Generated with Claude Code

---

Note

Low Risk
Test-only additions with no production code changes.

Overview
Adds fast unit tests so CharacterCausalLMWrapper’s gradient_checkpointing_enable / gradient_checkpointing_disable are verified to affect the inner causal LM, not only the wrapper.

A tiny_wrapper fixture builds a 2-layer LlamaConfig model in-process (no Hub download). TestGradientCheckpointing checks enable/disable round-trip, gradient_checkpointing_kwargs forwarding, and that forward + backward still produce gradients on char_embedding when checkpointing is on.

Test-only change; no library code in this diff.

^{Reviewed by Cursor Bugbot for commit 73ce446. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

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

^{Reviewed by Cursor Bugbot for commit 73ce446. Configure here.}

[cursor]

Kwargs forwarding never asserted

Low Severity

test_enable_forwards_kwargs only checks that checkpointing is enabled after the call, same as the plain enable test. It never inspects whether gradient_checkpointing_kwargs reached the wrapped model, so dropping kwargs in the override would still pass this safety-net test.

 

^{Reviewed by Cursor Bugbot for commit 73ce446. Configure here.}

[cursor]

Grad check skips checkpointed model

Low Severity

test_backward_still_produces_gradients asserts grads on char_embedding.embedding.bit_proj_w, but that parameter gets gradients from the decode matmul without flowing through the wrapped, checkpointed model. A broken checkpointed backward can still leave this assertion green.

 

^{Reviewed by Cursor Bugbot for commit 73ce446. Configure here.}