torch.distributions.MultivariateNormal.log_prob throws RuntimeError: CUDA error: CUBLAS_STATUS_NOT_SUPPORTED for large batch sizes

charlio23: <p>I am using torch.distributions.MultivariateNormal with large batches since I am working with long sequential data. For large batches, a runtime error from CUDA is thrown when computing the log_prob. Here is a snippet to reproduce the issue.</p>
<pre><code class="lang-auto">import torch
cov = torch.eye(3).to('cuda:1')
means = torch.randn(540000,3).to('cuda:1')
distribs = torch.distributions.MultivariateNormal(means, cov)

obs = torch.randn(540000,3).to('cuda:1')

log_prob = distribs.log_prob(obs)
</code></pre>
<p>This is the error message:</p>
<p>Traceback (most recent call last):<br>
File “/data2/users/cb221/ar-hmm/ar-hmm/code_snippet.py”, line 11, in <br>
log_prob = distribs.log_prob(obs)<br>
File “/data2/packages/anaconda3/envs/py38_pytorch/lib/python3.8/site-packages/torch/distributions/multivariate_normal.py”, line 216, in log_prob<br>
M = _batch_mahalanobis(self._unbroadcasted_scale_tril, diff)<br>
File “/data2/packages/anaconda3/envs/py38_pytorch/lib/python3.8/site-packages/torch/distributions/multivariate_normal.py”, line 59, in _batch_mahalanobis<br>
M_swap = torch.linalg.solve_triangular(flat_L, flat_x_swap, upper=False).pow(2).sum(-2)  # shape = b x c<br>
RuntimeError: CUDA error: CUBLAS_STATUS_NOT_SUPPORTED when calling <code>cublasStrsmBatched( handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb, batchCount)</code></p>
<p>For context, this is the ouput from <code>python3 -m torch.utils.collect_env</code>.</p>
<p>Collecting environment information…<br>
PyTorch version: 1.13.1+cu117<br>
Is debug build: False<br>
CUDA used to build PyTorch: 11.7<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: Ubuntu 20.04.6 LTS (x86_64)<br>
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0<br>
Clang version: Could not collect<br>
CMake version: Could not collect<br>
Libc version: glibc-2.31</p>
<p>Python version: 3.8.16 | packaged by conda-forge | (default, Feb  1 2023, 16:01:55)  [GCC 11.3.0] (64-bit runtime)<br>
Python platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.10<br>
Is CUDA available: True<br>
CUDA runtime version: Could not collect<br>
CUDA_MODULE_LOADING set to: LAZY<br>
GPU models and configuration:<br>
GPU 0: NVIDIA GeForce RTX 2080 Ti<br>
GPU 1: NVIDIA GeForce RTX 2080 Ti<br>
GPU 2: NVIDIA GeForce RTX 2080 Ti<br>
GPU 3: NVIDIA GeForce RTX 2080 Ti</p>
<p>Nvidia driver version: 515.105.01<br>
cuDNN version: Could not collect<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A<br>
Is XNNPACK available: True</p>
<p>Versions of relevant libraries:<br>
[pip3] numpy==1.24.2<br>
[pip3] torch==1.13.1<br>
[conda] numpy                     1.24.2                   pypi_0    pypi<br>
[conda] torch                     1.13.1                   pypi_0    pypi</p>
<p><strong>Note:</strong> I am not experiencing similar issues when using smaller batch sizes. The current solution I am using is to segment the batch into smaller pieces, run it iteratively, and then concatenate (which is obviously slower).</p>

ptrblck: <p>Thanks for forwarding the issue! I can reproduce it in <code>2.0.0+cu118</code> and we’ll look into it.</p>

eqy: <p>A similar issue was reported upstream: <a href="https://github.com/pytorch/pytorch/issues/97211" rel="noopener nofollow ugc">Triangular solve fails on batches of matrices of size &gt; (*, 524281) · Issue #97211 · pytorch/pytorch (github.com)</a> and fixed in cuBLAS which should hopefully be landing in a newer version soon.</p>

timw: <p>I got the same error on the latest Pytorch+CUDA, but it ran fine for me on Pytorch 1.7.1+CUDA toolkit 11.0.</p>

timw: <p>A similar problem was reported <a href="https://discuss.pytorch.org/t/how-to-use-torch-distributions-multivariate-normal-multivariatenormal-in-multi-gpu-mode/135030">here</a>. The fix there was to repeat the covariance matrices.</p>