Repository: pytorch
Issue: Can't use JIT modules traced with AMP autocast, with Triton Server (or any C++ environment) - freeze() issue ?
Labels: oncall: jit

### 🐛 Describe the bug

We have recently found out at NVIDIA that Triton's libtorch backend can't run TorchScript modules exported via torch.jit.trace() with amp.autocast. 
There is still outstanding issue about lack of true autocast support in C++ here:  https://github.com/pytorch/pytorch/issues/44710.
For our purposes(inference), we would actually prefer a different solution: 
I think after jit.freeze() and jit.optimize_for_inference, there should be no more _autocast_to_xxx nodes in the graph - they should be resolved according to current settings and frozen - like onnx.export does. Now we still see those after jit.freeze() and jit.optimize_for_inference: 

      context = torch.index(_4, _9)
      max_len = (torch.size(context))[-1]
      ids = torch.arange(0, max_len, dtype=None, layout=None, device=ops.prim.device(lens_sorted))
      mask = torch.lt(ids, torch.unsqueeze(lens_sorted, 1))
      _10 = torch.to(mask, ops.prim.dtype(context))
      mask0 = torch.unsqueeze(_10, 1)
      _11 = torch._autocast_to_full_precision(context, True, False)
      _12 = torch._autocast_to_full_precision(mask0, True, False)
      input = torch.mul(_11, _12)
      _13 = __torch__.torch.autograd.grad_mode.no_grad.__new__(__torch__.torch.autograd.grad_mode.no_grad)
      _13.prev = False
      with _13:
        _14 = torch._autocast_to_reduced_precision(mask0, True, False, 5, 15)
        update_mask = torch.conv1d(_14, CONSTANTS.c2, None, [1], [2])
        update_mask_filled = torch.masked_fill(update_mask, torch.eq(update_mask, 0), 5)
        _15 = torch._autocast_to_full_precision(update_mask_filled, True, False)
        _16 = torch._autocast_to_full_precision(torch.reciprocal(_15), True, False)
        mask_ratio = torch.mul(_16, 5)
        update_mask0 = torch.clamp(update_mask, 0, 1)
        _17 = torch._autocast_to_full_precision(mask_ratio, True, False)
        _18 = torch._autocast_to_full_precision(update_mask0, True, False)
        mask_ratio0 = torch.mul(_17, _18)
      _19 = torch._autocast_to_reduced_precision(input, True, False, 5, 15)
      _20 = torch.conv1d(_19, CONSTANTS.c3, CONSTANTS.c4, [1], [2])
      _21 = torch._autocast_to_full_precision(_20, True, False)


### Versions

PyTorch version: 1.13.0a0+d0d6b1f
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.2
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.8.89
GPU models and configuration: GPU 0: NVIDIA RTX A5000
Nvidia driver version: 515.65.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] functorch==0.3.0a0
[pip3] k2==1.21.dev20221109+cuda11.8.torch1.13.0a0
[pip3] numpy==1.22.2
[pip3] pytorch-lightning==1.7.7
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.13.0a0+d0d6b1f
[pip3] torch-tensorrt==1.3.0a0
[pip3] torchaudio==0.13.0
[pip3] torchmetrics==0.10.2
[pip3] torchvision==0.14.0a0
[conda] functorch                 0.3.0a0                  pypi_0    pypi
[conda] k2                        1.21.dev20221109+cuda11.8.torch1.13.0a0          pypi_0    pypi
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.22.2           py38h6ae9a64_0    conda-forge
[conda] pytorch-lightning         1.7.7                    pypi_0    pypi
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.13.0a0+d0d6b1f          pypi_0    pypi
[conda] torch-tensorrt            1.3.0a0                  pypi_0    pypi
[conda] torchaudio                0.13.0                   pypi_0    pypi
[conda] torchmetrics              0.10.2                   pypi_0    pypi
[conda] torchvision               0.14.0a0                 pypi_0    pypi


cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel