[Compile] Error with FullGraph "eager" compile on TensorSubclass version

[drisspg]

Summary

Using nightly: torch==2.2.0.dev20230922+cu121

Repro

Running the following tests:
pytest tests/test_compile.py -k "eager"

As well I created a PyTorch Issue tracking what I think is a minimal repro:

Issue: pytorch/pytorch#109910

Output

E       torch._dynamo.exc.Unsupported: call_function UserDefinedObjectVariable(to_float8) [TensorVariable(), TensorVariable(), ConstantVariable(dtype), TensorVariable()] {}
E
E       from user code:
E          File "/home/drisspg/meta/float8_experimental/float8_experimental/float8_linear.py", line 252, in forward
E           x_fp8 = self.cast_x_to_float8(x, self.is_amax_initialized)
E         File "/home/drisspg/meta/float8_experimental/float8_experimental/float8_linear.py", line 204, in cast_x_to_float8
E           x_fp8 = Float8Tensor.to_float8(
E
E       Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
E
E
E       You can suppress this exception and fall back to eager by setting:
E           import torch._dynamo
E           torch._dynamo.config.suppress_errors = True

../../miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/exc.py:176: Unsupported

[drisspg]

@zou3519 Is this error because:

float8_experimental/float8_experimental/float8_tensor.py

Line 113 in 4ba9e06

def to_float8(cls, tensor, scale, float8_dtype, amax_buffer=None):

is passing in a torch.dtype to an autograd function?

[drisspg]

Note that this test will run succesful for fullgraph=False

[drisspg]

Update:
When running the tests with: torch==2.2.0.dev20231003+cu121 the error appears to not be reproducing anymore but instead throws the following warnings:

tests/test_compile.py::test_eager_only[dtype0-True-False-False] [2023-10-03 10:28:09,331] [2/1] torch._dynamo.variables.higher_order_ops: [WARNING] speculate_subgraph: while introspecting the user-defined autograd.Function, we were unable to trace function `trampoline_autograd_fwd` into a single graph. This means that Dynamo was unable to prove safety for this API and will fall back to eager-mode PyTorch, which could lead to a slowdown.
[2023-10-03 10:28:09,331] [2/1] torch._dynamo.variables.higher_order_ops: [ERROR] Unexpected type in sourceless builder <class 'torch.dtype'>
[2023-10-03 10:28:09,359] [8/0] torch._dynamo.variables.higher_order_ops: [WARNING] speculate_subgraph: while introspecting the user-defined autograd.Function, we were unable to trace function `trampoline_autograd_fwd` into a single graph. This means that Dynamo was unable to prove safety for this API and will fall back to eager-mode PyTorch, which could lead to a slowdown.
[2023-10-03 10:28:09,359] [8/0] torch._dynamo.variables.higher_order_ops: [ERROR] call_method GetAttrVariable(GetAttrVariable(TensorVariable(), _data), shape) __getitem__ (ConstantVariable(int),) {}
[2023-10-03 10:28:09,370] [9/0] torch._dynamo.variables.higher_order_ops: [WARNING] speculate_subgraph: while introspecting the user-defined autograd.Function, we were unable to trace function `trampoline_autograd_fwd` into a single graph. This means that Dynamo was unable to prove safety for this API and will fall back to eager-mode PyTorch, which could lead to a slowdown.
[2023-10-03 10:28:09,370] [9/0] torch._dynamo.variables.higher_order_ops: [ERROR] call_method GetAttrVariable(GetAttrVariable(TensorVariable(), _data), shape) __getitem__ (ConstantVariable(int),) {}

[drisspg]

hmm idk whats going on:
When I use backend = "eager" I get no errors but the warnings are thrown.
When I use backend = EagerAndRecordGraphs() I get no error on full_graph = False but get the original error for full_graph=True

[zou3519]

It sounds like we don't support torch.compile with torch.autograd.Function that accepts dtype inputs. torch.compile is falling back to eager-mode when fullgraph = False

[drisspg]

@zou3519 is there a way to enable support for this? or is there some invariant I am unaware of that prevents this?

[zou3519]

This is mostly just a self-imposed limitation. We can lift it