Disable autocast

[anijain2305]

With AMP, AOT Autograd traced graph already reflects the AMP modifications.

However, Torchscript does not know that, and can try to AMP-ify already AMP-ified AOTAutograd traced graph, resulting in weird type promotion errors.

Concern (and that's why WIP) - Calling with for each forward and backward pass, might add overhead. Is there any better way?

[eellison]

We can probably provide a private api to run without amp

[anijain2305]

We can probably provide a private api to run without amp

Yes, this would be super helpful.

Just to confirm that we would like private API to be applicable for a torchscript graph. We already have global flags that we can disable/enable. But, that changes the behavior of the user code.

[eellison]

@anijain2305 i guess using jit compiler for forward/backward isnt in tree currently?

[anijain2305]

@anijain2305 i guess using jit compiler for forward/backward isnt in tree currently?

Not sure I fully understand.

We can definitely use Torchscript - https://github.com/pytorch/functorch/blob/main/functorch/_src/compilers.py#L23

This is the place where we call it - https://github.com/pytorch/functorch/blob/main/functorch/_src/aot_autograd.py#L169

(basically fw_compiler = ts_compile)

This is the AOTAutogrard return object - https://github.com/pytorch/functorch/blob/main/functorch/_src/aot_autograd.py#L143-L185. An autograd.Function with its forward and backward set to the compiled graphs. It is this place, where this PR wraps the forward and backward calls in with torch.cuda.amp.autocast(enabled=False)

[eellison]

Thanks! Yea was just looking for the ts_compile function to know where I need to add hooks to set this. I missed invocation here

[rwightman]

I was about to file a separate issue for an AMP problem, but might be covered here. Testing PT 1.12 w/ the 0.2 release branch built locally I can no longer use aot-autograd with AMP (and that's really the only combo I'm interested in). PT 1.11 w/ the current pypi release of functorch seemed fine. Now, w/ AMP enabled I get errors like

WARNING: "The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.9", line 54, in forward
    getitem_174 = native_batch_norm_backward_2[1]
    getitem_175 = native_batch_norm_backward_2[2];  native_batch_norm_backward_2 = None
    convolution_backward_2 = torch.ops.aten.convolution_backward(getitem_173, relu__45, _to_copy_51, [0], [1, 1], [0, 0], [1, 1], False, [0, 0], 1, [True, True, False]);  getitem_173 = _to_copy_51 = None
                             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    getitem_176 = convolution_backward_2[0]
    getitem_177 = convolution_backward_2[1];  convolution_backward_2 = None
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same

and

Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.3", line 39, in forward
    _unsafe_view_2 = torch.ops.aten._unsafe_view(bmm, [256, 12, 197, 197]);  bmm = None
    mul = torch.ops.aten.mul(_unsafe_view_2, 0.125);  _unsafe_view_2 = None
    _softmax = torch.ops.aten._softmax(mul, -1, True);  mul = None
               ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _to_copy_6 = torch.ops.aten._to_copy(_softmax, dtype = torch.float16);  _softmax = None
    expand_3 = torch.ops.aten.expand(_to_copy_6, [256, 12, 197, 197]);  _to_copy_6 = None
RuntimeError: conversion is supported for Half type only

[anijain2305]

I was about to file a separate issue for an AMP problem, but might be covered here. Testing PT 1.12 w/ the 0.2 release branch built locally I can no longer use aot-autograd with AMP (and that's really the only combo I'm interested in). PT 1.11 w/ the current pypi release of functorch seemed fine. Now, w/ AMP enabled I get errors like

WARNING: "The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.9", line 54, in forward
    getitem_174 = native_batch_norm_backward_2[1]
    getitem_175 = native_batch_norm_backward_2[2];  native_batch_norm_backward_2 = None
    convolution_backward_2 = torch.ops.aten.convolution_backward(getitem_173, relu__45, _to_copy_51, [0], [1, 1], [0, 0], [1, 1], False, [0, 0], 1, [True, True, False]);  getitem_173 = _to_copy_51 = None
                             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    getitem_176 = convolution_backward_2[0]
    getitem_177 = convolution_backward_2[1];  convolution_backward_2 = None
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same

and

Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.3", line 39, in forward
    _unsafe_view_2 = torch.ops.aten._unsafe_view(bmm, [256, 12, 197, 197]);  bmm = None
    mul = torch.ops.aten.mul(_unsafe_view_2, 0.125);  _unsafe_view_2 = None
    _softmax = torch.ops.aten._softmax(mul, -1, True);  mul = None
               ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _to_copy_6 = torch.ops.aten._to_copy(_softmax, dtype = torch.float16);  _softmax = None
    expand_3 = torch.ops.aten.expand(_to_copy_6, [256, 12, 197, 197]);  _to_copy_6 = None
RuntimeError: conversion is supported for Half type only

Hi @rwightman, does this PR work for you? If you have a script, I can try on my end as well.

We did not merge this one because there is a slightly better fix that @eellison is working on in PyTorch core. But, its still work in progress. So, if this PR works for you, I am inclined to merge this in and bring in functorch 0.2 release to unblock.

[Chillee]

LGTM, one minor nit.

[rwightman]

@anijain2305 I cherry-picked this onto the 0.2 release branch locally and appears to resolve the AMP issues for me.