Dynamo error for autograd function

[drisspg]

Summary

Encountering Dynamo error when attempting to compile an autograd function that returns a dtype.

Repro

import torch

class dtype_test(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        tensor: torch.Tensor,
        dtype=torch.dtype,
    ):
        orig_precision = tensor.dtype
        ctx.orig_precision = orig_precision
        return tensor.to(dtype), dtype
    @staticmethod
    def backward(ctx, dOut):
        return dOut.to(ctx.orig_precision), None, None,

def main():
    x = torch.randn(16, 16, device="cpu", dtype=torch.float32)
    out = dtype_test.apply(x, torch.bfloat16)
    def test_func(x, dtype):
        return dtype_test.apply(x, dtype)
    compiled_func = torch.compile(test_func, fullgraph=True)
    y = compiled_func(x, torch.bfloat16)

if __name__ == "__main__":
    main()

Output

 File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 329, in call_method
   raise unimplemented(f"call_method {self} {name} {args} {kwargs}")
 File "/home/drisspg/miniconda3/envs/nightly/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 176, in unimplemented
   raise Unsupported(msg)
torch._dynamo.exc.Unsupported: call_method TupleVariable() to [ConstantVariable(dtype)] {}

from user code:
  File "/home/drisspg/meta/float8_experimental/../scripts/compile/autograd_dtype_float.py", line 21, in test_func
   return dtype_test.apply(x, dtype)
 File "/home/drisspg/meta/float8_experimental/../scripts/compile/autograd_dtype_float.py", line 12, in forward
   return tensor.to(dtype), dtype

Note

This is a proxy issue I am creating for meta-pytorch/float8_experimental#108

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7 @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

[ezyang]

cc @voznesenskym

[drisspg]

Any update here? Is there anything more pro-active I can be doing to push this along?

[ezyang]

I can't access the other issue, so none of us has any idea how high prio this is. Is it possible to work around?

[zou3519]

We shouldn't be erroring out on this in fullgraph=False. Also cc @yanboliang who's been poking around at the autograd.Function pieces.

Workaround should be to set torch._dynamo.config.capture_autograd_function = False, or allow_in_graph the call to the autograd.Function's apply

[drisspg]

Hmmm, so just updated to nightly the repro, is not reproing for me anymore. And now on the main issue the error still persists except throws the following warnings:

tests/test_compile.py::test_eager_only[dtype0-True-False-False] [2023-10-03 10:07:22,429] [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:07:22,429] [2/1] torch._dynamo.variables.higher_order_ops: [ERROR] Unexpected type in sourceless builder <class 'torch.dtype'>
[2023-10-03 10:07:22,457] [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:07:22,458] [8/0] torch._dynamo.variables.higher_order_ops: [ERROR] call_method GetAttrVariable(GetAttrVariable(TensorVariable(), _data), shape) __getitem__ (ConstantVariable(int),) {}
[2023-10-03 10:07:22,468] [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:07:22,468] [9/0] torch._dynamo.variables.higher_order_ops: [ERROR] call_method GetAttrVariable(GetAttrVariable(TensorVariable(), _data), shape) __getitem__ (ConstantVariable(int),) {}

so I can probably close this issue, not sure though if there has been changes that would have fixed this

[zou3519]

Kind of unclear. But we can still leave this issue open (for the fullgraph=True case) -- Dynamo is still not capturing the autograd.Function, it's just that on the nightly it has learned to gracefully fall back to eager-mode

[yanboliang]

The repro works well for me, no graph break.