# AOTAutograd incorrect lowering composite ops in inference\_mode

**URL:** <https://dev-discuss.pytorch.org/t/aotautograd-incorrect-lowering-composite-ops-in-inference-mode/2275>\
**Category:** compiler\
**Created:** [July 15, 2024, 8:29am UTC](https://dev-discuss.pytorch.org/t/aotautograd-incorrect-lowering-composite-ops-in-inference-mode/2275 "2024-07-15T08:29:45Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![wdziurdz](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/wdziurdz/32/2125_2.png) [@wdziurdz](https://dev-discuss.pytorch.org/u/wdziurdz)\
**Post date:** [July 15, 2024, 8:29am UTC](https://dev-discuss.pytorch.org/t/aotautograd-incorrect-lowering-composite-ops-in-inference-mode/2275/1 "2024-07-15T08:29:45Z")

</div>

I have a problem when using **aot\_autograd** with my own backend and **torch.compile** in in inference mode.

Example code:

```auto
def test1():
    with torch.inference_mode(True):
        @torch.compile(backend='my_backend')
        def fn(inp):
            instance_norm = torch.nn.InstanceNorm2d(3, device='myDevice')
            return instance_norm(inp)
        inp = torch.randn(1, 3, 4, 2).to(device='myDevice')
        print(fn(inp))

```

My backend:

```auto
@register_backend
def my_backend(graph_module: torch.fx.GraphModule, example_inputs: List[torch.Tensor], **kwargs):
	return aot_autograd(
		fw_compiler=training_compiler_fw,
		bw_compiler=training_compiler_bw,
		inference_compiler=inference_compiler,
	)(graph_module, example_inputs)

```

Without **aot\_autograd** , using only **inference\_compiler** , **instance\_norm** is normally lowered to my custom op under DispatchKey::MyPrivateKey.  
However, when using aot\_autograd, it is lowered using the native lowering and happens under CompositeImplicitAutograd.

Example of registration:

```auto

TORCH_LIBRARY_IMPL(aten, MyPrivateKey, m) {
    m.impl("instance_norm_wrap", instance_norm_wrap);
}

```

Could you help resolve this issue?

Or How should a backend override **CompositeImplicitAutograd** only for a specific backend? For other backends,  
we would like to dispatch via the default path, but it causes recursion. Is there guidance on how this should be done?

The following code works, but we still need the native path for non-MyDevice devices:

```auto
@torch.ops.aten.instance_norm.default.py_impl(torch.ops.aten.instance_norm.default, DispatchKey.CompositeImplicitAutograd)
def instance_norm(input, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled):
    return torch.ops.my_dev.instance_norm_wrap.default(input, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled)

@register_meta([torch.ops.my_dev.instance_norm_wrap.default])
def instance_norm(input, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled):
    out = torch.empty_like(input)
    return out

```

---

<div class="post-metadata">

**Author:** ![bdhirsh](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/bdhirsh/32/575_2.png) [@bdhirsh](https://dev-discuss.pytorch.org/u/bdhirsh)\
**Post date:** [July 16, 2024, 1:51am UTC](https://dev-discuss.pytorch.org/t/aotautograd-incorrect-lowering-composite-ops-in-inference-mode/2275/2 "2024-07-16T01:51:07Z")

</div>

You’re right that AOTDispatcher (previously AOTAutograd, since we think the new name is a bit more fitting) doesn’t have a great out-of-the-box way to prevent specific CompositeImplicitAutograd ops from getting decomposed during tracing under an inference setting - they will always get decomposed during tracing.

We could offer a better API for this - but for now, one thing you can do is modify the state of the python dispatcher (similar to what you did about with `py_impl`, but under a context manager so you will only temporarily modify it during your compilation):

```auto
@contextmanager
def override_instance_norm():
    op = torch.ops.aten.instance_norm.default
    old_table = op.py_kernels.copy()

    def new_impl(*args, **kwargs):
        return torch.ops.my_dev.instance_norm_wrap.default(*args, **kwargs)

    # temporarily put a decomp to your custom op in the dispatcher
    # careful - note that your new op will not work with training/autograd
    # unless you manually write and register your own derivative rule for it.
    op.py_impl(torch._C.DispatchKey.CompositeImplicitAutograd)(new_impl)

    try:
        yield
    finally:
        op.py_kernels.clear()
        op.py_kernels.update(old_table)
        # private APIs
        op._dispatch_cache.clear()

...

@register_backend
def my_backend(graph_module: torch.fx.GraphModule, example_inputs: List[torch.Tensor], **kwargs):
        # temporarily tweak how the dispatcher traces instance_norm when AOTDispatcher traces
        with override_instance_norm():
	        return aot_autograd(
	    	      fw_compiler=training_compiler_fw,
		      bw_compiler=training_compiler_bw,
		      inference_compiler=inference_compiler,
	        )(graph_module, example_inputs)

```

---

<div class="post-metadata">

**Author:** ![Sujoy\_Saraswati](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/sujoy_saraswati/32/783_2.png) [@Sujoy\_Saraswati](https://dev-discuss.pytorch.org/u/Sujoy_Saraswati)\
**Post date:** [July 17, 2024, 4:53am UTC](https://dev-discuss.pytorch.org/t/aotautograd-incorrect-lowering-composite-ops-in-inference-mode/2275/3 "2024-07-17T04:53:58Z")

</div>

@bdhirsh , thank you for your reply. The example you shared shows how to create a custom CompositImplicitAutograd op for instance\_norm. However, to enable it only for a specific backend, we would still require dispatches on other devices (example: CPU) to go through the default path even during graph compilation, in case there are instance\_norm on CPU still for some reason.

One possible way is to write the custom op “new\_impl” to handle this, by checking on the input tensor device and calling backend specific op when the inputs are on the device, and redirect op dispatch for other devices back to the default. Since we already create our own CompositImplicitAutograd dispatch with py\_impl, we need to create another context for other devices to pop our CompositImplicitAutograd entry, dispatch to the default instance\_norm and then push our CompositImplicitAutograd afterwards, to avoid the recursion @wdziurdz mentioned.

While this is doable, do you suggest any better way to define this? For example, is there any way to qualify custom op for dispatch keys like CompositImplicitAutograd for a specific backend only?

---

<div class="post-metadata">

**Author:** ![bdhirsh](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/bdhirsh/32/575_2.png) [@bdhirsh](https://dev-discuss.pytorch.org/u/bdhirsh)\
**Post date:** [July 17, 2024, 2:36pm UTC](https://dev-discuss.pytorch.org/t/aotautograd-incorrect-lowering-composite-ops-in-inference-mode/2275/4 "2024-07-17T14:36:51Z")

</div>

> While this is doable, do you suggest any better way to define this? For example, is there any way to qualify custom op for dispatch keys like CompositImplicitAutograd for a specific backend only?

that’s a good question - I don’t think there’s an easy way to do this today using the “usual” per-backend dispatch key registration. Our tracing logic today has some special logic to automatically hook into `CompositeImplicitAutograd` decompositions and run them, and `CompositeImplicitAutograd` decomps are usually expected to be backend agnostic. So I think having the “new decomp” that you register branch on device is probably… a reasonable workaround:

```auto
    def new_impl(*args, **kwargs):
        if get_device(args, kwargs) == 'my_backend_name':
            return torch.ops.my_dev.instance_norm_wrap.default(*args, **kwargs)
        else:
            # otherwise, dispatch to the original C++ decomp from core
            # I think this should avoid the recursion issues - op_dk() will call directly
            # into the CompositeImplicitAutograd registration from C++
            return op._op_dk(torch._C.DispatchKey.CompositeImplicitAutograd, *args, **kwargs)

```