WuYunfan:

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