[DiffusionGemma] Support gradient checkpointing (#46572)
* Support gradient checkpointing in DiffusionGemma
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Test that checkpointing reproduces the exact loss and gradients
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* update use_cache test now that the decoder ignores it instead of raising
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* use a single read-only kv cache under gradient checkpointing
* drop the unused encoder kv return and guard the decoder cache
* reuse BaseModelOutputWithPast and drop the unrelated mask cast
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
