Using DMRG with GPU

[EliasCombarro]

Hi!

I'm trying to use the DMRG algorithm with jax as the backend in order to try to speedup the computation with a GPU.

My code is really, really simple, just the following (adapted from one of the examples in the documentation):

H = qtn.tensor_builder.MPO_ham_heis(L=100)
dmrg = qtn.tensor_dmrg.DMRG2(H, bond_dims=[10, 20, 100, 100, 200], cutoffs=1e-10)
dmrg.solve(tol=1e-6, verbosity=1)

When I print H, I can see that the backend being used is numpy, but I do not know how to change it to jax to use the GPU.

Thanks in advance!

[jcmgray]

Hi @elias. To answer your question, there is not a plug and play way to do DMRG on the GPU using quimb currently. While almost all of the constituent parts should work, I wrote it before autoray was integrated into quimb and particularly the lazy LinearOperator eigen-decomposition stuff would require some special treatment since the iterative solvers are currently in scipy/slepc, i.e. cpu.

I'd note:

1. jax might not be the best choice for jitting the whole computation, as espeically for DMRG2, dynamic shapes and cutoff will prevent it compiling - torch or cupy might be better things to investigate.
2. You could use jax just for the individual contractions, by setting a contraction backend. I'm not sure who performant this would be overall though, swapping to the gpu and back and so forth.
3. You could try using TNOptimizer directly instead of DMRG to do an optimization, using gpu.

[jcmgray]

Hi @elias. To answer your question, there is not a plug and play way to do DMRG on the GPU using quimb currently. While almost all of the constituent parts should work, I wrote it before autoray was integrated into quimb and particularly the lazy LinearOperator eigen-decomposition stuff would require some special treatment since the iterative solvers are currently in scipy/slepc, i.e. cpu.

I'd note:

1. jax might not be the best choice for jitting the whole computation, as espeically for DMRG2, dynamic shapes and cutoff will prevent it compiling - torch or cupy might be better things to investigate.
2. You could use jax just for the individual contractions, by setting a contraction backend. I'm not sure who performant this would be overall though, swapping to the gpu and back and so forth.
3. You could try using TNOptimizer directly instead of DMRG to do an optimization, using gpu.

[EliasCombarro]

Thanks so much for the detailed answer. I will try to follow your suggestions!