"CUDA error: invalid argument" on model test for vit_h_14

[YosuaMichael]

We got the following error on the CI:

___________________ test_classification_model[cuda-vit_h_14] ___________________
Traceback (most recent call last):
  File "/work/test/test_models.py", line 732, in test_classification_model
    _check_input_backprop(model, x)
  File "/work/test/test_models.py", line 226, in _check_input_backprop
    out[0].sum().backward()
  File "/work/ci_env/lib/python3.8/site-packages/torch/_tensor.py", line 488, in backward
    torch.autograd.backward(
  File "/work/ci_env/lib/python3.8/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: CUDA error: invalid argument
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

This error only happened on vit_h_14 model in cuda device (the cpu is fine). Also I cannot reproduce the error on AWS cluster machine. Seems like this error is either machine or environment dependent and likely to be pytorch-core issue.

Note: I have tried running the test with CUDA_LAUNCH_BLOCKING=1 but the error trace seems pretty similar (see here):

___________________ test_classification_model[cuda-vit_h_14] ___________________
Traceback (most recent call last):
  File "/work/test/test_models.py", line 725, in test_classification_model
    _check_input_backprop(model, x)
  File "/work/test/test_models.py", line 225, in _check_input_backprop
    out[0].sum().backward()
  File "/work/ci_env/lib/python3.8/site-packages/torch/_tensor.py", line 488, in backward
    torch.autograd.backward(
  File "/work/ci_env/lib/python3.8/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: CUDA error: invalid argument
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

cc @pmeier @seemethere @atalman @osalpekar

[NicolasHug]

Thank you Yosua for keeping track of this

I opened #7218 to skip this test.

This is quite bad, but I don't know how to solve it right now. I can't reproduce it on the AWS cluster either. I suspected that it could be a CUDA version issue: it passes on AWS cluster with CUDA 11.6, while our CI is 11.7 and #7214 shows that it also fails on CI x 11.8. But I can't tell whether CI x 11.6 also fails, and I can't tell whether AWS x 11.7 fails either. #7208 suggests that it's not a memory issue (although it could still be).

This could be a CUDA issue, or an upstream core issue. At this point and with all of what we have to do for the release (#7217), our last resort is to skip the test as done in #7218.

CC @atalman @malfet , just letting you know that this might be a core or CI issue, I'm just unable to tell at this point.

[weiwangmeta]

I just tried today's cu118 nightly torch and latest vision. The test seems to be able to pass

Command: pytest test/test_models.py::test_classification_model[cuda-vit_h_14] -v
test/test_models.py::test_classification_model[cuda-vit_h_14] PASSED

Also, looking into what is blocking #7221

@weiwangmeta update: this was not reproducible on A100