[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 independent RNG values alongside the original handles, rewind those states in place for recompute, and restore live values on exit. Eager snapshots stay on CPU. During capture, save seed and base offset on GPU via PyTorch's public Generator.philox_state API and save the intragraph counter separately in the CPU state returned by Generator.get_state. Recompute restores the GPU values and rewinds the CPU counter with Generator.set_state, then restores both live states on exit; no offset addition or subtraction is needed. Both TE checkpoint engines remain in place; their RNG state objects own the saved snapshots directly. No C++ or CUDA changes remain in the final diff.

Builds without Generator.philox_state (available in PyTorch 2.14+) reject captured seeded PyTorch operations and native TE dropout/attention inside checkpoint regions with a clear error. The missing-API guard does not cover native NVFP4 stochastic rounding. Eager dropout and deterministic captured checkpointing remain supported on older builds.

Fixes #3833
Fixes #3834

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

• Preserve RNG values, generator identities, aliases, and existing graph registrations in both TE checkpoint modes.
• Keep RNG snapshots on the checkpoint RNG state object, without a separate autograd storage function.
• Reject unsupported captured seeded PyTorch operations and native TE dropout/attention; native checks use the checkpoint-region flag and capture state rather than inspecting the dispatch-mode stack. Restore graph wrappers and recompute context on errors.
• Remove RNG compatibility paths for unsupported PyTorch versions while retaining tensor-based downstream tracker states.
• Cover default, tracked, mixed, shared, and no-tracker streams, graph replay, and nested checkpoints.

Validated with regular python -m pytest on RTX 5880 Ada (PyTorch 2.15 nightly, CUDA 13.4, pytest 8.2.1), importing the complete test module and TE Python sources from the checkout: tests/pytorch/test_numerics.py -k "checkpoint and not gpt" passed all 46 selected tests (1113 deselected, no skips). This includes both checkpoint modes, separate forward/backward graphs, manual capture of forward and backward in one graph. Core and PyTorch native libraries were rebuilt from the merged sources (main ab4b9222) before this run; the PR adds no C++ changes. Four existing ops-registration cases also passed (50 tests in the combined selection). The Python L0 lint launcher, pre-commit hooks, formatting and diff checks passed. Full-suite and older-PyTorch runtime testing were not run.

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 two non-blocking concerns from earlier reviews.

Findings

1.  RNG buffers outlive backward ▶
2.  Older builds lose test coverage ▶

Summary

The PR saves independent RNG values so checkpoint recompute can repeat forward draws without changing live generator states.

• Checkpoint recompute reuses the forward pass’s CUDA random values.
• Captured checkpoint dropout gets a clear error on older PyTorch builds.
• Checkpoint and graph setup restores temporary state after errors.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Save RNG values and handles] --> B[Run checkpoint forward]
  B --> C[Save live RNG values before recompute]
  C --> D[Restore forward RNG values]
  D --> E[Recompute forward]
  E --> F[Restore live RNG values and handles]

Loading

_{Reviews (17) · Last reviewed commit: "Revert "Test the checkpoint RNG guard wi..." · Reviewed by Greptile}

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[greptile-apps]

RNG buffers outlive backward

RNG snapshots now stay in self.saved_states instead of autograd's saved tensors. In reentrant mode, keeping a checkpoint output after backward also keeps these snapshots alive, even though backward has released the saved inputs. Each checkpoint retains a CPU RNG buffer and one buffer per CUDA generator.

This is a non-blocking memory cleanup concern. Keep the snapshots in save_for_backward so backward can release them and saved-tensor hooks can manage their storage.

[greptile-apps]

Older builds lose test coverage

Removing test_checkpoint_graph_rng_without_philox_api and test_checkpoint_native_attention_requires_philox_api removes the only tests of capture without philox_state. The remaining graph-replay test skips builds without that API, so regressions in the guards or cleanup after rejected capture can pass unnoticed.

This is non-blocking, but please keep these tests or replace them. They check that random operations raise, deterministic capture works, and rejected capture restores the changed state.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[pggPL]

TE CI — selected tests

Pipeline ID: 82005005 · Commit: 5ea84ec6927eff420d8900aee9c803c03e2be17c

Selected CI jobs and pytest filters are listed below.

Frameworks: PyTorch

Mode: normal CI build.

Selected jobs: 6

• L0_pytorch_debug_unittest--B200_1GPU
• L0_pytorch_debug_unittest--H100_1GPU
• L0_pytorch_lint
• L0_pytorch_unittest--A100_1GPU
• L0_pytorch_unittest--B200_1GPU
• L0_pytorch_unittest--H100_1GPU

Not run (manual/skipped): 9

Show jobs not run

• L1_pytorch_distributed_unittest--B200_8GPU
• L1_pytorch_distributed_unittest--H100_4GPU
• L1_pytorch_mcore_fsdp_integration--B200_8GPU
• L1_pytorch_mcore_integration--B200_1GPU
• L1_pytorch_mcore_integration--H100_1GPU
• L1_pytorch_onnx_unittest--B200_1GPU
• L1_pytorch_onnx_unittest--H100_1GPU
• L3_pytorch_FA_versions_test--B200_8GPU
• L3_pytorch_FA_versions_test--H100_4GPU

Results are reported as individual commit statuses on this PR.