Memory Issue with torchsort.soft_rank on CUDA

[liuquant]

I encountered a CUDA memory error when using torchsort.soft_rank during parallel training on the GPU. The error message is as follows:
File "/home/xxx/anaconda3/envs/DL2/lib/python3.10/site-packages/torchsort/ops.py", line 121, in backward
).gather(1, inv_permutation)
RuntimeError: CUDA error: an illegal memory access was encountered

However, when I switched the code to run torchsort.soft_rank on the CPU, while keeping the other parts of the code on the GPU, the error disappeared.

For example, when I modify the code like this:
if pred_2d.device != 'cpu':
pred_2d = pred_2d.to('cpu')
rank_2d = torchsort.soft_rank(pred_2d, regularization_strength=0.1)
rank_2d = rank_2d.to(y_true.device)

The error is resolved. But if I directly run:
rank_2d = torchsort.soft_rank(pred_2d, regularization_strength=0.1)
The error occurs again.

Could you provide guidance on how to resolve this issue when using torchsort.soft_rank with CUDA?
Thank you so much!

[liuquant]

By the way, my versions are:
cuda==12.6
torch==2.4.0
python==3.10.10

[teddykoker]

Hi @liuquant, would you be able to provide the input tensor necessary to reproduce this?

[liuquant]

Hi @liuquant, would you be able to provide the input tensor necessary to reproduce this?

Unfortunately, I am unable to provide the exact input tensor because the error occurred during the backward pass. The full error message is as follows: :
File "/home/xxx/DLpy_product_new/HyperPretrainForGlobal_Simple.py", line 287, in train
scaler.scale(l).backward()
File "/home/xxx/anaconda3/envs/DL2/lib/python3.10/site-packages/torch/_tensor.py", line 521, in backward
torch.autograd.backward(
File "/home/xxx/anaconda3/envs/DL2/lib/python3.10/site-packages/torch/autograd/init.py", line 289, in backward
_engine_run_backward(
File "/home/xxx/anaconda3/envs/DL2/lib/python3.10/site-packages/torch/autograd/graph.py", line 768, in _engine_run_backward
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
File "/home/xxx/anaconda3/envs/DL2/lib/python3.10/site-packages/torch/autograd/function.py", line 306, in apply
return user_fn(self, *args)
File "/home/xxx/anaconda3/envs/DL2/lib/python3.10/site-packages/torchsort/ops.py", line 121, in backward
).gather(1, inv_permutation)
RuntimeError: CUDA error: an illegal memory access was encountered
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

Unfortunately, I cannot pinpoint which specific tensor caused the issue. Apologies for the inconvenience.

[liuquant]

@teddykoker I suspect that this error might be related to gradient issues (nan or inf), as I noticed that it tends to occur whenever there is a problem with the gradients. This error usually doesn’t appear in the earlier stages of training.

[teddykoker]

I would recommend using torch.autograd.detect_anomaly to investigate the potential source of nan. Torchsort will have undefined behavior if I gets nan's for inputs on the forward or backward pass, which it does not check for. Without a reproducible example I'm not sure I'll be able to debug this very successfully.

[liuquant]

Thank you for the suggestion. I’ll use torch.autograd.detect_anomaly to investigate the potential source of NaNs. I’ll follow your advice and see if I can pinpoint the issue. I appreciate your help and patience.