From 3816a260d458300e1d5a4edcc6e8344fbddcdf74 Mon Sep 17 00:00:00 2001
From: losewayy <104015127+losewayy@users.noreply.github.com>
Date: Sun, 8 Nov 2026 20:00:16 +0800
Subject: [PATCH 01/23] [PyTorch] Preserve CUDA RNG state during checkpoint
 recompute

Graph-safe generator handles share state with the active CUDA generator.
Saving only the handles lets forward advance the saved RNG position, so
dropout recomputation can produce a different mask and incorrect input
gradients.

Save seed and offset values separately and rewind the original states
only while recomputing. Restore the live values on exit so aliases and
CUDA graph registrations stay valid. During capture, keep using the
graph-safe handles. Captured dropout recomputation remains a separate
limitation.

Cover both checkpoint modes with default, tracked, and shared RNG
streams, and verify that an existing graph can replay after an eager
checkpoint.

Signed-off-by: losewayy <104015127+losewayy@users.noreply.github.com>
---
 tests/pytorch/test_numerics.py            |  96 +++++++++++++
 transformer_engine/pytorch/distributed.py | 164 ++++++++++------------
 2 files changed, 171 insertions(+), 89 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 395476dbd1d..12373858382 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -10,6 +10,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch.nn import Parameter
 
 from transformer_engine.pytorch.quantization import (
@@ -45,6 +46,8 @@
 )
 from transformer_engine.pytorch import checkpoint as te_checkpoint
 from transformer_engine.pytorch.distributed import (
+    CudaRNGStatesTracker,
+    graph_safe_rng_available,
     is_fp8_activation_recompute_enabled,
     in_fp8_activation_recompute_phase,
 )
@@ -1132,6 +1135,99 @@ def outer_body(value):
     _assert_fp8_recompute_state_drained(inner, tail)
 
 
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+@pytest.mark.parametrize("rng_source", ("default", "tracker", "shared"))
+def test_checkpoint_recompute_preserves_tracker_rng_state(use_reentrant, rng_source):
+    """Recomputed dropout must match forward for default, tracked, and shared RNG."""
+
+    def run(checkpointed):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        tracker = CudaRNGStatesTracker()
+        tracker.add("model-parallel-rng", 4321)
+        tracker_state = tracker.get_states()["model-parallel-rng"]
+        inp = torch.ones(4096, device="cuda", requires_grad=True)
+        draws = []
+
+        def body(value):
+            if rng_source == "tracker":
+                with tracker.fork():
+                    out = F.dropout(value, p=0.5, training=True)
+            else:
+                out = F.dropout(value, p=0.5, training=True)
+                if rng_source == "shared":
+                    with tracker.fork():
+                        out = out + 2 * F.dropout(value, p=0.5, training=True)
+            draws.append(out.detach().clone())
+            return out
+
+        # Leave the fork before backward so the checkpoint must restore the
+        # sharing between the default generator and the tracker for recompute.
+        with tracker.fork() if rng_source == "shared" else nullcontext():
+            if checkpointed:
+                out = te_checkpoint(
+                    body,
+                    inp,
+                    use_reentrant=use_reentrant,
+                    get_rng_state_tracker=lambda: tracker,
+                    distribute_saved_activations=False,
+                    tp_group=None,
+                )
+            else:
+                out = body(inp)
+        out.sum().backward()
+        return out.detach(), inp.grad.detach(), draws, tracker, tracker_state
+
+    ref_out, ref_grad, _, _, _ = run(checkpointed=False)
+    out, grad, draws, tracker, tracker_state = run(checkpointed=True)
+    assert torch.equal(out, ref_out)
+    assert len(draws) == 2
+    assert torch.equal(draws[0], draws[1])
+    assert torch.equal(grad, ref_grad)
+    if isinstance(tracker_state, torch.Generator):
+        assert tracker.get_states()["model-parallel-rng"] is tracker_state
+
+
+@pytest.mark.skipif(not graph_safe_rng_available(), reason="Requires graph-safe CUDA RNG")
+@pytest.mark.parametrize("fork_tracker", all_boolean)
+def test_checkpoint_preserves_captured_rng_state(fork_tracker):
+    """Graph replay must still advance its registered RNG after eager recompute."""
+    tracker = CudaRNGStatesTracker()
+    tracker.add("model-parallel-rng", 4321)
+    rng_context = tracker.fork if fork_tracker else nullcontext
+    generator = (
+        tracker.get_states()["model-parallel-rng"]
+        if fork_tracker
+        else torch.cuda.default_generators[torch.cuda.current_device()]
+    )
+    sample = torch.empty(128, device="cuda")
+    with rng_context():
+        sample.uniform_()
+    graph = torch.cuda.CUDAGraph()
+    if fork_tracker:
+        graph.register_generator_state(generator)
+    with torch.cuda.graph(graph), rng_context():
+        sample.uniform_()
+
+    offset = generator.get_offset()
+    graph.replay()
+    increment = generator.get_offset() - offset
+    assert increment > 0
+
+    def body(value):
+        with rng_context():
+            return F.dropout(value, p=0.5, training=True)
+
+    inp = torch.ones(4096, device="cuda", requires_grad=True)
+    out = te_checkpoint(body, inp, use_reentrant=True, get_rng_state_tracker=lambda: tracker)
+    out.sum().backward()
+    assert torch.equal(inp.grad, out)
+
+    offset = generator.get_offset()
+    graph.replay()
+    assert generator.get_offset() - offset == increment
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 55b1b3ae0fe..63e33390090 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -328,6 +328,63 @@ def _get_active_autocast_contexts():
     return gpu_autocast_ctx, cpu_autocast_ctx
 
 
+class _CheckpointRNGState:  # pylint: disable=too-few-public-methods
+    """Save the forward RNG state and restore it temporarily during recompute."""
+
+    def __init__(self, get_rng_state_tracker):
+        self.get_rng_state_tracker = get_rng_state_tracker
+        self.cpu_state = torch.get_rng_state()
+        self.tracker_states = (
+            get_rng_state_tracker().get_states() if get_rng_state_tracker is not None else None
+        )
+        graph_safe = bool(self.tracker_states) and is_graph_safe_rng_state(
+            next(iter(self.tracker_states.values()))
+        )
+        self.cuda_state = _get_cuda_rng_state(graph_safe=graph_safe)
+        self.generator_states = []
+        if (
+            self.tracker_states
+            and graph_safe_rng_available()
+            and not torch.cuda.is_current_stream_capturing()
+        ):
+            # Keep the original handles to preserve aliases and CUDA graph registrations.
+            generators = (self.cuda_state, *self.tracker_states.values())
+            self.generator_states = [
+                (gen, gen.get_state()) for gen in generators if isinstance(gen, torch.Generator)
+            ]
+
+    @contextmanager
+    def fork(self):
+        """Rewind RNG for recompute, then restore the live state."""
+        graph_safe = is_graph_safe_rng_state(self.cuda_state)
+        cpu_state = torch.get_rng_state()
+        cuda_state = _get_cuda_rng_state(graph_safe=graph_safe)
+        tracker = self.get_rng_state_tracker() if self.get_rng_state_tracker is not None else None
+        tracker_states = tracker.get_states() if tracker is not None else None
+
+        generator_states = self.generator_states
+        if generator_states and torch.cuda.is_current_stream_capturing():
+            # Keep capture-local offsets when the forward saved an eager snapshot.
+            generator_states = []
+        live_generator_states = [(gen, gen.get_state()) for gen, _ in generator_states]
+
+        try:
+            for gen, state in generator_states:
+                gen.set_state(state)
+            torch.set_rng_state(self.cpu_state)
+            _set_cuda_rng_state(self.cuda_state, graph_safe=graph_safe)
+            if tracker is not None:
+                tracker.set_states(self.tracker_states.copy())
+            yield
+        finally:
+            for gen, state in live_generator_states:
+                gen.set_state(state)
+            torch.set_rng_state(cpu_state)
+            _set_cuda_rng_state(cuda_state, graph_safe=graph_safe)
+            if tracker is not None:
+                tracker.set_states(tracker_states)
+
+
 class _CheckpointFunction(torch.autograd.Function):
     """This function is adapted from torch.utils.checkpoint with
     two main changes:
@@ -352,18 +409,7 @@ def forward(
         ctx.run_function = run_function
         ctx.distribute_saved_activations = distribute_saved_activations
 
-        # Copy the rng states.
-        ctx.fwd_cpu_rng_state = torch.get_rng_state()
-        if get_rng_state_tracker is not None:
-            ctx.fwd_cuda_rng_state_tracker = get_rng_state_tracker().get_states()
-            ctx.graph_safe_rng_state = (
-                is_graph_safe_rng_state(next(iter(ctx.fwd_cuda_rng_state_tracker.values())))
-                if ctx.fwd_cuda_rng_state_tracker
-                else False
-            )
-        else:
-            ctx.graph_safe_rng_state = False
-        ctx.fwd_cuda_rng_state = _get_cuda_rng_state(graph_safe=ctx.graph_safe_rng_state)
+        ctx.rng_state = _CheckpointRNGState(get_rng_state_tracker)
 
         if context_fn is not None:
             forward_ctx, recompute_ctx = context_fn()
@@ -397,7 +443,6 @@ def forward(
         ctx.save_for_backward(*tensor_inputs)
 
         fp8 = FP8GlobalStateManager.is_fp8_enabled()
-        ctx.get_rng_state_tracker = get_rng_state_tracker
         ctx.tp_group = tp_group
         ctx.recompute_ctx = recompute_ctx
         ctx.torch_gpu_amp_ctx = torch_gpu_amp_ctx
@@ -422,41 +467,25 @@ def backward(
             t if t is not None else arg for (t, arg) in zip(ctx.saved_tensors, ctx.inputs)
         )
 
-        get_rng_state_tracker = ctx.get_rng_state_tracker
-
         if ctx.distribute_saved_activations:
             safely_set_viewless_tensor_data(
                 inputs[0],
                 gather_split_1d_tensor(inputs[0].data, ctx.tp_group).view(ctx.input_0_shape),
             )
 
-        # Store the current states.
-        bwd_cpu_rng_state = torch.get_rng_state()
-        bwd_cuda_rng_state = _get_cuda_rng_state(graph_safe=ctx.graph_safe_rng_state)
-        if get_rng_state_tracker is not None:
-            bwd_cuda_rng_state_tracker = get_rng_state_tracker().get_states()
-
-        # Set the states to what it used to be before the forward pass.
-        torch.set_rng_state(ctx.fwd_cpu_rng_state)
-        _set_cuda_rng_state(ctx.fwd_cuda_rng_state, graph_safe=ctx.graph_safe_rng_state)
-        if get_rng_state_tracker is not None:
-            get_rng_state_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)
-
         # Compute the forward pass.
         detached_inputs = detach_variable(inputs)
-        with torch.enable_grad(), ctx.recompute_ctx, ctx.torch_gpu_amp_ctx, ctx.torch_cpu_amp_ctx, activation_recompute_forward(
-            activation_recompute=True, recompute_phase=True
-        ), autocast(
-            enabled=ctx.fp8, recipe=ctx.fp8_recipe
+        with (
+            ctx.rng_state.fork(),
+            torch.enable_grad(),
+            ctx.recompute_ctx,
+            ctx.torch_gpu_amp_ctx,
+            ctx.torch_cpu_amp_ctx,
+            activation_recompute_forward(activation_recompute=True, recompute_phase=True),
+            autocast(enabled=ctx.fp8, recipe=ctx.fp8_recipe),
         ):
             outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
 
-        # Set the states back to what it was at the start of this function.
-        torch.set_rng_state(bwd_cpu_rng_state)
-        _set_cuda_rng_state(bwd_cuda_rng_state, graph_safe=ctx.graph_safe_rng_state)
-        if get_rng_state_tracker is not None:
-            get_rng_state_tracker().set_states(bwd_cuda_rng_state_tracker)
-
         if isinstance(outputs, torch.Tensor):
             outputs = (outputs,)
 
@@ -481,7 +510,7 @@ def backward(
         return (None, None, None, None, None, None) + grads
 
 
-class _CheckpointFrame:
+class _CheckpointFrame:  # pylint: disable=too-few-public-methods
     """
     Storage frame for forward RNG states and detached activations from the forward recompute.
     """
@@ -490,44 +519,7 @@ def __init__(self, recompute_fn: Callable, get_rng_state_tracker: Callable):
         self.recompute_fn = recompute_fn
         self.recomputed = []
         self.count = 0
-        self.get_rng_state_tracker = get_rng_state_tracker
-        self.fwd_rng_states = None
-        self.bwd_rng_states = None
-
-    def cache_rng_states(self, forward=True):
-        """Cache fwd/bwd RNG states in the frame to restore later."""
-        rng_states = (torch.get_rng_state(),)
-        if self.get_rng_state_tracker is not None:
-            tracker_states = self.get_rng_state_tracker().get_states()
-            self.graph_safe_rng_state = (
-                is_graph_safe_rng_state(next(iter(tracker_states.values())))
-                if tracker_states
-                else False
-            )
-            rng_states += (
-                _get_cuda_rng_state(graph_safe=self.graph_safe_rng_state),
-                tracker_states,
-            )
-        else:
-            self.graph_safe_rng_state = False
-            rng_states += (_get_cuda_rng_state(graph_safe=self.graph_safe_rng_state),)
-
-        if forward:
-            self.fwd_rng_states = rng_states
-        else:
-            self.bwd_rng_states = rng_states
-
-    def restore_rng_states(self, forward=True):
-        """Restore fwd/bwd RNG states that were previously cached into the frame."""
-        if forward:
-            rng_states = self.fwd_rng_states
-        else:
-            rng_states = self.bwd_rng_states
-
-        torch.set_rng_state(rng_states[0])
-        _set_cuda_rng_state(rng_states[1], graph_safe=self.graph_safe_rng_state)
-        if self.get_rng_state_tracker is not None:
-            self.get_rng_state_tracker().set_states(rng_states[2])
+        self.rng_state = _CheckpointRNGState(get_rng_state_tracker)
 
 
 class _recomputation_hook(
@@ -583,19 +575,9 @@ def unpack_hook(idx):
             """
 
             if not frame.recomputed:
-                # Store current RNG states in the backward pass
-                frame.cache_rng_states(forward=False)
-
-                # Set RNG states to what we saved before the forward pass
-                frame.restore_rng_states(forward=True)
-
-                # Recompute the forward pass
-                with _recomputation_hook(frame):
+                with frame.rng_state.fork(), _recomputation_hook(frame):
                     frame.recompute_fn(*args, **kwargs)
 
