Using RNG Generator in torch.compile

torch.compile
1

Hi, I was trying to find some info on how random number generation works in torch compile. I have some code that generates quite often random numbers (an RL environment), and I am trying to speed everything up using compile, but I noticed calls to torch.rand and similar introduce graph breaks when using the generator argument. Is there a way to avoid this? What is actually happening under the hood here?

Here is an example:

import torch as th

th._logging.set_logs(recompiles=True, graph_breaks=True)

rng = th.Generator(device="cuda")

@th.compile(mode="max-autotune")
def sum_random( a, b):
    br = th.randn((1,), device= "cuda")*b
    return a+br

@th.compile(mode="max-autotune")
def sum_random_rng( a, b):
    br = th.randn((1,), device= "cuda", generator = rng)*b
    return a+br


a = th.as_tensor([10.,20,30], device = "cuda")
b = th.as_tensor([1.,2,3], device= "cuda")
print("Without generator:")
print(sum_random(a,b))
print("With generator:")
print(sum_random_rng(a,b))

And the output (on torch 2.8.0) is:

Without generator:
tensor([ 8.9647, 17.9294, 26.8942], device='cuda:0')
With generator:
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] Graph break in user code at /home/host/lr_ws/src/torch_rng_compile.py:19
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] Graph Break Reason: Failed to convert args/kwargs to proxy
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks]   Explanation: Missing `as_proxy()` implementation for some arg/kwarg.
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] 
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] 
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks]   Developer debug context: call_function args: TupleVariable(length=1) ConstantVariable(str: 'cuda') UserDefinedObjectVariable(Generator)
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] 
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] User code traceback:
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks]   File "/home/host/lr_ws/src/torch_rng_compile.py", line 28, in <module>
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks]     print(sum_random_rng(a,b))
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks]   File "/home/host/lr_ws/src/torch_rng_compile.py", line 19, in sum_random_rng
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks]     br = th.randn((1,), device= "cuda", generator = rng)*b
V0819 16:18:25.308000 694852 torch/_dynamo/symbolic_convert.py:560] [1/0] [__graph_breaks] 
tensor([ 8.8488, 17.6975, 26.5463], device='cuda:0')
2

I believe this issue is tracked here so feel free to share your use case in this GitHub issue.

3

Thanks, I dug a bit in the code starting from issue 88576 you linked, and it looks to me that things have been rearranged and support for Generators removed. This PR removed the lines mentioned in issue 88576. From what I understand the current handling of random functions in dynamo is at replace_random.py (see PR 100064). And from what I understand Generators are not handled and only the global seed can be used, I’m not sure if there’s some alternative, or if there are further discussions somewhere, it would be nice to have some mechanism to manage an RNG or have some equivalent method.