I want to multiply two dense matrices A(N, d) and B(d, N). The resulting matrix will be of size N x N, where N is very large and I can’t use normal matrix multiplication function to do this.
However, I only want values at a few positions which are specified in a sparse matrix C(N, N). For example, I only want E non-zero entries in C.
My problem is the same as discussed in Matrix multiplication with values only at masked entries.
However, I don’t find a memory efficient way to solve this problem.
The solution provided by tom take a memory complexity of O(Ed).
Spliting it into chunks does not help as it still occupies memory for gradient ackpropagation.
matmul doesn’t store output AB, but your further indexing does. so, you can try writing custom autograd.Function (A,B) (returning (A@B)[indexes]) that would use sparse gradient of AB in backward (then do sparse-dense matmuls to compute gradients of A and B).
PS: I mean using chunked matmuls in forward, can’t think of a better way atm
I can’t do matmuls A@B there as it takes a time complexity of O(N^2 d).
Specifically, I have a gpraph with N nodes and E edges.
N is several millions and E is dozens of millions.
if you have memory problem with this:
vals = (A[rows]*B[cols]).sum(1)
maybe make it a function and wrap it with utils.checkpoint
It works well with my code, thanks a lot.