-                # Restore RNG states back to the backward pass
-                frame.restore_rng_states(forward=False)
-
             # Return the already recomputed activation tensor at the given index
             activation = frame.recomputed[idx]
             frame.recomputed[idx] = None
@@ -822,8 +804,6 @@ def recompute_fn(*args, **kwargs):
         recompute_fn,
         get_rng_state_tracker,
     )
-    new_frame.cache_rng_states(forward=True)
-
     with _checkpoint_hook(new_frame, args, kwargs), te_forward_ctx, user_forward_ctx:
         out = function(*args, **kwargs)
 
@@ -858,6 +838,12 @@ def get_states(self) -> Dict[str, torch.Tensor]:
         """
         Get rng states. Copy the dictionary so we have direct pointers
         to the states, not just a pointer to the dictionary.
+
+        .. note::
+            Graph-safe RNG states are live ``torch.Generator`` handles rather
+            than snapshots. They advance in place as random numbers are drawn;
+            call ``Generator.get_state()`` to save the current seed and offset
+            outside CUDA graph capture.
         """
         states = {}
         for name in self.states_:

From 613a44987761140befff41a449f3312a2b70ca09 Mon Sep 17 00:00:00 2001
From: losewayy <104015127+losewayy@users.noreply.github.com>
Date: Sun, 8 Nov 2026 22:48:10 +0800
Subject: [PATCH 02/23] [PyTorch] Test CUDA RNG restoration after checkpoint
 backward

Draw from the default and tracked CUDA generators between forward and backward. Compare their states before and after backward in both checkpoint modes to catch any change to the live streams.

Signed-off-by: losewayy <104015127+losewayy@users.noreply.github.com>
---
 tests/pytorch/test_numerics.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 12373858382..b2425155bd1 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1138,7 +1138,7 @@ def outer_body(value):
 @pytest.mark.parametrize("use_reentrant", all_boolean)
 @pytest.mark.parametrize("rng_source", ("default", "tracker", "shared"))
 def test_checkpoint_recompute_preserves_tracker_rng_state(use_reentrant, rng_source):
-    """Recomputed dropout must match forward for default, tracked, and shared RNG."""
+    """Recompute must match forward and preserve the live RNG states."""
 
     def run(checkpointed):
         torch.manual_seed(seed)
@@ -1149,6 +1149,10 @@ def run(checkpointed):
         inp = torch.ones(4096, device="cuda", requires_grad=True)
         draws = []
 
+        def get_tracked_rng_state():
+            state = tracker.get_states()["model-parallel-rng"]
+            return state.get_state() if isinstance(state, torch.Generator) else state.clone()
+
         def body(value):
             if rng_source == "tracker":
                 with tracker.fork():
@@ -1175,7 +1179,20 @@ def body(value):
                 )
             else:
                 out = body(inp)
+
+        # Backward must preserve the RNG positions after unrelated draws.
+        torch.rand(8, device=inp.device)
+        with tracker.fork():
+            torch.rand(8, device=inp.device)
+        default_rng_before = torch.cuda.get_rng_state(inp.device)
+        tracked_rng_before = get_tracked_rng_state()
         out.sum().backward()
+        assert torch.equal(
+            torch.cuda.get_rng_state(inp.device), default_rng_before
+        ), "Default RNG changed during backward"
+        assert torch.equal(
+            get_tracked_rng_state(), tracked_rng_before
+        ), "Tracked RNG changed during backward"
         return out.detach(), inp.grad.detach(), draws, tracker, tracker_state
 
     ref_out, ref_grad, _, _, _ = run(checkpointed=False)

From 26a46cc450c7979b7dbd6db7e94558c72cb254e6 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 11:44:38 +0100
Subject: [PATCH 03/23] [PyTorch] Restore captured checkpoint RNG state on
 device

Snapshot effective Philox state on CUDA during capture and rewind the original registered generators for recompute. Save RNG tensors with autograd in both TE checkpoint implementations, and remove RNG compatibility paths for unsupported PyTorch versions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py                |  59 +++++-
 transformer_engine/common/fused_attn/utils.cu |  15 ++
 .../include/transformer_engine/fused_attn.h   |  12 ++
 .../dot_product_attention.py                  |   4 -
 transformer_engine/pytorch/csrc/common.cpp    |  38 ++++
 transformer_engine/pytorch/csrc/common.h      |   3 +
 .../pytorch/csrc/extensions/pybind.cpp        |   5 +
 transformer_engine/pytorch/distributed.py     | 195 +++++++++---------
 transformer_engine/pytorch/graph.py           |  33 ++-
 9 files changed, 233 insertions(+), 131 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index b2425155bd1..284b07b9ab2 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -47,7 +47,6 @@
 from transformer_engine.pytorch import checkpoint as te_checkpoint
 from transformer_engine.pytorch.distributed import (
     CudaRNGStatesTracker,
-    graph_safe_rng_available,
     is_fp8_activation_recompute_enabled,
     in_fp8_activation_recompute_phase,
 )
@@ -1205,9 +1204,9 @@ def body(value):
         assert tracker.get_states()["model-parallel-rng"] is tracker_state
 
 
-@pytest.mark.skipif(not graph_safe_rng_available(), reason="Requires graph-safe CUDA RNG")
+@pytest.mark.parametrize("use_reentrant", all_boolean)
 @pytest.mark.parametrize("fork_tracker", all_boolean)
-def test_checkpoint_preserves_captured_rng_state(fork_tracker):
+def test_checkpoint_preserves_captured_rng_state(fork_tracker, use_reentrant):
     """Graph replay must still advance its registered RNG after eager recompute."""
     tracker = CudaRNGStatesTracker()
     tracker.add("model-parallel-rng", 4321)
@@ -1236,7 +1235,9 @@ def body(value):
             return F.dropout(value, p=0.5, training=True)
 
     inp = torch.ones(4096, device="cuda", requires_grad=True)
-    out = te_checkpoint(body, inp, use_reentrant=True, get_rng_state_tracker=lambda: tracker)
+    out = te_checkpoint(
+        body, inp, use_reentrant=use_reentrant, get_rng_state_tracker=lambda: tracker
+    )
     out.sum().backward()
     assert torch.equal(inp.grad, out)
 
@@ -1245,6 +1246,56 @@ def body(value):
     assert generator.get_offset() - offset == increment
 
 
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+@pytest.mark.parametrize("rng_source", ("default", "tracker", "mixed", "shared", "no_tracker"))
+def test_checkpoint_dropout_gradients_during_graph_replay(use_reentrant, rng_source):
+    """Each captured recompute must reproduce its corresponding forward masks."""
+    torch.manual_seed(seed)
+    tracker = CudaRNGStatesTracker()
+    tracker.add("model-parallel-rng", 4321)
+    saved_gpu_states = []
+
+    def pack(tensor):
+        if tensor.is_cuda and tensor.dtype == torch.int64 and tensor.shape == (2,):
+            saved_gpu_states.append(torch.cuda.is_current_stream_capturing())
+        return tensor
+
+    def body(value):
+        if rng_source == "tracker":
+            with tracker.fork():
+                return F.dropout(value, p=0.5, training=True)
+        out = F.dropout(value, p=0.5, training=True)
+        if rng_source in ("mixed", "shared"):
+            with tracker.fork():
+                out = out + 2 * F.dropout(value, p=0.5, training=True)
+        return out
+
+    class Model(torch.nn.Module):
+        def forward(self, value):
+            with tracker.fork() if rng_source == "shared" else nullcontext():
+                return te_checkpoint(
+                    body,
+                    value,
+                    use_reentrant=use_reentrant,
+                    get_rng_state_tracker=None if rng_source == "no_tracker" else lambda: tracker,
+                )
+
+    sample = torch.ones(4096, device="cuda", requires_grad=True)
+    with torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor):
+        model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
+    previous = None
+    for _ in range(3):
+        value = torch.ones_like(sample, requires_grad=True)
+        output = model(value)
+        forward_output = output.detach().clone()
+        output.sum().backward()
+        assert torch.equal(value.grad, forward_output), "Recompute used a different dropout mask"
+        if previous is not None:
+            assert not torch.equal(forward_output, previous), "Graph replay did not advance RNG"
+        previous = forward_output
+    assert any(saved_gpu_states), "Captured RNG snapshots must use autograd's saved-tensor hooks"
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""
diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu
index b74dcfe4344..0810f4d2ecc 100644
--- a/transformer_engine/common/fused_attn/utils.cu
+++ b/transformer_engine/common/fused_attn/utils.cu
@@ -520,6 +520,12 @@ __global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, i
   }
 }
 
+__global__ void restore_seed_and_offset(const int64_t *rng_state_ptr, int64_t *seed_ptr,
+                                        int64_t *offset_ptr, uint32_t offset_intragraph) {
+  *seed_ptr = rng_state_ptr[0];
+  *offset_ptr = static_cast<int64_t>(static_cast<uint64_t>(rng_state_ptr[1]) - offset_intragraph);
+}
+
 }  // namespace fused_attn
 }  // namespace transformer_engine
 
@@ -533,3 +539,12 @@ void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t
       rng_state_ptr, captured, seed_ptr, seed_val, offset_ptr, offset_val, offset_intragraph);
   NVTE_CHECK_CUDA(cudaGetLastError());
 }
+
+void nvte_restore_seed_and_offset(const int64_t *rng_state_ptr, int64_t *seed_ptr,
+                                  int64_t *offset_ptr, uint32_t offset_intragraph,
+                                  cudaStream_t stream) {
+  NVTE_API_CALL(nvte_restore_seed_and_offset);
+  transformer_engine::fused_attn::restore_seed_and_offset<<<1, 1, 0, stream>>>(
+      rng_state_ptr, seed_ptr, offset_ptr, offset_intragraph);
+  NVTE_CHECK_CUDA(cudaGetLastError());
+}
diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h
index 7d1ffe13198..9a6f580d53d 100644
--- a/transformer_engine/common/include/transformer_engine/fused_attn.h
+++ b/transformer_engine/common/include/transformer_engine/fused_attn.h
@@ -693,6 +693,18 @@ void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t
                                   uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val,
                                   uint32_t offset_intragraph, cudaStream_t stream);
 
