Unfortunately it is not computationally feasible to run this on CPU. If I run the model with parameters that are too small the error does not appear, and even in the most scaled down case that I have tried where I can reproduce the error, running it on CPU would be at best highly inconvenient.
I have tried CUDA_LAUNCH_BLOCKING=1, and when I did that I got an error message that was several hundred lines long and very difficult to parse. It looks something like this:
/pytorch/torch/lib/THC/THCTensorIndex.cu:325: void indexSelectLargeIndex(TensorInfo<T, IndexType>, TensorInfo<T, IndexType>, TensorInfo<long, IndexType>, int, int, IndexType, IndexType, long) [with T = float, IndexType = unsigned int, DstDim = 2, SrcDim = 2, IdxDim = -2]: block: [50,0,0], thread: [30,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/torch/lib/THC/THCTensorIndex.cu:325: void indexSelectLargeIndex(TensorInfo<T, IndexType>, TensorInfo<T, IndexType>, TensorInfo<long, IndexType>, int, int, IndexType, IndexType, long) [with T = float, IndexType = unsigned int, DstDim = 2, SrcDim = 2, IdxDim = -2]: block: [50,0,0], thread: [31,0,0] Assertion srcIndex < srcSelectDimSize failed.
Segmentation fault
except with that first part repeated several hundred times.
I have considered that the segmentation fault I saw there and the one I get when I try the hook are caused by the same thing but I don’t know enough about what could cause a segfault in PyTorch to know for sure. From what I have looked at it seems like the CUDA error I get happens somewhere during the forward pass when there are NaNs in my weights, and somewhere in there CUDA goes apeshit and segfaults. I am trying to track where the gradient becomes NaN in the previous backward pass before this happens, that is why I am attempting to use the backward hooks. I am not sure why/whether the hook might trigger the same issue with the segfault - though when I get it from the forward pass I get a massively long error message and when I get it from the hook it literally just says “Segmentation Fault” and nothing else.
Unfortunately giving a small code sample will also be difficult as the code base for this model is running in the thousands of lines now. I could try just giving all the code for the forward pass? That would still be very long though.