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 1/2] [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 2/2] [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)