Activation Checkpointing PT2 - AOTAutograd cannot handle set_rng_state 
Labels: module: checkpoint, triaged, oncall: pt2, module: dynamo

### 🐛 Describe the bug

### What am I trying to do?

My end goal is to support activation checkpointing using TorchDynamo escape hatch - `allow_in_graph`. As shown in the snippet below, `torch.utils.checkpoint.checkpoint` is inserted in the Dynamo generated FX graph as it is. Now, AOT Autograd traces it, and gets `set_rng_state` related error. 

I am first showing the full code that triggers the problem. And later, I will show a smaller repro

~~~
import logging
import torch
import torch._dynamo
import weakref
import traceback
import torch.utils.checkpoint

from functorch.compile import aot_function, aot_module, draw_graph, nop, print_compile  
# torch._dynamo.config.log_level = logging.DEBUG
class Mock(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.relu = torch.nn.ReLU()
    
    def forward(self, x, y):
        # return self.relu(x + y)
        a = torch.mm(x, y)
        # What happens if we do a graph break here :(
        # torch._dynamo.graph_break()
        a = torch.mm(a, y)
        b = self.relu(a)
        return b
mod = Mock()

@torch._dynamo.allow_in_graph
def allowed_checkpoint(x, y, preserve_rng_state):
    return torch.utils.checkpoint.checkpoint(mod, x, y, preserve_rng_state=preserve_rng_state, use_reentrant=False)

def fn(x, y, preserve_rng_state=True):
    a = allowed_checkpoint(x, y, preserve_rng_state=preserve_rng_state)
    # a = allowed_checkpoint(mod, a, y)
    return a

x = torch.randn(4, 4, device="cuda", requires_grad=True)
y = torch.randn(4, 4, device="cuda", requires_grad=True)

# Case 3 - Try aot_eager backend - Dont preserve rng state
opt_fn = torch.compile(fn, backend="aot_eager")
opt_z = opt_fn(x, y, True)
opt_z.sum().backward() 
print("Success")
~~~

This triggers following error. 

~~~
Traceback (most recent call last):
  File "/scratch/anijain/work/pytorch/torch/random.py", line 133, in fork_rng
    yield
  File "/scratch/anijain/work/pytorch/torch/utils/checkpoint.py", line 414, in unpack
    set_device_states(fwd_gpu_devices, fwd_gpu_states)
  File "/scratch/anijain/work/pytorch/torch/utils/checkpoint.py", line 58, in set_device_states
    torch.cuda.set_rng_state(state)
  File "/scratch/anijain/work/pytorch/torch/cuda/random.py", line 64, in set_rng_state
    _lazy_call(cb)
  File "/scratch/anijain/work/pytorch/torch/cuda/__init__.py", line 192, in _lazy_call
    callable()
  File "/scratch/anijain/work/pytorch/torch/cuda/random.py", line 62, in cb
    default_generator.set_state(new_state_copy)
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/scratch/anijain/work/pytorch/torch/_dynamo/output_graph.py", line 708, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.fake_example_inputs())
  File "/scratch/anijain/work/pytorch/torch/_dynamo/debug_utils.py", line 1055, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/scratch/anijain/work/pytorch/torch/_dynamo/backends/common.py", line 48, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 2810, in aot_module_simplified
    compiled_fn = create_aot_dispatcher_function(
  File "/scratch/anijain/work/pytorch/torch/_dynamo/utils.py", line 164, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 2503, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 1717, in aot_wrapper_dedupe
    return compiler_fn(flat_fn, leaf_flat_args, aot_config)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 2092, in aot_dispatch_autograd
    fx_g = make_fx(joint_forward_backward, aot_config.decompositions)(
  File "/scratch/anijain/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 714, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/scratch/anijain/work/pytorch/torch/_dynamo/eval_frame.py", line 215, in _fn
    return fn(*args, **kwargs)
  File "/scratch/anijain/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 443, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/scratch/anijain/work/pytorch/torch/_dynamo/eval_frame.py", line 215, in _fn
    return fn(*args, **kwargs)
  File "/scratch/anijain/work/pytorch/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "/scratch/anijain/work/pytorch/torch/fx/_symbolic_trace.py", line 652, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/scratch/anijain/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 459, in wrapped
    out = f(*tensors)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 1156, in traced_joint
    return functionalized_f_helper(primals, tangents)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 1108, in functionalized_f_helper
    f_outs = flat_fn_no_input_mutations(fn, f_primals, f_tangents, meta, keep_input_mutations)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 1076, in flat_fn_no_input_mutations
    outs = flat_fn_with_synthetic_bases_expanded(fn, primals, primals_after_cloning, maybe_tangents, meta, keep_input_mutations)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 1048, in flat_fn_with_synthetic_bases_expanded
    outs = forward_or_joint(fn, primals_before_cloning, primals, maybe_tangents, meta, keep_input_mutations)
  File "/scratch/anijain/work/pytorch/torch/_functorch/aot_autograd.py", line 1017, in forward_or_joint
    backward_out = torch.autograd.grad(
  File "/scratch/anijain/work/pytorch/torch/autograd/__init__.py", line 307, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/scratch/anijain/work/pytorch/torch/utils/checkpoint.py", line 420, in unpack
    _unused = function(*args, **kwargs)
  File "/scratch/anijain/work/env/lib/python3.9/contextlib.py", line 137, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/scratch/anijain/work/pytorch/torch/random.py", line 137, in fork_rng
    torch.cuda.set_rng_state(gpu_rng_state, device)
  File "/scratch/anijain/work/pytorch/torch/cuda/random.py", line 64, in set_rng_state
    _lazy_call(cb)
  File "/scratch/anijain/work/pytorch/torch/cuda/__init__.py", line 192, in _lazy_call
    callable()
  File "/scratch/anijain/work/pytorch/torch/cuda/random.py", line 62, in cb
    default_generator.set_state(new_state_copy)
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
~~~


### Why am I using allow_in_graph? Why not do a graph break?

* In the case of checkpointing, Dynamo tracing checkpointing code does not really add any benefit. If there is a graph break in the mod thats getting checkpointed, Dynamo generated subgraphs will not respect checkpointing. This is because the second subgraph will need to store the intermediate tensor for computing its backward graph. This storing of intermediate tensor breaks the eager checkpointing requirements. So, a module must not have graph breaks in order for PT2 stack to compile through the checkpoint-generated code. So, Dynamo does not add a benefit here.

* Checkpointing uses saved_tensors_hooks, which internally is a pair of hooks - pack (called during fwd) and unpack (called during bwd). These hooks are called from C++, and are not visible to Dynamo. We can ignore these hooks in Dynamo (assuming that hooks are traceable etc) and let AOTAutograd trace through them. However, in the case of checkpointing, the hooks have an unusual nonlocal mutation, which further upsets AOT Autograd.

* Graph break means less performance opportunity. However, this is less of an issue if the checkpointed code is decently large.

### Show me smaller code to repro the problem

Alright, here it is. I have removed all the checkpointing related code. This one directly uses `aot_function` and avoid Dynamo stack. 

~~~

import torch
from functorch.compile import aot_function, aot_module, draw_graph, nop, print_compile  

def fn(x):
    state = torch.cuda.get_rng_state()
    out = torch.randn(4) + x
    torch.cuda.set_rng_state(state)
    return out

x = torch.randn(4)
opt_fn = aot_function(fn, nop)
z = opt_fn(x)
print(z)
~~~

So, the question is - how can we support `set_rng_state`? IIUC, these need to be in the AOT generated graph, so that we do preserve the rng state (in the actual checkpointed code). But, these rng state ops are not aten ops. 

Looking for ideas.


### Error logs

_No response_

### Minified repro

_No response_

### Versions

NA

cc @ezyang @msaroufim @wconstab @ngimel @bdhirsh @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @soumith