Deterministic integer indexing operation fails in indices size check / missing broadcast

[dbalchev]

🐛 Bug

When I run advanced integer indexing on cuda with deterministic algorithms I get

RuntimeError: linearIndex.numel()*sliceSize*nElemBefore == value.numel()INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/native/cuda/Indexing.cu":250, please report a bug to PyTorch. number of flattened indices did not match number of elements in the value tensor21

This bug is similar to #61032, but it affects integer indexing instead of boolean indexing.

To Reproduce

Steps to reproduce the behavior:

1. Run the following script

   import torch
   
   torch.use_deterministic_algorithms(True)
   
   x = torch.zeros(5).cuda()
   x[torch.tensor([1, 3]).cuda()] = 2
   print(x)

2. Get the error

   Traceback (most recent call last):
     File "Distilled Example.py", line 6, in <module>
       x[torch.tensor([1, 3]).cuda()] = 2
   RuntimeError: linearIndex.numel()*sliceSize*nElemBefore == value.numel()INTERNAL ASSERT FAILED at "/pytorch/aten/src/ATen/native/cuda/Indexing.cu":250, please report a bug to PyTorch. number of flattened indices did not match number of elements in the value tensor21
   

Expected behavior

Not getting a runtime error

Environment

Collecting environment information...
PyTorch version: 1.10.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.26

Python version: 3.7.9 (default, Nov 15 2021, 19:06:29)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.141-67.229.amzn2.x86_64-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: Tesla K80
Nvidia driver version: 460.73.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.18.0
[pip3] torch==1.10.0
[pip3] torch-scatter==2.0.9
[pip3] torchvision==0.11.0
[conda] Could not collect

Additional context

I originally get this error when I try to run torchvision FasterRCNN on cuda with deterministic algorithms:

import torch

torch.use_deterministic_algorithms(True)

from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn

model = fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False).cuda()

model(
    torch.zeros(1, 3, 800, 800).cuda(), [{
        'boxes': torch.tensor([[100, 200, 300, 400]]).cuda(), 
        'labels': torch.tensor([1]).cuda(),
    }])

cc @ngimel @mruberry @kurtamohler

[ngimel]

Should be fixed on master, cc @vfdev-5

[vfdev-5]

@dbalchev @ngimel yes, the first code snippet is working with pytorch nightly:

1.11.0.dev20211118+cu111
tensor([0., 2., 0., 2., 0.], device='cuda:0')

The second code snippet with fasterrcnn_resnet50_fpn should be run like :

CUBLAS_WORKSPACE_CONFIG=:4096:8 python script.py

to avoid errors like:

RuntimeError: Deterministic behavior was enabled with either `torch.use_deterministic_algorithms(True)` or `at::Context::setDeterministicAlgorithms(true)`, but this operation is not deterministic because it uses CuBLAS and you have CUDA >= 10.2. To enable deterministic behavior in this case, you must set an environment variable before running your PyTorch application: CUBLAS_WORKSPACE_CONFIG=:4096:8 or CUBLAS_WORKSPACE_CONFIG=:16:8. For more information, go to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility

I think we can close this issue as solved.

[dbalchev]

Thanks for the quick response (:

[andrei-rusu]

Is there any preferred workaround for this apart from being required to update the environment to PyTorch 1.11? I'm getting the same error on PyTorch 1.9, and I'd like to avoid updating the environment for this singular issue. Thanks!

EDIT: For others that are stumbling upon the same issue and don't want to update their PyTorch version, one could execute the indexing operation prior to sending the Tensor to CUDA or completely circumvent the broadcasting logic by assigning a CUDA tensor rather than a scalar:
e.g.

x[torch.tensor([1, 3]).cuda()] = torch.tensor([2,2]).cuda() 

[Starkguava]

为什么我在torch-2.6.0+cuda126上运行也会出现相同的问题？