Repository: pytorch
Issue: JIT function crashing or failing (depending on profiling)
Labels: oncall: jit,NNC

### 🐛 Describe the bug

We've run in to issues using this simple function:
```
@torch.jit.script
    def jit_function(scores: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        return scores - (mask - 1) * 1e6
```

The following test leads to an error.
The second time `jit_function` is evaluated (with the same inputs), the result is `NaN` (rather than 1, which it should be).
```
def test1():
    @torch.jit.script
    def jit_function(scores: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        return scores - (mask - 1) * 1e6

    scores = torch.as_tensor([1], dtype=torch.float16, device="cuda")
    mask = scores.clone()

    result1 = jit_function(scores, mask)  # always 1?
    result2 = jit_function(scores, mask)  # nan sometimes

    assert result1.item() == result2.item()  # fails; 1 != NaN

test1()
```

The following configuration changes fix the above issue:
```
torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)
```

But with those changes, `jit_function` will fail when provided with large inputs, per this second test:
```
torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)


def test2():
    @torch.jit.script
    def jit_function(scores: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        return scores - (mask - 1) * 1e6

    # 2147483520 (= 128 * 16777215??)  is ok but 2147483521 isn't
    # Also fails with torch.randn not but torch.ones
    scores = torch.randn((2147483521,), dtype=torch.float16, device="cuda")
    mask = torch.ones_like(scores)
    jit_function(scores, mask)  # crashes sometimes


test2()
# RuntimeError: The following operation failed in the TorchScript interpreter.
# Traceback of TorchScript (most recent call last):
# RuntimeError: CUDA driver error: invalid argument
```

Both tests pass if only `torch._C._jit_set_profiling_mode(False)` is called (not `torch._C._jit_set_profiling_executor(False)` ).
However, this causes a bunch of minor numerics changes, which I find surprising.

I’m not really sure what the significance of these flags is, and I’m surprised that changing them has an impact on numerics. It seems like there might be multiple issues and I'm not sure how to disentangle them.

### Versions

```
Collecting environment information...
PyTorch version: 1.10.0a0+gitee11606
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.27

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.156-83.273.amzn2.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.4.152
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 470.57.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.921
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] pytorch-memlab==0.2.4
[pip3] torch==1.10.0a0+gitee11606
[pip3] torch-ops==0.0.0
[pip3] torchtyping==0.1.4
[conda] blas                      1.0                         mkl
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-include               2022.0.1           h06a4308_117
[conda] mkl-service               2.4.0            py38h7f8727e_0
[conda] mkl_fft                   1.3.1            py38hd3c417c_0
[conda] mkl_random                1.2.2            py38h51133e4_0
[conda] mypy                      0.921                    pypi_0    pypi
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.21.2           py38h20f2e39_0
[conda] numpy-base                1.21.2           py38h79a1101_0
[conda] pytorch-memlab            0.2.4                    pypi_0    pypi
[conda] torch                     1.10.0a0+gitee11606          pypi_0    pypi
[conda] torch-ops                 0.0.0                     dev_0    <develop>
[conda] torchtyping               0.1.4                    pypi_0    pypi
```