+/*! \brief Restore a captured RNG's seed and effective offset from a device snapshot.
+ *
+ * \param[in] rng_state_ptr Device array containing the seed and effective Philox offset.
+ * \param[out] seed_ptr Captured generator's seed.
+ * \param[out] offset_ptr Captured generator's base offset.
+ * \param[in] offset_intragraph Current intragraph offset, subtracted from the saved offset.
+ * \param[in] stream CUDA stream used for this operation.
+ */
+void nvte_restore_seed_and_offset(const int64_t *rng_state_ptr, int64_t *seed_ptr,
+                                  int64_t *offset_ptr, uint32_t offset_intragraph,
+                                  cudaStream_t stream);
+
 /*!  \brief Copy keys and values into the KV cache.
  *
  * \warning   This API is **experimental** and subject to change.
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 9a3c8a44dbf..80844adca43 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -51,7 +51,6 @@
     checkpoint,
     set_all_rng_states,
     CudaRNGStatesTracker,
-    graph_safe_rng_available,
 )
 from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.utils import SplitAlongDim
@@ -2342,9 +2341,6 @@ def forward(
                 assert isinstance(
                     self.rng_states_tracker, CudaRNGStatesTracker
                 ), "Unsupported RNG states tracker."
-                assert (
-                    graph_safe_rng_available()
-                ), "Upgrade PyTorch version to get RNG manipulation support for cuda graph capture."
 
             # checks for FP8
             if self.fp8:
diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index d85dcda1593..06259be9aab 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -6,6 +6,8 @@
 
 #include "common.h"
 
+#include <c10/cuda/CUDAGuard.h>
+
 #include "c10/util/ArrayRef.h"
 #include "pybind.h"
 #include "transformer_engine/transformer_engine.h"
@@ -364,4 +366,40 @@ at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_pe
   return philox_args;
 }
 
+at::Tensor checkpoint_rng_state(const at::Generator& generator) {
+  auto* gen = at::check_generator<at::CUDAGeneratorImpl>(generator);
+  const c10::cuda::CUDAGuard device_guard(generator.device());
+  if (at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None) {
+    std::lock_guard<std::mutex> lock(gen->mutex_);
+    return generator.get_state();
+  }
+  auto state = at::empty({2}, at::TensorOptions().dtype(at::kLong).device(generator.device()));
+  philox_unpack(init_philox_state(gen, 0), state.data_ptr<int64_t>());
+  return state;
+}
+
+void restore_checkpoint_rng_state(const at::Generator& generator, const at::Tensor& state) {
+  auto* gen = at::check_generator<at::CUDAGeneratorImpl>(generator);
+  const c10::cuda::CUDAGuard device_guard(generator.device());
+  const bool capturing = at::cuda::currentStreamCaptureStatus() != at::cuda::CaptureStatus::None;
+  TORCH_CHECK(
+      state.is_cuda() == capturing,
+      "Checkpoint forward and recompute must both use eager execution or both be captured.");
+  if (!capturing) {
+    std::lock_guard<std::mutex> lock(gen->mutex_);
+    auto mutable_generator = generator;
+    mutable_generator.set_state(state);
+    return;
+  }
+  TORCH_CHECK(state.device() == generator.device() && state.scalar_type() == at::kLong &&
+                  state.numel() == 2 && state.is_contiguous(),
+              "Expected a contiguous int64 CUDA RNG snapshot with two elements.");
+  auto philox = init_philox_state(gen, 0);
+  // Recompute has its own intragraph position; preserve the saved effective offset.
+  NVTE_SCOPED_GIL_RELEASE({
+    nvte_restore_seed_and_offset(state.data_ptr<int64_t>(), philox.seed_.ptr, philox.offset_.ptr,
+                                 philox.offset_intragraph_, at::cuda::getCurrentCUDAStream());
+  });
+}
+
 }  // namespace transformer_engine::pytorch
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
index 9a3c6314054..859e0562c30 100644
--- a/transformer_engine/pytorch/csrc/common.h
+++ b/transformer_engine/pytorch/csrc/common.h
@@ -629,6 +629,9 @@ void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr);
 // extract PhiloxCudaState from CUDA random number generator
 at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_per_thread);
 
+at::Tensor checkpoint_rng_state(const at::Generator& generator);
+void restore_checkpoint_rng_state(const at::Generator& generator, const at::Tensor& state);
+
 }  // namespace transformer_engine::pytorch
 
 namespace std {
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index 6fbed1436fb..41d468ca105 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -196,6 +196,11 @@ void bind_quantize_with_amax_extensions(py::module_ &m) {
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m)
 
+  m.def("checkpoint_rng_state", &transformer_engine::pytorch::checkpoint_rng_state,
+        py::arg("generator"));
+  m.def("restore_checkpoint_rng_state", &transformer_engine::pytorch::restore_checkpoint_rng_state,
+        py::arg("generator"), py::arg("state"));
+
   // Register __eq__/__ne__ on the pybind ``DType`` enum so it compares by integer.
   py::object dtype_class = m.attr("DType");
   dtype_class.attr("__eq__") = py::cpp_function(
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 63e33390090..b8ea8ff0e11 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -81,26 +81,10 @@ def set_all_rng_states(states: List) -> None:
     _ALL_ACTIVE_RNG_STATES = states
 
 
-def graph_safe_rng_available() -> bool:
-    """Returns whether cuda graph safe RNG state manipulation is supported."""
-    return (
-        hasattr(torch.cuda.CUDAGraph, "register_generator_state")
-        and hasattr(torch.Generator, "graphsafe_set_state")
-        and hasattr(torch.Generator, "graphsafe_get_state")
-        and hasattr(torch.Generator, "clone_state")
-    )
-
-
-def is_graph_safe_rng_state(state: Union[torch.Tensor, torch.Generator]) -> bool:
-    """Returns whether the rng state is a graph safe version."""
-    return graph_safe_rng_available() and isinstance(state, torch.Generator)
-
-
 def _get_cuda_rng_state(
     device: Union[int, str, torch.device] = te_device_type(),
     clone: bool = False,
-    graph_safe: bool = True,
-) -> torch.Tensor:
+) -> torch.Generator:
     """Return the random number generator state of the specified GPU."""
 
     _lazy_init()
@@ -112,19 +96,12 @@ def _get_cuda_rng_state(
     if idx is None:
         idx = torch.cuda.current_device()
     default_generator = torch.cuda.default_generators[idx]
-    if graph_safe_rng_available() and graph_safe:
-        if clone:
-            # Reference to the cloned generator state
-            return default_generator.clone_state()
-        # Reference to the current generator state
-        return default_generator.graphsafe_get_state()
-    return default_generator.get_state()
+    return default_generator.clone_state() if clone else default_generator.graphsafe_get_state()
 
 
 def _set_cuda_rng_state(
-    new_state: torch.Tensor,
+    new_state: torch.Generator,
     device: Union[int, str] = -1,
-    graph_safe=True,
 ) -> None:
     """Sets the random number generator state of the current GPU."""
 
@@ -140,10 +117,7 @@ def cb() -> None:
         if idx is None:
             idx = torch.cuda.current_device()
         default_generator = torch.cuda.default_generators[idx]
-        if graph_safe_rng_available() and graph_safe:
-            default_generator.graphsafe_set_state(new_state)
-            return
-        default_generator.set_state(new_state)
+        default_generator.graphsafe_set_state(new_state)
 
     _lazy_call(cb)
 
@@ -328,61 +302,67 @@ def _get_active_autocast_contexts():
     return gpu_autocast_ctx, cpu_autocast_ctx
 
 
-class _CheckpointRNGState:  # pylint: disable=too-few-public-methods
-    """Save the forward RNG state and restore it temporarily during recompute."""
+class _CheckpointRNGState:
+    """Save RNG values separately from live generator handles."""
 
     def __init__(self, get_rng_state_tracker):
-        self.get_rng_state_tracker = get_rng_state_tracker
-        self.cpu_state = torch.get_rng_state()
-        self.tracker_states = (
-            get_rng_state_tracker().get_states() if get_rng_state_tracker is not None else None
-        )
-        graph_safe = bool(self.tracker_states) and is_graph_safe_rng_state(
-            next(iter(self.tracker_states.values()))
-        )
-        self.cuda_state = _get_cuda_rng_state(graph_safe=graph_safe)
-        self.generator_states = []
-        if (
-            self.tracker_states
-            and graph_safe_rng_available()
-            and not torch.cuda.is_current_stream_capturing()
-        ):
-            # Keep the original handles to preserve aliases and CUDA graph registrations.
-            generators = (self.cuda_state, *self.tracker_states.values())
-            self.generator_states = [
-                (gen, gen.get_state()) for gen in generators if isinstance(gen, torch.Generator)
-            ]
+        self.tracker = get_rng_state_tracker() if get_rng_state_tracker is not None else None
+        self.cuda_state = _get_cuda_rng_state()
+        self.generators = [self.cuda_state]
+        self.tracker_states = {}
+        for name, state in (self.tracker.get_states() if self.tracker is not None else {}).items():
+            if isinstance(state, torch.Generator):
+                self.generators.append(state)
+                self.tracker_states[name] = state
+            else:
+                self.tracker_states[name] = None
+
+    def snapshot(self):
+        """Copy the CPU state, generator values, and tensor-based tracker states."""
+        states = [torch.get_rng_state(), *(tex.checkpoint_rng_state(g) for g in self.generators)]
+        if self.tracker is not None:
+            tracker_states = self.tracker.get_states()
+            states.extend(
+                tracker_states[name] for name, state in self.tracker_states.items() if state is None
+            )
+        return states
 
-    @contextmanager
-    def fork(self):
-        """Rewind RNG for recompute, then restore the live state."""
-        graph_safe = is_graph_safe_rng_state(self.cuda_state)
-        cpu_state = torch.get_rng_state()
-        cuda_state = _get_cuda_rng_state(graph_safe=graph_safe)
-        tracker = self.get_rng_state_tracker() if self.get_rng_state_tracker is not None else None
-        tracker_states = tracker.get_states() if tracker is not None else None
-
-        generator_states = self.generator_states
-        if generator_states and torch.cuda.is_current_stream_capturing():
-            # Keep capture-local offsets when the forward saved an eager snapshot.
-            generator_states = []
-        live_generator_states = [(gen, gen.get_state()) for gen, _ in generator_states]
+    def _restore(self, states):
+        torch.set_rng_state(states[0])
+        for generator, state in zip(self.generators, states[1:]):
+            tex.restore_checkpoint_rng_state(generator, state)
 
+    @contextmanager
+    def fork(self, states):
+        """Restore forward RNG values temporarily, then restore the live values."""
+        cuda_state = _get_cuda_rng_state()
+        tracker_states = self.tracker.get_states() if self.tracker is not None else None
+        live_states = [
+            torch.get_rng_state(),
+            *(tex.checkpoint_rng_state(g) for g in self.generators),
+        ]
         try:
-            for gen, state in generator_states:
-                gen.set_state(state)
-            torch.set_rng_state(self.cpu_state)
-            _set_cuda_rng_state(self.cuda_state, graph_safe=graph_safe)
-            if tracker is not None:
-                tracker.set_states(self.tracker_states.copy())
+            self._restore(states)
+            _set_cuda_rng_state(self.cuda_state)
+            if self.tracker is not None:
+                tensor_states = dict(
+                    zip(
+                        (name for name, state in self.tracker_states.items() if state is None),
+                        states[1 + len(self.generators) :],
+                    )
+                )
+                self.tracker.set_states(
+                    {
+                        name: state if state is not None else tensor_states[name]
+                        for name, state in self.tracker_states.items()
+                    }
+                )
             yield
         finally:
-            for gen, state in live_generator_states:
-                gen.set_state(state)
-            torch.set_rng_state(cpu_state)
-            _set_cuda_rng_state(cuda_state, graph_safe=graph_safe)
-            if tracker is not None:
-                tracker.set_states(tracker_states)
+            self._restore(live_states)
+            _set_cuda_rng_state(cuda_state)
+            if self.tracker is not None:
+                self.tracker.set_states(tracker_states)
 
 
 class _CheckpointFunction(torch.autograd.Function):
@@ -410,6 +390,7 @@ def forward(
         ctx.distribute_saved_activations = distribute_saved_activations
 
         ctx.rng_state = _CheckpointRNGState(get_rng_state_tracker)
+        rng_tensors = ctx.rng_state.snapshot()
 
         if context_fn is not None:
             forward_ctx, recompute_ctx = context_fn()
@@ -440,7 +421,8 @@ def forward(
         # Store everything.
         ctx.inputs = [arg if not torch.is_tensor(arg) else None for arg in args]
         tensor_inputs = [arg if torch.is_tensor(arg) else None for arg in args]
-        ctx.save_for_backward(*tensor_inputs)
+        ctx.num_inputs = len(tensor_inputs)
+        ctx.save_for_backward(*tensor_inputs, *rng_tensors)
 
         fp8 = FP8GlobalStateManager.is_fp8_enabled()
         ctx.tp_group = tp_group
@@ -463,8 +445,10 @@ def backward(
                 "Checkpointing is not compatible with .grad(), please use .backward() if possible"
             )
 
+        saved_tensors = ctx.saved_tensors
         inputs = tuple(
-            t if t is not None else arg for (t, arg) in zip(ctx.saved_tensors, ctx.inputs)
+            t if t is not None else arg
+            for t, arg in zip(saved_tensors[: ctx.num_inputs], ctx.inputs)
         )
 
         if ctx.distribute_saved_activations:
@@ -476,7 +460,7 @@ def backward(
         # Compute the forward pass.
         detached_inputs = detach_variable(inputs)
         with (
-            ctx.rng_state.fork(),
+            ctx.rng_state.fork(saved_tensors[ctx.num_inputs :]),
             torch.enable_grad(),
             ctx.recompute_ctx,
             ctx.torch_gpu_amp_ctx,
@@ -510,6 +494,21 @@ def backward(
         return (None, None, None, None, None, None) + grads
 
 
+class _CheckpointRNGStateHolder(torch.autograd.Function):
+    """Route non-reentrant RNG snapshots through autograd's saved-tensor hooks."""
+
+    @staticmethod
+    def forward(ctx, dummy, *states):
+        """Save snapshots without connecting them to the model's backward graph."""
+        ctx.save_for_backward(*states)
+        return dummy
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        """Reject differentiation through the storage-only output."""
+        raise RuntimeError("RNG snapshot storage is not part of the backward graph")
+
+
 class _CheckpointFrame:  # pylint: disable=too-few-public-methods
     """
     Storage frame for forward RNG states and detached activations from the forward recompute.
@@ -520,6 +519,9 @@ def __init__(self, recompute_fn: Callable, get_rng_state_tracker: Callable):
         self.recomputed = []
         self.count = 0
         self.rng_state = _CheckpointRNGState(get_rng_state_tracker)
+        self.rng_state_holder = _CheckpointRNGStateHolder.apply(
+            torch.empty(0, requires_grad=True), *self.rng_state.snapshot()
+        )
 
 
 class _recomputation_hook(
@@ -575,7 +577,9 @@ def unpack_hook(idx):
             """
 
             if not frame.recomputed:
-                with frame.rng_state.fork(), _recomputation_hook(frame):
+                with frame.rng_state.fork(
+                    frame.rng_state_holder.grad_fn.saved_tensors
+                ), _recomputation_hook(frame):
                     frame.recompute_fn(*args, **kwargs)
 
             # Return the already recomputed activation tensor at the given index
@@ -834,7 +838,7 @@ def reset(self):
         self.states_ = {}
         self.seeds_ = set()
 
-    def get_states(self) -> Dict[str, torch.Tensor]:
+    def get_states(self) -> Dict[str, torch.Generator]:
         """
         Get rng states. Copy the dictionary so we have direct pointers
         to the states, not just a pointer to the dictionary.
@@ -850,7 +854,7 @@ def get_states(self) -> Dict[str, torch.Tensor]:
             states[name] = self.states_[name]
         return states
 
-    def set_states(self, states: Dict[str, torch.Tensor]) -> None:
+    def set_states(self, states: Dict[str, torch.Generator]) -> None:
         """
         Set the rng states. For efficiency purposes, we do not
         check the size of seed for compatibility.
@@ -883,22 +887,10 @@ def add(self, name: str, seed: int) -> None:
         if name in self.states_:
             raise RuntimeError(f"cuda rng state {name} already exists")
 
-        if graph_safe_rng_available():
-            new_state = _get_cuda_rng_state(clone=True)
-            new_state.manual_seed(seed)
-            self.states_[name] = new_state
-            # Update global states.
-            set_all_rng_states(self.states_)
-        else:
-            # Get the current rng state.
-            orig_rng_state = _get_cuda_rng_state()
-            # Set the new state and store it.
-            torch.cuda.manual_seed(seed)
-            self.states_[name] = _get_cuda_rng_state(clone=True)
-            # Reset rng state to what it was.
-            _set_cuda_rng_state(orig_rng_state)
-            # Update global states.
-            set_all_rng_states(self.states_)
+        new_state = _get_cuda_rng_state(clone=True)
+        new_state.manual_seed(seed)
+        self.states_[name] = new_state
+        set_all_rng_states(self.states_)
 
     @contextmanager
     def fork(self, name: str = "model-parallel-rng"):
@@ -922,9 +914,6 @@ def fork(self, name: str = "model-parallel-rng"):
         try:
             yield
         finally:
-            # this is redundant with graph-safe API
-            if not graph_safe_rng_available():
-                self.states_[name] = _get_cuda_rng_state()
             # And set the state to the original state we started with.
             _set_cuda_rng_state(orig_cuda_rng_state)
 
diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py
index be88f730a18..206f86cafc5 100644
--- a/transformer_engine/pytorch/graph.py
+++ b/transformer_engine/pytorch/graph.py
@@ -22,7 +22,7 @@
     FP8GlobalStateManager,
     get_default_fp8_recipe,
 )
-from .distributed import get_all_rng_states, graph_safe_rng_available
+from .distributed import get_all_rng_states
 from .module.base import TransformerEngineBaseModule
 from .ops.op import BasicOperation
 from .ops import Sequential
@@ -482,12 +482,11 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params):
         )
 
     # For cases with multiple active RNG states, e.g. TP.
-    if graph_safe_rng_available():
-        for _, state in get_all_rng_states().items():
-            for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs):
-                fwd_graph.register_generator_state(state)
-                bwd_graph.register_generator_state(state)
-                bwd_dw_graph.register_generator_state(state)
+    for state in get_all_rng_states().values():
+        for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs):
+            fwd_graph.register_generator_state(state)
+            bwd_graph.register_generator_state(state)
+            bwd_dw_graph.register_generator_state(state)
 
     mempool = graph_pool_handle() if pool is None else pool
 
@@ -1712,14 +1711,11 @@ def call_func(self, *args, **kwargs):
         forward_funcs = tuple(forward_funcs)
 
     # Save RNG state.
-    if graph_safe_rng_available():
-        generators = [
-            torch.cuda.default_generators[torch.cuda.current_device()],
-            *get_all_rng_states().values(),
-        ]
-        original_rng_states = [state.get_state() for state in generators]
-    else:
-        original_rng_states = torch.cuda.get_rng_state()
+    generators = [
+        torch.cuda.default_generators[torch.cuda.current_device()],
+        *get_all_rng_states().values(),
+    ]
+    original_rng_states = [state.get_state() for state in generators]
 
     graphed_callables = _make_graphed_callables(
         forward_funcs,
@@ -1740,11 +1736,8 @@ def call_func(self, *args, **kwargs):
     )
 
     # Ensures warmup does not affect numerics for ops such as dropout.
-    if graph_safe_rng_available():
-        for gen, state in zip(generators, original_rng_states):
-            gen.set_state(state)
-    else:
-        torch.cuda.set_rng_state(original_rng_states)
+    for gen, state in zip(generators, original_rng_states):
+        gen.set_state(state)
 
     # Remove FP8 wrapper.
     for module_cls, old_call in old_call_funcs.items():

From e5d2d086cfc08e8681c6b88d15ef9edeb626d7fd Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 12:32:14 +0100
Subject: [PATCH 04/23] [PyTorch] Use public Philox API for captured checkpoint
 RNG

Replace native checkpoint RNG helpers with Python tensor operations. Require Generator.philox_state for stochastic CUDA graph capture and reject unsupported random operations clearly on older builds, while retaining eager dropout and deterministic capture. Restore graph and recompute contexts on errors.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py                | 63 ++++++++++++
 transformer_engine/common/fused_attn/utils.cu | 15 ---
 .../include/transformer_engine/fused_attn.h   | 12 ---
 .../dot_product_attention/backends.py         |  9 +-
 transformer_engine/pytorch/csrc/common.cpp    | 38 -------
 transformer_engine/pytorch/csrc/common.h      |  3 -
 .../pytorch/csrc/extensions/pybind.cpp        |  5 -
 transformer_engine/pytorch/distributed.py     | 99 +++++++++++++++++--
 transformer_engine/pytorch/graph.py           | 72 +++++++-------
 .../pytorch/ops/basic/dropout.py              |  6 +-
 10 files changed, 205 insertions(+), 117 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 284b07b9ab2..dd99c8de9cd 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -45,6 +45,7 @@
     is_nvfp4_available,
 )
 from transformer_engine.pytorch import checkpoint as te_checkpoint
+import transformer_engine.pytorch.distributed as te_distributed
 from transformer_engine.pytorch.distributed import (
     CudaRNGStatesTracker,
     is_fp8_activation_recompute_enabled,
@@ -1248,6 +1249,9 @@ def body(value):
 
 @pytest.mark.parametrize("use_reentrant", all_boolean)
 @pytest.mark.parametrize("rng_source", ("default", "tracker", "mixed", "shared", "no_tracker"))
+@pytest.mark.skipif(
+    not hasattr(torch.Generator, "philox_state"), reason="Captured RNG requires PyTorch 2.14+"
+)
 def test_checkpoint_dropout_gradients_during_graph_replay(use_reentrant, rng_source):
     """Each captured recompute must reproduce its corresponding forward masks."""
     torch.manual_seed(seed)
@@ -1296,6 +1300,65 @@ def forward(self, value):
     assert any(saved_gpu_states), "Captured RNG snapshots must use autograd's saved-tensor hooks"
 
 
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+@pytest.mark.parametrize("dropout", (0.0, 0.5))
+@pytest.mark.parametrize("te_dropout", all_boolean)
+def test_checkpoint_graph_rng_without_philox_api(use_reentrant, dropout, te_dropout, monkeypatch):
+    """Older builds allow deterministic capture and reject stochastic recompute."""
+    monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
+    layer = te.ops.Dropout(dropout) if te_dropout else torch.nn.Dropout(dropout)
+
+    class Model(torch.nn.Module):
+        def forward(self, value):
+            return 2 * te_checkpoint(layer, value + 0, use_reentrant=use_reentrant)
+
+    sample = torch.ones(4096, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    # The same dropout remains supported outside capture.
+    output = Model()(sample)
+    expected = output.detach().clone()
+    output.sum().backward()
+    assert torch.equal(sample.grad, expected)
+    del output, expected
+    sample.grad = None
+
+    original_call = Model.__call__
+    original_recompute_flags = tuple(
+        te_distributed.activation_recompute_forward._is_first_fp8_module
+    )
+    if dropout:
+        with pytest.raises(RuntimeError, match="philox_state"):
+            te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
+        assert Model.__call__ is original_call
+        assert (
+            tuple(te_distributed.activation_recompute_forward._is_first_fp8_module)
+            == original_recompute_flags
+        )
+    else:
+        model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
+        value = torch.ones_like(sample, requires_grad=True)
+        model(value).sum().backward()
+        assert torch.equal(value.grad, torch.full_like(value, 2))
+
+
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+@pytest.mark.parametrize("backend", ("FusedAttention", "FlashAttention"))
+def test_checkpoint_native_attention_requires_philox_api(use_reentrant, backend, monkeypatch):
+    """Native attention dropout must reject capture before entering its CUDA backend."""
+    from transformer_engine.pytorch.attention.dot_product_attention import backends
+
+    monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
+    attention = getattr(backends, backend)(softmax_scale=0.125, attention_dropout=0.5)
+    value = torch.ones(16, 1, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+
+    def body(query):
+        return attention(query, query, query)
+
+    graph = torch.cuda.CUDAGraph()
+    with pytest.raises(RuntimeError, match="philox_state"):
+        with torch.cuda.graph(graph):
+            te_checkpoint(body, value + 0, use_reentrant=use_reentrant)
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""
diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu
index 0810f4d2ecc..b74dcfe4344 100644
--- a/transformer_engine/common/fused_attn/utils.cu
+++ b/transformer_engine/common/fused_attn/utils.cu
@@ -520,12 +520,6 @@ __global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, i
   }
 }
 
-__global__ void restore_seed_and_offset(const int64_t *rng_state_ptr, int64_t *seed_ptr,
-                                        int64_t *offset_ptr, uint32_t offset_intragraph) {
-  *seed_ptr = rng_state_ptr[0];
-  *offset_ptr = static_cast<int64_t>(static_cast<uint64_t>(rng_state_ptr[1]) - offset_intragraph);
-}
-
 }  // namespace fused_attn
 }  // namespace transformer_engine
 
@@ -539,12 +533,3 @@ void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t
       rng_state_ptr, captured, seed_ptr, seed_val, offset_ptr, offset_val, offset_intragraph);
   NVTE_CHECK_CUDA(cudaGetLastError());
 }
-
-void nvte_restore_seed_and_offset(const int64_t *rng_state_ptr, int64_t *seed_ptr,
-                                  int64_t *offset_ptr, uint32_t offset_intragraph,
-                                  cudaStream_t stream) {
-  NVTE_API_CALL(nvte_restore_seed_and_offset);
-  transformer_engine::fused_attn::restore_seed_and_offset<<<1, 1, 0, stream>>>(
-      rng_state_ptr, seed_ptr, offset_ptr, offset_intragraph);
-  NVTE_CHECK_CUDA(cudaGetLastError());
-}
diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h
index 9a6f580d53d..7d1ffe13198 100644
--- a/transformer_engine/common/include/transformer_engine/fused_attn.h
+++ b/transformer_engine/common/include/transformer_engine/fused_attn.h
@@ -693,18 +693,6 @@ void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t
                                   uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val,
                                   uint32_t offset_intragraph, cudaStream_t stream);
 
-/*! \brief Restore a captured RNG's seed and effective offset from a device snapshot.
- *
- * \param[in] rng_state_ptr Device array containing the seed and effective Philox offset.
- * \param[out] seed_ptr Captured generator's seed.
- * \param[out] offset_ptr Captured generator's base offset.
- * \param[in] offset_intragraph Current intragraph offset, subtracted from the saved offset.
- * \param[in] stream CUDA stream used for this operation.
- */
-void nvte_restore_seed_and_offset(const int64_t *rng_state_ptr, int64_t *seed_ptr,
-                                  int64_t *offset_ptr, uint32_t offset_intragraph,
-                                  cudaStream_t stream);
-
 /*!  \brief Copy keys and values into the KV cache.
  *
  * \warning   This API is **experimental** and subject to change.
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index c3503d0f1bf..2da9cac28f3 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -54,7 +54,7 @@
     META_QKV,
 )
 from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
-from transformer_engine.pytorch.distributed import get_distributed_world_size
+from transformer_engine.pytorch.distributed import get_distributed_world_size, _check_checkpoint_rng
 from transformer_engine.pytorch.attention.custom_ops import (
     convert_bshd_to_thd,
     convert_thd_to_bshd,
@@ -1047,6 +1047,9 @@ def forward(
     ) -> torch.Tensor:
         """flash-attn fprop"""
 
+        if self.training and self.attention_dropout > 0:
+            _check_checkpoint_rng()
+
         assert all(
             x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, Float8Tensor)
             for x in [query_layer, key_layer, value_layer]
@@ -2724,6 +2727,10 @@ def forward(
         load_balancing_strategy=CPLoadBalancingStrategy.DUAL_CHUNK_SWAP,
     ) -> torch.Tensor:
         """fused attention fprop"""
+
+        if self.training and self.attention_dropout > 0:
+            _check_checkpoint_rng()
+
         assert (
             fused_attention_backend != FusedAttnBackend["No_Backend"]
         ), "No fused attention backend supports this input combination!"
diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index 06259be9aab..d85dcda1593 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -6,8 +6,6 @@
 
 #include "common.h"
 
-#include <c10/cuda/CUDAGuard.h>
-
 #include "c10/util/ArrayRef.h"
 #include "pybind.h"
 #include "transformer_engine/transformer_engine.h"
@@ -366,40 +364,4 @@ at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_pe
   return philox_args;
 }
 
-at::Tensor checkpoint_rng_state(const at::Generator& generator) {
-  auto* gen = at::check_generator<at::CUDAGeneratorImpl>(generator);
-  const c10::cuda::CUDAGuard device_guard(generator.device());
-  if (at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None) {
-    std::lock_guard<std::mutex> lock(gen->mutex_);
-    return generator.get_state();
-  }
-  auto state = at::empty({2}, at::TensorOptions().dtype(at::kLong).device(generator.device()));
-  philox_unpack(init_philox_state(gen, 0), state.data_ptr<int64_t>());
-  return state;
-}
-
-void restore_checkpoint_rng_state(const at::Generator& generator, const at::Tensor& state) {
-  auto* gen = at::check_generator<at::CUDAGeneratorImpl>(generator);
-  const c10::cuda::CUDAGuard device_guard(generator.device());
-  const bool capturing = at::cuda::currentStreamCaptureStatus() != at::cuda::CaptureStatus::None;
-  TORCH_CHECK(
-      state.is_cuda() == capturing,
-      "Checkpoint forward and recompute must both use eager execution or both be captured.");
-  if (!capturing) {
-    std::lock_guard<std::mutex> lock(gen->mutex_);
-    auto mutable_generator = generator;
-    mutable_generator.set_state(state);
-    return;
-  }
-  TORCH_CHECK(state.device() == generator.device() && state.scalar_type() == at::kLong &&
-                  state.numel() == 2 && state.is_contiguous(),
-              "Expected a contiguous int64 CUDA RNG snapshot with two elements.");
-  auto philox = init_philox_state(gen, 0);
-  // Recompute has its own intragraph position; preserve the saved effective offset.
-  NVTE_SCOPED_GIL_RELEASE({
-    nvte_restore_seed_and_offset(state.data_ptr<int64_t>(), philox.seed_.ptr, philox.offset_.ptr,
-                                 philox.offset_intragraph_, at::cuda::getCurrentCUDAStream());
-  });
-}
-
 }  // namespace transformer_engine::pytorch
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
index 859e0562c30..9a3c6314054 100644
--- a/transformer_engine/pytorch/csrc/common.h
+++ b/transformer_engine/pytorch/csrc/common.h
@@ -629,9 +629,6 @@ void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr);
 // extract PhiloxCudaState from CUDA random number generator
 at::PhiloxCudaState init_philox_state(at::CUDAGeneratorImpl* gen, size_t elts_per_thread);
 
-at::Tensor checkpoint_rng_state(const at::Generator& generator);
-void restore_checkpoint_rng_state(const at::Generator& generator, const at::Tensor& state);
-
 }  // namespace transformer_engine::pytorch
 
 namespace std {
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index 41d468ca105..6fbed1436fb 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -196,11 +196,6 @@ void bind_quantize_with_amax_extensions(py::module_ &m) {
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m)
 
-  m.def("checkpoint_rng_state", &transformer_engine::pytorch::checkpoint_rng_state,
-        py::arg("generator"));
-  m.def("restore_checkpoint_rng_state", &transformer_engine::pytorch::restore_checkpoint_rng_state,
-        py::arg("generator"), py::arg("state"));
-
   // Register __eq__/__ne__ on the pybind ``DType`` enum so it compares by integer.
   py::object dtype_class = m.attr("DType");
   dtype_class.attr("__eq__") = py::cpp_function(
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index b8ea8ff0e11..132dfb01c9b 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -17,6 +17,7 @@
 import torch
 from torch.cuda import _lazy_call, _lazy_init
 from torch.utils.checkpoint import detach_variable, noop_context_fn
+from torch.utils._python_dispatch import TorchDispatchMode, _get_current_dispatch_mode_stack
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp._common_utils import _get_module_fsdp_state
 from torch.distributed.fsdp._traversal_utils import _get_fsdp_states_with_modules
@@ -231,6 +232,8 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F
 
     def __enter__(self):
         global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE
+        self.previous_state = (_IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE)
+        self.saved_state_count = len(activation_recompute_forward._is_first_fp8_module)
         # Track the checkpoint region independently of the FP8 state at entry.
         # A checkpointed callable may open its own FP8 autocast context (for
         # example, to select precision per layer). Delayed-scaling modules in
@@ -247,8 +250,9 @@ def __enter__(self):
 
     def __exit__(self, *exc_details):
         global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE
-        _IN_ACTIVATION_RECOMPUTE_REGION = False
-        _ACTIVATION_RECOMPUTE_PHASE = False
+        if exc_details[0] is not None and self.activation_recompute and not self.recompute_phase:
+            del activation_recompute_forward._is_first_fp8_module[self.saved_state_count :]
+        _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE = self.previous_state
 
 
 def is_fp8_activation_recompute_enabled() -> bool:
@@ -302,6 +306,65 @@ def _get_active_autocast_contexts():
     return gpu_autocast_ctx, cpu_autocast_ctx
 
 
+def _checkpoint_rng_available():
+    """Whether captured Philox values are exposed by this PyTorch build."""
+    return hasattr(torch.Generator, "philox_state")
+
+
+_CHECKPOINT_RNG_ERROR = (
+    "Checkpointing dropout or other random operations during CUDA graph capture requires "
+    "torch.Generator.philox_state (available in PyTorch 2.14+). "
+    "Upgrade PyTorch, disable dropout, or disable CUDA graph capture."
+)
+
+
+class _CheckpointRNGGuard(TorchDispatchMode):
+    """Reject captured RNG use when its state cannot be preserved."""
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        if torch.Tag.nondeterministic_seeded in func.tags:
+            raise RuntimeError(_CHECKPOINT_RNG_ERROR)
+        return func(*args, **(kwargs or {}))
+
+
+def _check_checkpoint_rng():
+    """Check native TE RNG consumers that bypass the PyTorch dispatcher."""
+    if not _checkpoint_rng_available() and any(
+        isinstance(mode, _CheckpointRNGGuard) for mode in _get_current_dispatch_mode_stack()
+    ):
+        raise RuntimeError(_CHECKPOINT_RNG_ERROR)
+
+
+def _checkpoint_rng_state(generator):
+    """Snapshot RNG values on their current device."""
+    with torch.cuda.device(generator.device):
+        if not torch.cuda.is_current_stream_capturing():
+            return generator.get_state()
+        if not _checkpoint_rng_available():
+            return torch.empty(0, device=generator.device, dtype=torch.int64)
+        seed, offset, intragraph = generator.philox_state(0)
+        state = torch.cat((seed, offset))
+        state[1].add_(intragraph.item())
+        return state
+
+
+def _restore_checkpoint_rng_state(generator, state):
+    """Restore a snapshot without moving captured state to the CPU."""
+    with torch.cuda.device(generator.device):
+        capturing = torch.cuda.is_current_stream_capturing()
+        if state.is_cuda != capturing:
+            raise RuntimeError(
+                "Checkpoint forward and recompute must both use eager execution or both be"
+                " captured."
+            )
+        if not capturing:
+            generator.set_state(state)
+        elif state.numel():
+            seed, offset, intragraph = generator.philox_state(0)
+            seed.copy_(state[:1])
+            offset.copy_(state[1:] - intragraph.item())
+
+
 class _CheckpointRNGState:
     """Save RNG values separately from live generator handles."""
 
@@ -319,7 +382,7 @@ def __init__(self, get_rng_state_tracker):
 
     def snapshot(self):
         """Copy the CPU state, generator values, and tensor-based tracker states."""
-        states = [torch.get_rng_state(), *(tex.checkpoint_rng_state(g) for g in self.generators)]
+        states = [torch.get_rng_state(), *(_checkpoint_rng_state(g) for g in self.generators)]
         if self.tracker is not None:
             tracker_states = self.tracker.get_states()
             states.extend(
@@ -330,7 +393,13 @@ def snapshot(self):
     def _restore(self, states):
         torch.set_rng_state(states[0])
         for generator, state in zip(self.generators, states[1:]):
-            tex.restore_checkpoint_rng_state(generator, state)
+            _restore_checkpoint_rng_state(generator, state)
+
+    def guard(self):
+        """Keep deterministic capture available on older PyTorch builds."""
+        if torch.cuda.is_current_stream_capturing() and not _checkpoint_rng_available():
+            return _CheckpointRNGGuard()
+        return nullcontext()
 
     @contextmanager
     def fork(self, states):
@@ -339,7 +408,7 @@ def fork(self, states):
         tracker_states = self.tracker.get_states() if self.tracker is not None else None
         live_states = [
             torch.get_rng_state(),
-            *(tex.checkpoint_rng_state(g) for g in self.generators),
+            *(_checkpoint_rng_state(g) for g in self.generators),
         ]
         try:
             self._restore(states)
@@ -357,7 +426,8 @@ def fork(self, states):
                         for name, state in self.tracker_states.items()
                     }
                 )
-            yield
+            with self.guard():
+                yield
         finally:
             self._restore(live_states)
             _set_cuda_rng_state(cuda_state)
@@ -406,7 +476,7 @@ def forward(
             if any(ctx.needs_input_grad)
             else nullcontext()
         )
-        with torch.no_grad(), forward_ctx, activation_recompute_ctx:
+        with torch.no_grad(), forward_ctx, activation_recompute_ctx, ctx.rng_state.guard():
             outputs = run_function(*args, **kwargs)
 
         # Divide hidden states across model parallel group and only keep
@@ -671,6 +741,10 @@ def checkpoint(
         `use_reentrant=False` does not support early stopping, and will execute the entire forward
         pass for the checkpointed module when recomputing activations in the backward pass.
 
+    .. note::
+        Captured dropout and other random operations require ``torch.Generator.philox_state``
+        (PyTorch 2.14+). Older builds support deterministic capture and eager checkpointing.
+
     Parameters
     ----------
     function : Callable
@@ -723,7 +797,9 @@ def checkpoint(
     context_fn = kwargs.pop("context_fn", noop_context_fn)
     determinism_check = kwargs.pop("determinism_check", "default")
     debug = kwargs.pop("debug", False)
-    if not has_te_modules(function):
+    if not has_te_modules(function) and (
+        not torch.cuda.is_initialized() or not torch.cuda.is_current_stream_capturing()
+    ):
         return torch.utils.checkpoint.checkpoint(
             function,
             *args,
@@ -808,7 +884,12 @@ def recompute_fn(*args, **kwargs):
         recompute_fn,
         get_rng_state_tracker,
     )
-    with _checkpoint_hook(new_frame, args, kwargs), te_forward_ctx, user_forward_ctx:
+    with (
+        _checkpoint_hook(new_frame, args, kwargs),
+        te_forward_ctx,
+        user_forward_ctx,
+        new_frame.rng_state.guard(),
+    ):
         out = function(*args, **kwargs)
 
     return out
diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py
index 206f86cafc5..24c0248467a 100644
--- a/transformer_engine/pytorch/graph.py
+++ b/transformer_engine/pytorch/graph.py
@@ -85,9 +85,11 @@ def _none_grad_context_wrapper(inputs):
     for input_tensor in inputs:
         original_input_grads.append(input_tensor.grad)
         input_tensor.grad = None
-    yield
-    for input_tensor, original_grad in zip(inputs, original_input_grads):
-        input_tensor.grad = original_grad
+    try:
+        yield
+    finally:
+        for input_tensor, original_grad in zip(inputs, original_input_grads):
+            input_tensor.grad = original_grad
 
 
 @contextlib.contextmanager
@@ -103,10 +105,12 @@ def _graph_context_wrapper(*args, **kwargs):
     gc_is_enabled = gc.isenabled()
     if gc_is_enabled:
         gc.disable()
-    with torch.cuda.graph(*args, **kwargs):
-        yield
-    if gc_is_enabled:
-        gc.enable()
+    try:
+        with torch.cuda.graph(*args, **kwargs):
+            yield
+    finally:
+        if gc_is_enabled:
+            gc.enable()
 
 
 def _canonicalize_capture_time_hooks(
@@ -1717,34 +1721,36 @@ def call_func(self, *args, **kwargs):
     ]
     original_rng_states = [state.get_state() for state in generators]
 
-    graphed_callables = _make_graphed_callables(
-        forward_funcs,
-        sample_args,
-        num_warmup_iters=num_warmup_iters,
-        allow_unused_input=allow_unused_input,
-        cache_quantized_params=cache_quantized_params,
-        sample_kwargs=sample_kwargs,
-        _order=_order,
-        _num_layers_per_chunk=_num_layers_per_chunk,
-        pool=pool,
-        retain_graph_in_backward=retain_graph_in_backward,
-        _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers,
-        clone_param_grads_on_return=clone_param_grads_on_return,
-        pre_warmup_hook=pre_warmup_hook,
-        post_warmup_hook=post_warmup_hook,
-        capture_time_hooks=capture_time_hooks,
-    )
+    try:
+        graphed_callables = _make_graphed_callables(
+            forward_funcs,
+            sample_args,
+            num_warmup_iters=num_warmup_iters,
+            allow_unused_input=allow_unused_input,
+            cache_quantized_params=cache_quantized_params,
+            sample_kwargs=sample_kwargs,
+            _order=_order,
+            _num_layers_per_chunk=_num_layers_per_chunk,
+            pool=pool,
+            retain_graph_in_backward=retain_graph_in_backward,
+            _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers,
+            clone_param_grads_on_return=clone_param_grads_on_return,
+            pre_warmup_hook=pre_warmup_hook,
+            post_warmup_hook=post_warmup_hook,
+            capture_time_hooks=capture_time_hooks,
+        )
+    finally:
+        # Ensures warmup does not affect numerics for ops such as dropout.
+        for gen, state in zip(generators, original_rng_states):
+            gen.set_state(state)
 
-    # Ensures warmup does not affect numerics for ops such as dropout.
-    for gen, state in zip(generators, original_rng_states):
-        gen.set_state(state)
+        # Remove FP8 wrapper.
+        for module_cls, old_call in old_call_funcs.items():
+            module_cls.__call__ = old_call
 
-    # Remove FP8 wrapper.
-    for module_cls, old_call in old_call_funcs.items():
-        module_cls.__call__ = old_call
+        # Restore FP8 state.
+        restore_fp8_tensors(modules, saved_fp8_tensors)
 
-    # Restore FP8 state.
-    restore_fp8_tensors(modules, saved_fp8_tensors)
+        set_capture_end()
 
-    set_capture_end()
     return graphed_callables
diff --git a/transformer_engine/pytorch/ops/basic/dropout.py b/transformer_engine/pytorch/ops/basic/dropout.py
index 8850604aadc..df4947f2c2a 100644
--- a/transformer_engine/pytorch/ops/basic/dropout.py
+++ b/transformer_engine/pytorch/ops/basic/dropout.py
@@ -10,6 +10,7 @@
 import torch
 import transformer_engine_torch as tex
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload
+from ...distributed import _check_checkpoint_rng
 from ...tensor import Quantizer
 from ...tensor.storage.float8_tensor_storage import Float8TensorStorage
 from .._common import maybe_autocast_dtype, maybe_dequantize
@@ -40,9 +41,12 @@ def op_forward(
         # Output dtype
         dtype = maybe_autocast_dtype(default_dtype=input_.dtype)
 
+        if self.training and self.dropout_probability > 0:
+            _check_checkpoint_rng()
+
         # Choose implementation
         impl = None
-        if not self.training:
+        if not self.training or self.dropout_probability == 0:
             impl = "evaluation"
         elif input_.numel() % 16 == 0 and dtype in (torch.float16, torch.bfloat16):
             impl = "fused"

From 1149f52a79e49391b43a695c46988c49ecfa8ef2 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 12:50:49 +0100
Subject: [PATCH 05/23] [PyTorch] Guard captured stochastic NVFP4 quantization

Check stochastic rounding when native quantizers read their configuration so direct and fused quantization paths cannot bypass the checkpoint RNG guard. Preserve the serialized field name, fix the missing test import, and cover both checkpoint modes through regular pytest.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py                | 27 +++++++++++++++++++
 .../pytorch/tensor/nvfp4_tensor.py            | 16 +++++++++++
 2 files changed, 43 insertions(+)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index dd99c8de9cd..7c8cbf1094b 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -12,6 +12,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn import Parameter
+import transformer_engine.pytorch as te
 
 from transformer_engine.pytorch.quantization import (
     FP8GlobalStateManager,
@@ -1359,6 +1360,32 @@ def body(query):
             te_checkpoint(body, value + 0, use_reentrant=use_reentrant)
 
 
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+@pytest.mark.parametrize("entry", ("quantize", "update", "native", "fused"))
+def test_checkpoint_stochastic_nvfp4_requires_philox_api(use_reentrant, entry, monkeypatch):
+    """Stochastic quantization must reject capture even through native entry points."""
+    import transformer_engine_torch as tex
+
+    monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
+    quantizer = te.NVFP4Quantizer(rowwise=True, columnwise=False, stochastic_rounding=True)
+    value = torch.ones(32, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    destination = quantizer.make_empty(value.shape, dtype=value.dtype, device=value.device)
+
+    def body(tensor):
+        if entry == "quantize":
+            return quantizer(tensor)
+        if entry == "update":
+            return quantizer.update_quantized(tensor, destination)
+        if entry == "native":
+            return tex.quantize(tensor, quantizer)
+        return tex.swiglu(tensor, quantizer)
+
+    graph = torch.cuda.CUDAGraph()
+    with pytest.raises(RuntimeError, match="philox_state"):
+        with torch.cuda.graph(graph):
+            te_checkpoint(body, value + 0, use_reentrant=use_reentrant)
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 78dea6cbee0..6b82def877c 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -206,6 +206,22 @@ def __getstate__(self):
         state["amax_reduction_group"] = None
         return state
 
+    @property
+    def stochastic_rounding(self) -> bool:
+        """Whether native quantization may use stochastic rounding in this context."""
+        enabled = self.__dict__["stochastic_rounding"]
+        if enabled:
+            # Native and fused quantizers read this flag without calling quantize_impl.
+            from ..distributed import _check_checkpoint_rng
+
+            _check_checkpoint_rng()
+        return enabled
+
+    @stochastic_rounding.setter
+    def stochastic_rounding(self, enabled: bool) -> None:
+        # Preserve the field name in serialized quantizers.
+        self.__dict__["stochastic_rounding"] = enabled
+
     def _rebuild_derived_state(self) -> None:
         """Build the derived ``rht_matrix`` (also used after value-key reconstruction).
 

From 24f7ce58946287873475630f1acbe9cc1dc88a4d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 14:38:11 +0100
Subject: [PATCH 06/23] Trim redundant checkpoint RNG test coverage

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py | 29 ++++++-----------------------
 1 file changed, 6 insertions(+), 23 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 7c8cbf1094b..66cec6261d0 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1206,9 +1206,8 @@ def body(value):
         assert tracker.get_states()["model-parallel-rng"] is tracker_state
 
 
-@pytest.mark.parametrize("use_reentrant", all_boolean)
 @pytest.mark.parametrize("fork_tracker", all_boolean)
-def test_checkpoint_preserves_captured_rng_state(fork_tracker, use_reentrant):
+def test_checkpoint_preserves_captured_rng_state(fork_tracker):
     """Graph replay must still advance its registered RNG after eager recompute."""
     tracker = CudaRNGStatesTracker()
     tracker.add("model-parallel-rng", 4321)
@@ -1237,9 +1236,7 @@ def body(value):
             return F.dropout(value, p=0.5, training=True)
 
     inp = torch.ones(4096, device="cuda", requires_grad=True)
-    out = te_checkpoint(
-        body, inp, use_reentrant=use_reentrant, get_rng_state_tracker=lambda: tracker
-    )
+    out = te_checkpoint(body, inp, use_reentrant=True, get_rng_state_tracker=lambda: tracker)
     out.sum().backward()
     assert torch.equal(inp.grad, out)
 
@@ -1258,12 +1255,6 @@ def test_checkpoint_dropout_gradients_during_graph_replay(use_reentrant, rng_sou
     torch.manual_seed(seed)
     tracker = CudaRNGStatesTracker()
     tracker.add("model-parallel-rng", 4321)
-    saved_gpu_states = []
-
-    def pack(tensor):
-        if tensor.is_cuda and tensor.dtype == torch.int64 and tensor.shape == (2,):
-            saved_gpu_states.append(torch.cuda.is_current_stream_capturing())
-        return tensor
 
     def body(value):
         if rng_source == "tracker":
@@ -1286,8 +1277,7 @@ def forward(self, value):
                 )
 
     sample = torch.ones(4096, device="cuda", requires_grad=True)
-    with torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor):
-        model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
+    model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
     previous = None
     for _ in range(3):
         value = torch.ones_like(sample, requires_grad=True)
@@ -1298,7 +1288,6 @@ def forward(self, value):
         if previous is not None:
             assert not torch.equal(forward_output, previous), "Graph replay did not advance RNG"
         previous = forward_output
-    assert any(saved_gpu_states), "Captured RNG snapshots must use autograd's saved-tensor hooks"
 
 
 @pytest.mark.parametrize("use_reentrant", all_boolean)
@@ -1341,9 +1330,8 @@ def forward(self, value):
         assert torch.equal(value.grad, torch.full_like(value, 2))
 
 
-@pytest.mark.parametrize("use_reentrant", all_boolean)
 @pytest.mark.parametrize("backend", ("FusedAttention", "FlashAttention"))
-def test_checkpoint_native_attention_requires_philox_api(use_reentrant, backend, monkeypatch):
+def test_checkpoint_native_attention_requires_philox_api(backend, monkeypatch):
     """Native attention dropout must reject capture before entering its CUDA backend."""
     from transformer_engine.pytorch.attention.dot_product_attention import backends
 
@@ -1357,11 +1345,11 @@ def body(query):
     graph = torch.cuda.CUDAGraph()
     with pytest.raises(RuntimeError, match="philox_state"):
         with torch.cuda.graph(graph):
-            te_checkpoint(body, value + 0, use_reentrant=use_reentrant)
+            te_checkpoint(body, value + 0, use_reentrant=True)
 
 
 @pytest.mark.parametrize("use_reentrant", all_boolean)
-@pytest.mark.parametrize("entry", ("quantize", "update", "native", "fused"))
+@pytest.mark.parametrize("entry", ("native", "fused"))
 def test_checkpoint_stochastic_nvfp4_requires_philox_api(use_reentrant, entry, monkeypatch):
     """Stochastic quantization must reject capture even through native entry points."""
     import transformer_engine_torch as tex
@@ -1369,13 +1357,8 @@ def test_checkpoint_stochastic_nvfp4_requires_philox_api(use_reentrant, entry, m
     monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
     quantizer = te.NVFP4Quantizer(rowwise=True, columnwise=False, stochastic_rounding=True)
     value = torch.ones(32, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-    destination = quantizer.make_empty(value.shape, dtype=value.dtype, device=value.device)
 
     def body(tensor):
-        if entry == "quantize":
-            return quantizer(tensor)
-        if entry == "update":
-            return quantizer.update_quantized(tensor, destination)
         if entry == "native":
             return tex.quantize(tensor, quantizer)
         return tex.swiglu(tensor, quantizer)

From 8afc77a50f19c738dc18abf303abe16a18d9183e Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 14:42:49 +0100
Subject: [PATCH 07/23] Expose custom op tags to checkpoint RNG guards

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 51 +++++++++++++++++++
 .../pytorch/dynamo/custom_op.py               |  4 ++
 2 files changed, 55 insertions(+)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 6868de7ef18..934d66a66d1 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -60,6 +60,8 @@
     Quantizer,
 )
 from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec
+from transformer_engine.pytorch.dynamo.custom_op import register_custom_op
+from transformer_engine.pytorch.distributed import _CheckpointRNGGuard
 from transformer_engine.pytorch.dynamo.fake import fake_bgrad
 from transformer_engine.pytorch import (
     is_fp8_available,
@@ -2785,6 +2787,55 @@ def fuse(ops, **unused):
         _assert_custom_ops(graphs, "_affineop", backward=False)
 
 
+@pytest.mark.parametrize("seeded", [False, True])
+def test_custom_op_rng_tags(seeded, monkeypatch):
+    @dataclasses.dataclass
+    class Args:
+        inp: torch.Tensor
+
+    calls = []
+
+    def impl(args):
+        calls.append(True)
+        return torch.rand_like(args.inp)
+
+    name = f"test_checkpoint_rng_tags_{seeded}"
+    op = register_custom_op(
+        op_name=name,
+        arg_type=Args,
+        impl=impl,
+        fake_impl=lambda args: args.inp,
+        tags=(torch.Tag.nondeterministic_seeded,) if seeded else (),
+    )
+    assert op is not None
+    for suffix in ("", "_base"):
+        overload = getattr(torch.ops.transformer_engine_compile, name + suffix).default
+        assert (torch.Tag.nondeterministic_seeded in overload.tags) == seeded
+
+    args = Args(torch.ones(16))
+    op(args)
+    assert len(calls) == 1
+    with _CheckpointRNGGuard():
+        if seeded:
+            with pytest.raises(RuntimeError, match="philox_state"):
+                op(args)
+            assert len(calls) == 1
+        else:
+            op(args)
+            assert len(calls) == 2
+
+    if seeded:
+        monkeypatch.setattr(
+            "transformer_engine.pytorch.distributed._checkpoint_rng_available", lambda: False
+        )
+        value = torch.ones(16, device="cuda", requires_grad=True)
+        graph = torch.cuda.CUDAGraph()
+        with pytest.raises(RuntimeError, match="philox_state"):
+            with torch.cuda.graph(graph):
+                te.checkpoint(lambda x: op(Args(x)), value + 0, use_reentrant=True)
+        assert len(calls) == 1
+
+
 @pytest.mark.parametrize(
     "forward,backward", [(True, True), (True, False), (False, True), (False, False)]
 )
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 1ed84017ecd..bd49f4b84cf 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1274,12 +1274,15 @@ def register_custom_op(
     arg_type: type,
     impl: Callable[[Any], Any],
     fake_impl: Callable[[Any], Any],
+    tags: Sequence[torch.Tag] = (),
 ) -> Optional[Callable[[Any], Any]]:
     """Register one custom op without autograd wiring.
 
     arg_type is a dataclass defining the input schema. Results may be nested
     tuples/lists of fresh tensors or None; fake_impl mirrors them with TensorSpec.
     The returned callable takes an args instance and preserves the result structure.
+    ``tags`` apply to both base and wrapper operators; mark RNG consumers with
+    ``torch.Tag.nondeterministic_seeded``.
     Returns None for unsupported registration APIs so callers can fall back to eager.
     Duplicate operator names raise ValueError.
     """
@@ -1295,6 +1298,7 @@ def pack_result(result):
             impl=impl,
             fake_impl=fake_impl,
             pack_result=pack_result,
+            tags=tags,
         )
     except (ImportError, AttributeError, RuntimeError, TypeError) as e:
         record_compile_disabled(

From 5fa09af938c180b26d4b01b5b22d9c56a03e14a3 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 14:44:57 +0100
Subject: [PATCH 08/23] Clarify the missing Philox state sentinel

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/distributed.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 132dfb01c9b..045f65a0504 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -341,6 +341,7 @@ def _checkpoint_rng_state(generator):
         if not torch.cuda.is_current_stream_capturing():
             return generator.get_state()
         if not _checkpoint_rng_available():
+            # Empty sentinel: the guard rejects RNG use without Philox state access.
             return torch.empty(0, device=generator.device, dtype=torch.int64)
         seed, offset, intragraph = generator.philox_state(0)
         state = torch.cat((seed, offset))

From 41a4afc61955442fa24ad4b8792ce8982ba93cd1 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 15:18:13 +0100
Subject: [PATCH 09/23] Restore captured RNG bases and CPU offsets separately

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py            | 36 +++++++++++++++++++----
 transformer_engine/pytorch/distributed.py | 35 +++++++++++-----------
 2 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 66cec6261d0..3e8f2d0bddc 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1247,10 +1247,11 @@ def body(value):
 
 @pytest.mark.parametrize("use_reentrant", all_boolean)
 @pytest.mark.parametrize("rng_source", ("default", "tracker", "mixed", "shared", "no_tracker"))
+@pytest.mark.parametrize("capture_mode", ("graphed", "manual"))
 @pytest.mark.skipif(
     not hasattr(torch.Generator, "philox_state"), reason="Captured RNG requires PyTorch 2.14+"
 )
-def test_checkpoint_dropout_gradients_during_graph_replay(use_reentrant, rng_source):
+def test_checkpoint_dropout_gradients_during_graph_replay(use_reentrant, rng_source, capture_mode):
     """Each captured recompute must reproduce its corresponding forward masks."""
     torch.manual_seed(seed)
     tracker = CudaRNGStatesTracker()
@@ -1269,6 +1270,7 @@ def body(value):
     class Model(torch.nn.Module):
         def forward(self, value):
             with tracker.fork() if rng_source == "shared" else nullcontext():
+                torch.rand(8, device=value.device)
                 return te_checkpoint(
                     body,
                     value,
@@ -1277,14 +1279,36 @@ def forward(self, value):
                 )
 
     sample = torch.ones(4096, device="cuda", requires_grad=True)
-    model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
+    if capture_mode == "manual":
+        model = Model()
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                sample.grad = None
+                model(sample).sum().backward()
+        torch.cuda.current_stream().wait_stream(stream)
+        sample.grad = None
+        graph = torch.cuda.CUDAGraph()
+        graph.register_generator_state(tracker.get_states()["model-parallel-rng"])
+        with torch.cuda.graph(graph):
+            output = model(sample)
+            output.sum().backward()
+    else:
+        model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
     previous = None
     for _ in range(3):
-        value = torch.ones_like(sample, requires_grad=True)
-        output = model(value)
+        if capture_mode == "manual":
+            sample.grad.zero_()
+            graph.replay()
+            grad = sample.grad
+        else:
+            value = torch.ones_like(sample, requires_grad=True)
+            output = model(value)
+            output.sum().backward()
+            grad = value.grad
         forward_output = output.detach().clone()
-        output.sum().backward()
-        assert torch.equal(value.grad, forward_output), "Recompute used a different dropout mask"
+        assert torch.equal(grad, forward_output), "Recompute used a different dropout mask"
         if previous is not None:
             assert not torch.equal(forward_output, previous), "Graph replay did not advance RNG"
         previous = forward_output
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 045f65a0504..39787a48c8a 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -336,20 +336,18 @@ def _check_checkpoint_rng():
 
 
 def _checkpoint_rng_state(generator):
-    """Snapshot RNG values on their current device."""
+    """Snapshot RNG values and the capture's CPU counter separately."""
     with torch.cuda.device(generator.device):
         if not torch.cuda.is_current_stream_capturing():
-            return generator.get_state()
+            return generator.get_state(), None
         if not _checkpoint_rng_available():
             # Empty sentinel: the guard rejects RNG use without Philox state access.
-            return torch.empty(0, device=generator.device, dtype=torch.int64)
-        seed, offset, intragraph = generator.philox_state(0)
-        state = torch.cat((seed, offset))
-        state[1].add_(intragraph.item())
-        return state
+            return torch.empty(0, device=generator.device, dtype=torch.int64), None
+        seed, offset, _ = generator.philox_state(0)
+        return torch.cat((seed, offset)), generator.get_state()
 
 
-def _restore_checkpoint_rng_state(generator, state):
+def _restore_checkpoint_rng_state(generator, state, capture_state):
     """Restore a snapshot without moving captured state to the CPU."""
     with torch.cuda.device(generator.device):
         capturing = torch.cuda.is_current_stream_capturing()
@@ -361,9 +359,10 @@ def _restore_checkpoint_rng_state(generator, state):
         if not capturing:
             generator.set_state(state)
         elif state.numel():
-            seed, offset, intragraph = generator.philox_state(0)
+            generator.set_state(capture_state)
+            seed, offset, _ = generator.philox_state(0)
             seed.copy_(state[:1])
-            offset.copy_(state[1:] - intragraph.item())
+            offset.copy_(state[1:])
 
 
 class _CheckpointRNGState:
@@ -383,7 +382,10 @@ def __init__(self, get_rng_state_tracker):
 
     def snapshot(self):
         """Copy the CPU state, generator values, and tensor-based tracker states."""
-        states = [torch.get_rng_state(), *(_checkpoint_rng_state(g) for g in self.generators)]
+        states = [
+            torch.get_rng_state(),
+            *(state for g in self.generators for state in _checkpoint_rng_state(g)),
+        ]
         if self.tracker is not None:
             tracker_states = self.tracker.get_states()
             states.extend(
@@ -393,8 +395,8 @@ def snapshot(self):
 
     def _restore(self, states):
         torch.set_rng_state(states[0])
-        for generator, state in zip(self.generators, states[1:]):
-            _restore_checkpoint_rng_state(generator, state)
+        for index, generator in enumerate(self.generators):
+            _restore_checkpoint_rng_state(generator, *states[1 + 2 * index : 3 + 2 * index])
 
     def guard(self):
         """Keep deterministic capture available on older PyTorch builds."""
@@ -407,10 +409,7 @@ def fork(self, states):
         """Restore forward RNG values temporarily, then restore the live values."""
         cuda_state = _get_cuda_rng_state()
         tracker_states = self.tracker.get_states() if self.tracker is not None else None
-        live_states = [
-            torch.get_rng_state(),
-            *(_checkpoint_rng_state(g) for g in self.generators),
-        ]
+        live_states = self.snapshot()
         try:
             self._restore(states)
             _set_cuda_rng_state(self.cuda_state)
@@ -418,7 +417,7 @@ def fork(self, states):
                 tensor_states = dict(
                     zip(
                         (name for name, state in self.tracker_states.items() if state is None),
-                        states[1 + len(self.generators) :],
+                        states[1 + 2 * len(self.generators) :],
                     )
                 )
                 self.tracker.set_states(

From 52f18953e7979dadf19fc39b5fac52ea80c42aa2 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 15:40:06 +0100
Subject: [PATCH 10/23] Remove the RNG tracker snapshot note

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/distributed.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 39787a48c8a..841e42434d1 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -924,11 +924,6 @@ def get_states(self) -> Dict[str, torch.Generator]:
         Get rng states. Copy the dictionary so we have direct pointers
         to the states, not just a pointer to the dictionary.
 
-        .. note::
-            Graph-safe RNG states are live ``torch.Generator`` handles rather
-            than snapshots. They advance in place as random numbers are drawn;
-            call ``Generator.get_state()`` to save the current seed and offset
-            outside CUDA graph capture.
         """
         states = {}
         for name in self.states_:

From 6dab9c709d40eb6996b1ec27018a39d543cba84f Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 15:48:38 +0100
Subject: [PATCH 11/23] Remove the custom op RNG tag test

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 51 -------------------
 .../pytorch/tensor/nvfp4_tensor.py            | 16 ------
 2 files changed, 67 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 934d66a66d1..6868de7ef18 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -60,8 +60,6 @@
     Quantizer,
 )
 from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec
-from transformer_engine.pytorch.dynamo.custom_op import register_custom_op
-from transformer_engine.pytorch.distributed import _CheckpointRNGGuard
 from transformer_engine.pytorch.dynamo.fake import fake_bgrad
 from transformer_engine.pytorch import (
     is_fp8_available,
@@ -2787,55 +2785,6 @@ def fuse(ops, **unused):
         _assert_custom_ops(graphs, "_affineop", backward=False)
 
 
-@pytest.mark.parametrize("seeded", [False, True])
-def test_custom_op_rng_tags(seeded, monkeypatch):
-    @dataclasses.dataclass
-    class Args:
-        inp: torch.Tensor
-
-    calls = []
-
-    def impl(args):
-        calls.append(True)
-        return torch.rand_like(args.inp)
-
-    name = f"test_checkpoint_rng_tags_{seeded}"
-    op = register_custom_op(
-        op_name=name,
-        arg_type=Args,
-        impl=impl,
-        fake_impl=lambda args: args.inp,
-        tags=(torch.Tag.nondeterministic_seeded,) if seeded else (),
-    )
-    assert op is not None
-    for suffix in ("", "_base"):
-        overload = getattr(torch.ops.transformer_engine_compile, name + suffix).default
-        assert (torch.Tag.nondeterministic_seeded in overload.tags) == seeded
-
-    args = Args(torch.ones(16))
-    op(args)
-    assert len(calls) == 1
-    with _CheckpointRNGGuard():
-        if seeded:
-            with pytest.raises(RuntimeError, match="philox_state"):
-                op(args)
-            assert len(calls) == 1
-        else:
-            op(args)
-            assert len(calls) == 2
-
-    if seeded:
-        monkeypatch.setattr(
-            "transformer_engine.pytorch.distributed._checkpoint_rng_available", lambda: False
-        )
-        value = torch.ones(16, device="cuda", requires_grad=True)
-        graph = torch.cuda.CUDAGraph()
-        with pytest.raises(RuntimeError, match="philox_state"):
-            with torch.cuda.graph(graph):
-                te.checkpoint(lambda x: op(Args(x)), value + 0, use_reentrant=True)
-        assert len(calls) == 1
-
-
 @pytest.mark.parametrize(
     "forward,backward", [(True, True), (True, False), (False, True), (False, False)]
 )
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 6b82def877c..78dea6cbee0 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -206,22 +206,6 @@ def __getstate__(self):
         state["amax_reduction_group"] = None
         return state
 
-    @property
-    def stochastic_rounding(self) -> bool:
-        """Whether native quantization may use stochastic rounding in this context."""
-        enabled = self.__dict__["stochastic_rounding"]
-        if enabled:
-            # Native and fused quantizers read this flag without calling quantize_impl.
-            from ..distributed import _check_checkpoint_rng
-
-            _check_checkpoint_rng()
-        return enabled
-
-    @stochastic_rounding.setter
-    def stochastic_rounding(self, enabled: bool) -> None:
-        # Preserve the field name in serialized quantizers.
-        self.__dict__["stochastic_rounding"] = enabled
-
     def _rebuild_derived_state(self) -> None:
         """Build the derived ``rht_matrix`` (also used after value-key reconstruction).
 

From 65ab410a722a3285b7a749982f54215cad5526a1 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 15:49:25 +0100
Subject: [PATCH 12/23] Restore the unrelated NVFP4 guard change

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/tensor/nvfp4_tensor.py               | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 78dea6cbee0..6b82def877c 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -206,6 +206,22 @@ def __getstate__(self):
         state["amax_reduction_group"] = None
         return state
 
+    @property
+    def stochastic_rounding(self) -> bool:
+        """Whether native quantization may use stochastic rounding in this context."""
+        enabled = self.__dict__["stochastic_rounding"]
+        if enabled:
+            # Native and fused quantizers read this flag without calling quantize_impl.
+            from ..distributed import _check_checkpoint_rng
+
+            _check_checkpoint_rng()
+        return enabled
+
+    @stochastic_rounding.setter
+    def stochastic_rounding(self, enabled: bool) -> None:
+        # Preserve the field name in serialized quantizers.
+        self.__dict__["stochastic_rounding"] = enabled
+
     def _rebuild_derived_state(self) -> None:
         """Build the derived ``rht_matrix`` (also used after value-key reconstruction).
 

From 82c6e72891ead7a64a6bfc998184ae7684e79ade Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 15:53:09 +0100
Subject: [PATCH 13/23] Remove the NVFP4 stochastic rounding property guard

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/tensor/nvfp4_tensor.py               | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 6b82def877c..78dea6cbee0 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -206,22 +206,6 @@ def __getstate__(self):
         state["amax_reduction_group"] = None
         return state
 
-    @property
-    def stochastic_rounding(self) -> bool:
-        """Whether native quantization may use stochastic rounding in this context."""
-        enabled = self.__dict__["stochastic_rounding"]
-        if enabled:
-            # Native and fused quantizers read this flag without calling quantize_impl.
-            from ..distributed import _check_checkpoint_rng
-
-            _check_checkpoint_rng()
-        return enabled
-
-    @stochastic_rounding.setter
-    def stochastic_rounding(self, enabled: bool) -> None:
-        # Preserve the field name in serialized quantizers.
-        self.__dict__["stochastic_rounding"] = enabled
-
     def _rebuild_derived_state(self) -> None:
         """Build the derived ``rht_matrix`` (also used after value-key reconstruction).
 

From 1e8baf3f0fc0f3c36452628534f1da5f47311462 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 15:56:54 +0100
Subject: [PATCH 14/23] Integrate checkpoint RNG guard and snapshot
 simplifications

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py            | 21 -----------
 transformer_engine/pytorch/distributed.py | 44 +++++++----------------
 2 files changed, 12 insertions(+), 53 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 3e8f2d0bddc..058f53d23e8 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1372,27 +1372,6 @@ def body(query):
             te_checkpoint(body, value + 0, use_reentrant=True)
 
 
-@pytest.mark.parametrize("use_reentrant", all_boolean)
-@pytest.mark.parametrize("entry", ("native", "fused"))
-def test_checkpoint_stochastic_nvfp4_requires_philox_api(use_reentrant, entry, monkeypatch):
-    """Stochastic quantization must reject capture even through native entry points."""
-    import transformer_engine_torch as tex
-
-    monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
-    quantizer = te.NVFP4Quantizer(rowwise=True, columnwise=False, stochastic_rounding=True)
-    value = torch.ones(32, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-
-    def body(tensor):
-        if entry == "native":
-            return tex.quantize(tensor, quantizer)
-        return tex.swiglu(tensor, quantizer)
-
-    graph = torch.cuda.CUDAGraph()
-    with pytest.raises(RuntimeError, match="philox_state"):
-        with torch.cuda.graph(graph):
-            te_checkpoint(body, value + 0, use_reentrant=use_reentrant)
-
-
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 841e42434d1..81c92d899ba 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -17,7 +17,7 @@
 import torch
 from torch.cuda import _lazy_call, _lazy_init
 from torch.utils.checkpoint import detach_variable, noop_context_fn
-from torch.utils._python_dispatch import TorchDispatchMode, _get_current_dispatch_mode_stack
+from torch.utils._python_dispatch import TorchDispatchMode
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp._common_utils import _get_module_fsdp_state
 from torch.distributed.fsdp._traversal_utils import _get_fsdp_states_with_modules
@@ -329,8 +329,10 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
 
 def _check_checkpoint_rng():
     """Check native TE RNG consumers that bypass the PyTorch dispatcher."""
-    if not _checkpoint_rng_available() and any(
-        isinstance(mode, _CheckpointRNGGuard) for mode in _get_current_dispatch_mode_stack()
+    if (
+        _IN_ACTIVATION_RECOMPUTE_REGION
+        and not _checkpoint_rng_available()
+        and torch.cuda.is_current_stream_capturing()
     ):
         raise RuntimeError(_CHECKPOINT_RNG_ERROR)
 
@@ -379,6 +381,7 @@ def __init__(self, get_rng_state_tracker):
                 self.tracker_states[name] = state
             else:
                 self.tracker_states[name] = None
+        self.saved_states = self.snapshot()
 
     def snapshot(self):
         """Copy the CPU state, generator values, and tensor-based tracker states."""
@@ -405,8 +408,9 @@ def guard(self):
         return nullcontext()
 
     @contextmanager
-    def fork(self, states):
+    def fork(self):
         """Restore forward RNG values temporarily, then restore the live values."""
+        states = self.saved_states
         cuda_state = _get_cuda_rng_state()
         tracker_states = self.tracker.get_states() if self.tracker is not None else None
         live_states = self.snapshot()
@@ -460,7 +464,6 @@ def forward(
         ctx.distribute_saved_activations = distribute_saved_activations
 
         ctx.rng_state = _CheckpointRNGState(get_rng_state_tracker)
-        rng_tensors = ctx.rng_state.snapshot()
 
         if context_fn is not None:
             forward_ctx, recompute_ctx = context_fn()
@@ -491,8 +494,7 @@ def forward(
         # Store everything.
         ctx.inputs = [arg if not torch.is_tensor(arg) else None for arg in args]
         tensor_inputs = [arg if torch.is_tensor(arg) else None for arg in args]
-        ctx.num_inputs = len(tensor_inputs)
-        ctx.save_for_backward(*tensor_inputs, *rng_tensors)
+        ctx.save_for_backward(*tensor_inputs)
 
         fp8 = FP8GlobalStateManager.is_fp8_enabled()
         ctx.tp_group = tp_group
@@ -515,10 +517,8 @@ def backward(
                 "Checkpointing is not compatible with .grad(), please use .backward() if possible"
             )
 
-        saved_tensors = ctx.saved_tensors
         inputs = tuple(
-            t if t is not None else arg
-            for t, arg in zip(saved_tensors[: ctx.num_inputs], ctx.inputs)
+            t if t is not None else arg for (t, arg) in zip(ctx.saved_tensors, ctx.inputs)
         )
 
         if ctx.distribute_saved_activations:
@@ -530,7 +530,7 @@ def backward(
         # Compute the forward pass.
         detached_inputs = detach_variable(inputs)
         with (
-            ctx.rng_state.fork(saved_tensors[ctx.num_inputs :]),
+            ctx.rng_state.fork(),
             torch.enable_grad(),
             ctx.recompute_ctx,
             ctx.torch_gpu_amp_ctx,
@@ -564,21 +564,6 @@ def backward(
         return (None, None, None, None, None, None) + grads
 
 
-class _CheckpointRNGStateHolder(torch.autograd.Function):
-    """Route non-reentrant RNG snapshots through autograd's saved-tensor hooks."""
-
-    @staticmethod
-    def forward(ctx, dummy, *states):
-        """Save snapshots without connecting them to the model's backward graph."""
-        ctx.save_for_backward(*states)
-        return dummy
-
-    @staticmethod
-    def backward(ctx, *grad_outputs):
-        """Reject differentiation through the storage-only output."""
-        raise RuntimeError("RNG snapshot storage is not part of the backward graph")
-
-
 class _CheckpointFrame:  # pylint: disable=too-few-public-methods
     """
     Storage frame for forward RNG states and detached activations from the forward recompute.
@@ -589,9 +574,6 @@ def __init__(self, recompute_fn: Callable, get_rng_state_tracker: Callable):
         self.recomputed = []
         self.count = 0
         self.rng_state = _CheckpointRNGState(get_rng_state_tracker)
-        self.rng_state_holder = _CheckpointRNGStateHolder.apply(
-            torch.empty(0, requires_grad=True), *self.rng_state.snapshot()
-        )
 
 
 class _recomputation_hook(
@@ -647,9 +629,7 @@ def unpack_hook(idx):
             """
 
             if not frame.recomputed:
-                with frame.rng_state.fork(
-                    frame.rng_state_holder.grad_fn.saved_tensors
-                ), _recomputation_hook(frame):
+                with frame.rng_state.fork(), _recomputation_hook(frame):
                     frame.recompute_fn(*args, **kwargs)
 
             # Return the already recomputed activation tensor at the given index

From 9a72905a8f7572abb8764ec19d7984ad446aeb87 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 10 Nov 2026 16:50:36 +0100
Subject: [PATCH 15/23] Remove the missing Philox API checkpoint tests

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py | 59 ----------------------------------
 1 file changed, 59 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 058f53d23e8..d28948ebda1 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -46,7 +46,6 @@
     is_nvfp4_available,
 )
 from transformer_engine.pytorch import checkpoint as te_checkpoint
-import transformer_engine.pytorch.distributed as te_distributed
 from transformer_engine.pytorch.distributed import (
     CudaRNGStatesTracker,
     is_fp8_activation_recompute_enabled,
@@ -1314,64 +1313,6 @@ def forward(self, value):
         previous = forward_output
 
 
-@pytest.mark.parametrize("use_reentrant", all_boolean)
-@pytest.mark.parametrize("dropout", (0.0, 0.5))
-@pytest.mark.parametrize("te_dropout", all_boolean)
-def test_checkpoint_graph_rng_without_philox_api(use_reentrant, dropout, te_dropout, monkeypatch):
-    """Older builds allow deterministic capture and reject stochastic recompute."""
-    monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
-    layer = te.ops.Dropout(dropout) if te_dropout else torch.nn.Dropout(dropout)
-
-    class Model(torch.nn.Module):
-        def forward(self, value):
-            return 2 * te_checkpoint(layer, value + 0, use_reentrant=use_reentrant)
-
-    sample = torch.ones(4096, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-    # The same dropout remains supported outside capture.
-    output = Model()(sample)
-    expected = output.detach().clone()
-    output.sum().backward()
-    assert torch.equal(sample.grad, expected)
-    del output, expected
-    sample.grad = None
-
-    original_call = Model.__call__
-    original_recompute_flags = tuple(
-        te_distributed.activation_recompute_forward._is_first_fp8_module
-    )
-    if dropout:
-        with pytest.raises(RuntimeError, match="philox_state"):
-            te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
-        assert Model.__call__ is original_call
-        assert (
-            tuple(te_distributed.activation_recompute_forward._is_first_fp8_module)
-            == original_recompute_flags
-        )
-    else:
-        model = te.make_graphed_callables(Model(), (sample,), num_warmup_iters=3)
-        value = torch.ones_like(sample, requires_grad=True)
-        model(value).sum().backward()
-        assert torch.equal(value.grad, torch.full_like(value, 2))
-
-
-@pytest.mark.parametrize("backend", ("FusedAttention", "FlashAttention"))
-def test_checkpoint_native_attention_requires_philox_api(backend, monkeypatch):
-    """Native attention dropout must reject capture before entering its CUDA backend."""
-    from transformer_engine.pytorch.attention.dot_product_attention import backends
-
-    monkeypatch.setattr(te_distributed, "_checkpoint_rng_available", lambda: False)
-    attention = getattr(backends, backend)(softmax_scale=0.125, attention_dropout=0.5)
-    value = torch.ones(16, 1, 1, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-
-    def body(query):
-        return attention(query, query, query)
-
-    graph = torch.cuda.CUDAGraph()
-    with pytest.raises(RuntimeError, match="philox_state"):
-        with torch.cuda.graph(graph):
-            te_checkpoint(body, value + 0, use_reentrant=True)
-
-
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""

From f4c9a2996e0f17cc6cb5742826e4b41d1e17b5fd Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 20:40:58 +0100
Subject: [PATCH 16/23] Remove unused custom op tags parameter

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/custom_op.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index bd49f4b84cf..1ed84017ecd 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1274,15 +1274,12 @@ def register_custom_op(
     arg_type: type,
     impl: Callable[[Any], Any],
     fake_impl: Callable[[Any], Any],
-    tags: Sequence[torch.Tag] = (),
 ) -> Optional[Callable[[Any], Any]]:
     """Register one custom op without autograd wiring.
 
     arg_type is a dataclass defining the input schema. Results may be nested
     tuples/lists of fresh tensors or None; fake_impl mirrors them with TensorSpec.
     The returned callable takes an args instance and preserves the result structure.
-    ``tags`` apply to both base and wrapper operators; mark RNG consumers with
-    ``torch.Tag.nondeterministic_seeded``.
     Returns None for unsupported registration APIs so callers can fall back to eager.
     Duplicate operator names raise ValueError.
     """
@@ -1298,7 +1295,6 @@ def pack_result(result):
             impl=impl,
             fake_impl=fake_impl,
             pack_result=pack_result,
-            tags=tags,
         )
     except (ImportError, AttributeError, RuntimeError, TypeError) as e:
         record_compile_disabled(

From 4ecc5330a274bd3e0b5b66e0f7900e578046fe8d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 21:04:07 +0100
Subject: [PATCH 17/23] Keep deprecated graph-safe RNG helpers

graph_safe_rng_available and is_graph_safe_rng_state are imported directly
by external projects and were deprecated only in #3788. Keep them until a
release cycle has passed.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/distributed.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 2b3945f1c95..1ca105cbf25 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -81,6 +81,24 @@ def set_all_rng_states(states: List) -> None:
     _ALL_ACTIVE_RNG_STATES = states
 
 
+def graph_safe_rng_available() -> bool:
+    """Returns whether cuda graph safe RNG state manipulation is supported.
+
+    Deprecated: always True, since all supported PyTorch versions provide the graph-safe RNG APIs.
+    """
+    warnings.warn(
+        "graph_safe_rng_available() is deprecated and always returns True.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return True
+
+
+def is_graph_safe_rng_state(state: Union[torch.Tensor, torch.Generator]) -> bool:
+    """Returns whether the rng state is a graph safe version."""
+    return isinstance(state, torch.Generator)
+
+
 def _get_cuda_rng_state(
     device: Union[int, str, torch.device] = te_device_type(),
     clone: bool = False,

From 369376ba75e771f0f3283cd39099b83b3f75a5ec Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 21:05:17 +0100
Subject: [PATCH 18/23] Keep dropout implementation selection for p=0

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/ops/basic/dropout.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/ops/basic/dropout.py b/transformer_engine/pytorch/ops/basic/dropout.py
index df4947f2c2a..65e51444a32 100644
--- a/transformer_engine/pytorch/ops/basic/dropout.py
+++ b/transformer_engine/pytorch/ops/basic/dropout.py
@@ -46,7 +46,7 @@ def op_forward(
 
         # Choose implementation
         impl = None
-        if not self.training or self.dropout_probability == 0:
+        if not self.training:
             impl = "evaluation"
         elif input_.numel() % 16 == 0 and dtype in (torch.float16, torch.bfloat16):
             impl = "fused"

From 2a3d1d30171e283a42c4981f802f3c102017831e Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 21:19:03 +0100
Subject: [PATCH 19/23] Accept preserve_rng_state in checkpoint during capture

Modules without TE layers use the TE checkpoint while capturing, so the
torch.utils.checkpoint keyword must not reach the checkpointed function.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/distributed.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 1ca105cbf25..5c532f595d8 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -781,6 +781,7 @@ def checkpoint(
     context_fn = kwargs.pop("context_fn", noop_context_fn)
     determinism_check = kwargs.pop("determinism_check", "default")
     debug = kwargs.pop("debug", False)
+    preserve_rng_state = kwargs.pop("preserve_rng_state", True)
     if not has_te_modules(function) and (
         not torch.cuda.is_initialized() or not torch.cuda.is_current_stream_capturing()
     ):
@@ -791,6 +792,7 @@ def checkpoint(
             context_fn=context_fn,
             determinism_check=determinism_check,
             debug=debug,
+            preserve_rng_state=preserve_rng_state,
             **kwargs,
         )
 
@@ -815,7 +817,7 @@ def checkpoint(
     #       cannot be sure there are no TE modules inside the function. It also means we might run
     #       the TE checkpoint for non-TE modules, so the TE checkpoint has to support a potential
     #       user context function.
-    del determinism_check, debug
+    del determinism_check, debug, preserve_rng_state
     if _USE_REENTRANT_ACTIVATION_RECOMPUTE:
         # If saved activations need to be distributed but there is no process group,
         # default to the world group.

From b34da49df3e56ad8244a574646d40a2442f815e6 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 21:25:44 +0100
Subject: [PATCH 20/23] Drop cosmetic edits unrelated to checkpoint RNG

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/distributed.py | 6 ++++--
 transformer_engine/pytorch/graph.py       | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 5c532f595d8..a67651c86de 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -909,7 +909,6 @@ def get_states(self) -> Dict[str, torch.Generator]:
         """
         Get rng states. Copy the dictionary so we have direct pointers
         to the states, not just a pointer to the dictionary.
-
         """
         states = {}
         for name in self.states_:
@@ -952,6 +951,7 @@ def add(self, name: str, seed: int) -> None:
         new_state = _get_cuda_rng_state(clone=True)
         new_state.manual_seed(seed)
         self.states_[name] = new_state
+        # Update global states.
         set_all_rng_states(self.states_)
 
     @contextmanager
@@ -976,7 +976,9 @@ def fork(self, name: str = "model-parallel-rng"):
         try:
             yield
         finally:
-            # And set the state to the original state we started with.
+            # The default generator shared its state with `self.states_[name]`,
+            # so that state already reflects any RNG consumed in this region.
+            # Set the state to the original state we started with.
             _set_cuda_rng_state(orig_cuda_rng_state)
 
 
diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py
index 24c0248467a..ca55b990eb5 100644
--- a/transformer_engine/pytorch/graph.py
+++ b/transformer_engine/pytorch/graph.py
@@ -486,7 +486,7 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params):
         )
 
     # For cases with multiple active RNG states, e.g. TP.
-    for state in get_all_rng_states().values():
+    for _, state in get_all_rng_states().items():
         for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs):
             fwd_graph.register_generator_state(state)
             bwd_graph.register_generator_state(state)

From cd2fa8b3d89291edf23732cd0538156a78348003 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 21:32:58 +0100
Subject: [PATCH 21/23] Restore checkpoint RNG handles after a failed value
 restore

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/distributed.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index a67651c86de..6ac90c95e8b 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -437,10 +437,12 @@ def fork(self):
             with self.guard():
                 yield
         finally:
-            self._restore(live_states)
-            _set_cuda_rng_state(cuda_state)
-            if self.tracker is not None:
-                self.tracker.set_states(tracker_states)
+            try:
+                self._restore(live_states)
+            finally:
+                _set_cuda_rng_state(cuda_state)
+                if self.tracker is not None:
+                    self.tracker.set_states(tracker_states)
 
 
 class _CheckpointFunction(torch.autograd.Function):

From 26645f82cd7e6997eb3a915ed634cafc5a0171aa Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 21:32:58 +0100
Subject: [PATCH 22/23] Test the checkpoint RNG guard without Philox state
 access

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py | 54 ++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 6a14f4e9485..8cd873b8cb4 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1313,6 +1313,60 @@ def forward(self, value):
         previous = forward_output
 
 
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+@pytest.mark.parametrize("dropout_impl", ("torch", "te_op"))
+@pytest.mark.parametrize("dropout_p", (0.0, 0.5))
+def test_checkpoint_rejects_captured_rng_without_philox_state(
+    use_reentrant, dropout_impl, dropout_p, monkeypatch
+):
+    """Without Philox state access, only deterministic captured checkpoints may proceed."""
+    monkeypatch.setattr(
+        "transformer_engine.pytorch.distributed._checkpoint_rng_available", lambda: False
+    )
+    tracker = CudaRNGStatesTracker()
+    tracker.add("model-parallel-rng", 4321)
+    dropout = te.ops.Dropout(dropout_p) if dropout_impl == "te_op" else None
+
+    def body(value):
+        with tracker.fork():
+            if dropout is not None:
+                return dropout(value)
+            return F.dropout(value, p=dropout_p, training=True)
+
+    def run(value):
+        return te_checkpoint(
+            body, value, use_reentrant=use_reentrant, get_rng_state_tracker=lambda: tracker
+        )
+
+    inp = torch.ones(4096, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            inp.grad = None
+            out = run(inp)
+            out.sum().backward()
+    torch.cuda.current_stream().wait_stream(stream)
+    assert torch.equal(inp.grad, out), "Eager checkpointing must stay supported"
+    del out
+
+    inp.grad = None
+    graph = torch.cuda.CUDAGraph()
+    graph.register_generator_state(tracker.get_states()["model-parallel-rng"])
+    with torch.cuda.graph(graph):
+        if dropout_p > 0:
+            inp.mul(1)
+            with pytest.raises(RuntimeError, match="philox_state"):
+                run(inp)
+        else:
+            out = run(inp)
+            out.sum().backward()
+    if dropout_p == 0:
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(inp.grad, out)
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""

From 5ea84ec6927eff420d8900aee9c803c03e2be17c Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 13 Nov 2026 22:09:09 +0100
Subject: [PATCH 23/23] Revert "Test the checkpoint RNG guard without Philox
 state access"

This reverts commit 26645f82cd7e6997eb3a915ed634cafc5a0171aa.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py | 54 ----------------------------------
 1 file changed, 54 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 8cd873b8cb4..6a14f4e9485 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1313,60 +1313,6 @@ def forward(self, value):
         previous = forward_output
 
 
-@pytest.mark.parametrize("use_reentrant", all_boolean)
-@pytest.mark.parametrize("dropout_impl", ("torch", "te_op"))
-@pytest.mark.parametrize("dropout_p", (0.0, 0.5))
-def test_checkpoint_rejects_captured_rng_without_philox_state(
-    use_reentrant, dropout_impl, dropout_p, monkeypatch
-):
-    """Without Philox state access, only deterministic captured checkpoints may proceed."""
-    monkeypatch.setattr(
-        "transformer_engine.pytorch.distributed._checkpoint_rng_available", lambda: False
-    )
-    tracker = CudaRNGStatesTracker()
-    tracker.add("model-parallel-rng", 4321)
-    dropout = te.ops.Dropout(dropout_p) if dropout_impl == "te_op" else None
-
-    def body(value):
-        with tracker.fork():
-            if dropout is not None:
-                return dropout(value)
-            return F.dropout(value, p=dropout_p, training=True)
-
-    def run(value):
-        return te_checkpoint(
-            body, value, use_reentrant=use_reentrant, get_rng_state_tracker=lambda: tracker
-        )
-
-    inp = torch.ones(4096, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            inp.grad = None
-            out = run(inp)
-            out.sum().backward()
-    torch.cuda.current_stream().wait_stream(stream)
-    assert torch.equal(inp.grad, out), "Eager checkpointing must stay supported"
-    del out
-
-    inp.grad = None
-    graph = torch.cuda.CUDAGraph()
-    graph.register_generator_state(tracker.get_states()["model-parallel-rng"])
-    with torch.cuda.graph(graph):
-        if dropout_p > 0:
-            inp.mul(1)
-            with pytest.raises(RuntimeError, match="philox_state"):
-                run(inp)
-        else:
-            out = run(inp)
-            out.sum().backward()
-    if dropout_p == 0:
-        graph.replay()
-        torch.cuda.synchronize()
-        assert torch.equal(inp.grad, out)
-
-
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_checkpoint_without_autograd_preserves_forward_context():
     """The direct path must preserve a context that explicitly enables gradients."""