diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 395476dbd1..b2425155bd 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,116 @@ 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): + """Recompute must match forward and preserve the live RNG states.""" + + 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 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(): + 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) + + # 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) + 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 55b1b3ae0f..63e3339009 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_: