[losewayy]

Description

Checkpointing saves graph-safe generator handles before forward. Those handles share mutable CUDA RNG state, so forward advances the offsets that recompute later reads. Dropout can then produce a different mask and incorrect input gradients.

Save the seed and offset separately from the handles, rewind the original states in place for recompute, and restore the live values on exit. This preserves shared RNG streams and the state already registered with CUDA graphs. The captured dropout recompute mismatch reported in #3834 remains.

Fixes #3833

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Share RNG save and restore between reentrant and non-reentrant checkpointing.
• Cover default, tracked, and shared RNG streams, and replay a graph captured before an eager checkpoint.

On an RTX 5070 Ti Laptop GPU with PyTorch 2.14.1+cu130, tracked-dropout gradient mismatches fell from 2017 of 4096 elements to zero. The checkpoint test selection passed all 33 cases, including eight regression cases. An explicit generator registration probe also passed capture and replay on PyTorch 2.11.0+cu128. The graph regression exercises the registration API that PyTorch 2.14 reports as deprecated.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge, with no new actionable issues found.

Summary

Shares RNG save and restore code between both checkpoint modes while preserving the original generator handles.

• Checkpoint recompute now reuses the random draws from forward.

_{Reviews (2) · Last reviewed commit: "[PyTorch] Test CUDA RNG restoration afte..." · Reviewed by Greptile}