Sequential CTMRG is slow compared to Python (with PyTorch)

[Yue-Zhengyuan]

I use the CTMRG algorithm to measure the Heisenberg model ground state obtained from simple update. The algorithm settings are

trscheme = truncerr(1e-8) & truncdim(12)
ctm_alg = CTMRG(; tol=1e-12, miniter=4, maxiter=100, verbosity=3, trscheme=trscheme)

The PEPS bond dimension is D = 6, and the environment bond dimension is χ = 12. Starting from random CTMRGEnv, it takes about 1.3s to perform one CTMRG step:

[ Info: CTMRG init:     obj = +2.310478936455e-09       err = 1.0000e+00
[ Info: CTMRG   1:      obj = +2.532703514106e-01       err = 3.4933131838e-01  time = 4.43 sec
[ Info: CTMRG   1:      obj = +2.532703514106e-01       err = 3.4933131838e-01  time = 0.18 sec
[ Info: CTMRG   2:      obj = +7.808860223845e-01       err = 1.5017305263e-01  time = 1.46 sec
[ Info: CTMRG   2:      obj = +7.808860223845e-01       err = 1.5017305263e-01  time = 0.00 sec
[ Info: CTMRG   3:      obj = +9.380692547858e-01       err = 4.9406274229e-02  time = 1.21 sec
[ Info: CTMRG   3:      obj = +9.380692547858e-01       err = 4.9406274229e-02  time = 0.00 sec
[ Info: CTMRG   4:      obj = +9.836562322063e-01       err = 2.2671232670e-02  time = 1.19 sec
[ Info: CTMRG   4:      obj = +9.836562322063e-01       err = 2.2671232670e-02  time = 0.00 sec
[ Info: CTMRG   5:      obj = +9.991752105702e-01       err = 9.6623034220e-03  time = 1.35 sec
[ Info: CTMRG   5:      obj = +9.991752105702e-01       err = 9.6623034220e-03  time = 0.00 sec
[ Info: CTMRG   6:      obj = +1.004775679484e+00       err = 3.8899763085e-03  time = 1.34 sec
[ Info: CTMRG   6:      obj = +1.004775679484e+00       err = 3.8899763085e-03  time = 0.00 sec
...

However, using my own Python implementation (using PyTorch; the projectors are also found from the half-infinite environment), it only takes about 0.7s per step, about twice the speed of PEPSKit:

iter      svd_diff    time/s
0       2.2900e+01      0.64
1       2.6927e-01      0.64
2       4.4606e-02      0.65
3       8.5458e-03      0.65
4       1.3304e-03      0.66
5       1.8990e-04      0.65
6       2.6252e-05      0.75
...

Here svd_diff is the convergence criterion calculated as follows (a little bit different from the err of PEPSKit):

• Calculate singular values spectrum for each CTM tensor before and after the RG step
• Calculate the 2-norm of the spectrum difference for each CTM tensor
• Sum them up and devide by 8 * N_row * N_col

I tried to use the functions in PEPSKit to write a simpler version without the fancy autodiff stuff, then the speed can be improved to about 0.9s per RG step, but is still slower than PyTorch:

1         4.5357e-01    11.248 s
2         2.0233e-01     0.984 s
3         3.5182e-02     0.769 s
4         7.0871e-03     0.952 s
5         1.2966e-03     0.757 s
6         2.1065e-04     0.953 s

So my concern is that the auto-diff stuff from Zygote, etc may cause too much performance overhead for applications not using auto-diff of CTMRG.

[Yue-Zhengyuan]

For χ = 24 (with χ = 12 as initialization) the difference is more significant:

• PEPSKit: about 9.7s per step
• PEPSKit without AD: about 4.5s per step
• PyTorch: about 3s per step

I wonder how can I look for the bottleneck...

[lkdvos]

This is not entirely unsurprising. We developed this package with a mindset of "let's make it work first", mostly because Zygote poses a (rather large) number of restrictions on the optimizations that you would typically do to make a Julia algorithm faster. Now, we are indeed at the stage of thinking about what to optimise, but it is not so straightforward given these restrictions.

I'm more than happy to have a look at your implementation to think about more ways to speed up, but it's hard to give any kind of answers without further information: I don't know what the other implementations are, what the setup is, did you run this in a multithreaded environment or not, ... If you don't feel comfortable publicly sharing, we can also continue and discuss via email

[Yue-Zhengyuan]

For better control of setups, here I only attach the PEPSKit CTMRG with all AD stuff removed:

ctmrg_column.txt

The basic idea is still only to implement the left-move. The other moves are done by rotating the network by 90 degrees. One more change is that I write the functions to update only one column each time (which will be used in the full update algorithm), instead of handling all columns at once. I haven't figure out if there are some Python settings that may affect performance.

[lkdvos]

I dont really have enough information to tell you why the python implementation should have different performance. (since I dont really know what you are using in Python). Did you check if both implementations use 1 iteration to denote a move in every direction? Otherwise, I would advise to run a profiler and see if anything stands out

[Yue-Zhengyuan]

Yes, a move includes all four directions for all rows and columns. Actually I haven't fully removed Zygote overhead in the Julia version (in the function ctmrg_renormalize_col! I'm still using PEPSKit.renormalize_bottom_corner etc, which first creates a Zygote Buffer); I'll first see if it can still be further improved.

[lkdvos]

It might be a large variety of things, I would refrain from changing anything before you have a profiler view. The bottlenecks are very often not where you expect them to be. In this case, I don't think the Zygote buffers really do a lot, and I would expect the impact of not using in-place operations, because of the need to be AD compatible, be a much larger factor.

[pbrehmer]

Maybe to add a comment to this: When I was profiling leading_boundary, i.e. a complete CTMRG run, the runtime was dominated by the SVD calls. So the majority of the time was spent in LAPACKs gesvd. For example, for the Heisenberg model at $D=3$ and $\chi=48$, I obtained this profile view:

In principle the Python implementation should also just call some LAPACK routine for the SVD, right? I don't see much Zygote overhead in the forward pass here, so I am somewhat confused about these differences. Could it be that the difference in runtime has somehow to do with the underlying SVD?

[lkdvos]

One thing to try here is to use SDD instead. In fact, I am a bit confused that this is not our default, as this should be the faster option.

[pbrehmer]

I was doing a quick test with SDD and it really seems much faster. Here are some benchmarks for one CTMRG iteration on a random PEPS and environment at $D=3$ and $\chi=64$ using 4 Julia threads and 4 OpenBLAS threads:

The same thing at $D=5$ and $\chi=64$ already produces drastic differences between SVD and the rest:

It seems that doing a sparse SVD (although on a densely constructed tensor) still seems faster with the right parameters. Anyways, I will open up a PR which changes the default algorithm to TensorKit.SDD as it seems much more performant in general. This probably explains the runtime differences between the Python and Julia implementations!

[Yue-Zhengyuan]

Thanks for the tests! Indeed, the default SVD algorithm in PyTorch is gesdd. But I remember reading somewhere that although sdd is faster, it can be less stable in certain edge cases. I haven't got time to test it myself yet, and I'm not quite familiar with Julia's profiling tools. I will learn and do the testing in the following days.

[lkdvos]

In Julia, the default is also SDD, so I am okay with having that as the default. We have had issues with this in the past, for example in MPSKit, but this was related to taking the SVD of matrices which were close to unitary, in which case it seems like that algorithm indeed has convergence issues. In other cases I have rarely seen these issues show up.

[Yue-Zhengyuan]

I did some naive timing; with SDD the PEPSKit CTMRG speed is almost the same as my custom version without Zygote. The slight advantage of my Python version may be due to in-place rotation of the PEPS/CTMRGEnv. I now close this issue.