Followup on Google Collab issue for release 2.3

[atalman]

Followup on Google Collab issue for release 2.3

[PaliC]

Verified the suggested fix works for torchonly smoketests on colab
https://colab.research.google.com/drive/1Wi_Xvv1AUDU2ogfoKCIIIOsGahhiFfKo?usp=sharing

[PaliC]

Verified this works on one of our devservers with an H100 with smoketests (all this time)
In a conda env run

To initially test we run

import torch
import jax
print(jax.default_backend())
print(torch.cuda.is_available())

python -m pip install torch==2.3.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu121
python -m pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

If the test script outputs

CUDA backend failed to initialize: Found CUDA version 12010, but JAX was built against version 12030, which is newer. The copy of CUDA that is installed must be at least as new as the version against which JAX was built. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Then Run
python -m pip uninstall -y
nvidia-cublas-cu12
nvidia-cuda-cupti-cu12
nvidia-cuda-nvrtc-cu12
nvidia-cuda-runtime-cu12
nvidia-cudnn-cu12
nvidia-cufft-cu12
nvidia-curand-cu12
nvidia-cusolver-cu12
nvidia-cusparse-cu12
nvidia-nccl-cu12
nvidia-nvjitlink-cu12
nvidia-nvtx-cu12 ; \

If running the test script you might get the following error

  File "/home/sahanp/.conda/envs/test/lib/python3.10/site-packages/torch/__init__.py", line 237, in <module>
   from torch._C import *  # noqa: F403
ImportError: libnccl.so.2: cannot open shared object file: No such file or directory

Then run

python -m pip install torch==2.3.0 --index-url https://download.pytorch.org/whl/test/cu121

Afterwards the test script returns

gpu
True

Running the smoke_tests we get the following output

sahanp@devgpu086 ~/b/t/smoke_test (main)> python smoke_test.py                                                        (test) 
torch: 2.3.0+cu121
Skip version check for channel None as stable version is None
Testing smoke_test_conv2d
Testing smoke_test_conv2d with cuda
/home/sahanp/.conda/envs/test/lib/python3.10/site-packages/torch/nn/modules/conv.py:456: UserWarning: Plan failed with a cudnnException: CUDNN_BACKEND_EXECUTION_PLAN_DESCRIPTOR: cudnnFinalize Descriptor Failed cudnn_status: CUDNN_STATUS_NOT_SUPPORTED (Triggered internally at ../aten/src/ATen/native/cudnn/Conv_v8.cpp:919.)
  return F.conv2d(input, weight, bias, self.stride,
Testing smoke_test_conv2d with cuda for torch.float16
Testing smoke_test_conv2d with cuda for torch.float32
Testing smoke_test_conv2d with cuda for torch.float64
Testing smoke_test_linalg on cpu
Testing smoke_test_linalg on cuda
Testing smoke_test_linalg with cuda for torch.float32
Testing smoke_test_linalg with cuda for torch.float64
Output: 
torchvision: 0.18.0+cu121
torch.cuda.is_available: True
torch.ops.image._jpeg_version() = 62
Is torchvision usable? True
German shepherd (cpu): 37.6%
German shepherd (cuda): 37.6%
/home/sahanp/.conda/envs/test/lib/python3.10/site-packages/torch/_inductor/compile_fx.py:124: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
torch.compile model output: torch.Size([1, 1000])


Output: 
Skipping ffmpeg test.
Smoke test passed.


torchvision CUDA: 12010
torchaudio CUDA: 12010
torch cuda: 12.1
torch cudnn: 8902
cuDNN enabled? True
torch nccl version: (2, 20, 5)
Testing smoke_test_compile for torch.float16
True
Testing smoke_test_compile for torch.float32
True
Testing smoke_test_compile for torch.float64
True
Testing smoke_test_compile with mode 'max-autotune' for torch.float32
Testing test_cuda_runtime_errors_captured
../aten/src/ATen/native/cuda/TensorCompare.cu:106: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0] Assertion `input[0] != 0` failed.
Caught CUDA exception with success: CUDA error: device-side assert triggered
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.