memory usage in dmrg (julia)

Dear there,
I’ve a quick question. I wonder is there a way to know the memory usage for each sweep? (like some built-in function in itensor that can be used to extracted how much RAM is used in a sweep or sweep with different bonddim etc.) because I met some out-of-memory problem… I find a possible useful function called “write-to-disk” in itensor (didn’t try yet), but still want to know my actual RAM usage.
Thank you so much!

Unfortunately no, there isn’t such a straightforward way to check how much ram is used at a certain point of an ITensor calculation (or in a Julia code) that I know of. You could check the official Julia language Discourse also to see if there is some generic Julia language way to do this.

It’s a feature I’ve been wanting to add to ITensor to some time, but haven’t gotten around to it, namely to allow all of the ITensor types and related types, like MPS, to query how much memory they are using.

You may be able to find a “profiler” tool that reports memory usage at different parts of a program.

I see. Thanks!

Julia has a function Base.summarysize:

julia> Base.summarysize(randn(1000))
8040

julia> Base.summarysize(randomITensor(Index(1000)))
8232

You could probably pass that into DMRG through an Observer to get the memory footprint of the current MPS. However, the memory probably peaks inside the DMRG eigensolver, so we would need to allow users to query information about the eigensolver contraction step to get information like the memory usage out of that step.

Additionally, you can use Julia functions like Sys.free_memory and Sys.total_memory to get an idea of the total memory usage.

Finally, as Miles mentioned, the most general tool that wouldn’t require modifying or passing any code into ITensor functions would be memory profiling: Profiling · The Julia Language, which you can visualize with PProf.jl.

2 Likes

Building on Matt’s suggestion, I’ve now added a full code example showing how to use the observer system and Base.summarysize to monitor the memory usage of DMRG:

https://itensor.github.io/ITensors.jl/dev/examples/DMRG.html#Monitoring-the-Memory-Usage-of-DMRG

Full code below:

using ITensors

mutable struct SizeObserver <: AbstractObserver
end

function ITensors.measure!(o::SizeObserver; bond, half_sweep, psi, projected_operator, kwargs...)
  if bond==1 && half_sweep==2
    psi_size =  Base.format_bytes(Base.summarysize(psi))
    PH_size =  Base.format_bytes(Base.summarysize(projected_operator))
    println("|psi| = $psi_size, |PH| = $PH_size")
  end
end

let
  N = 100

  s = siteinds("S=1/2",N)

  a = OpSum()
  for n=1:N-1
    a += "Sz",n,"Sz",n+1
    a += 0.5,"S+",n,"S-",n+1
    a += 0.5,"S-",n,"S+",n+1
  end
  H = MPO(a,s)
  psi0 = randomMPS(s,linkdims=4)

  nsweeps = 5
  maxdim = [10,20,80,160]
  cutoff = 1E-8

  obs = SizeObserver()

  energy, psi = dmrg(H,psi0; nsweeps, maxdim, cutoff, observer=obs)

  return
end

Thank you for adding this!! It’s exactly what I need:) I just got time to try the sample code, but got a question. I just copy and paste it, didn’t change anyting, but I got an error message saing LoadError: UndefKeywordError: keyword argument projected_operator not assigned. I think projected_operator is the Hamiltonian H? Then does it mean I need somehow assign H to be projected_op?

Also, just a double check for syntex, can I use Sweeps function rather than setting nsweeps, maxdim, cutoff separately when using the observer? I mean doing this:

sweeps1 = Sweeps(5)
setmaxdim!(sweeps1,3000)
energy, psi = dmrg(H, psi0, sweeps1, observer=obs; write_when_maxdim_exceeds=1000)

I’ve tried this, but I got error message saying I miss keyword like nsweeps. Then I tried changing “,” to “;” after phi0, but still doens’t work, so just want to check.

Thank you!

Hm I’m not sure exactly why you are getting that error. It could be helpful if you posted a complete working code so I could reproduce your error.

My guess, though, is that your overload of the measure! function does not allow all of the keyword arguments that are being passed by dmrg (projected_operator is a new one that was added recently). A solution could be to put ,kwargs...) at the end of your measure! function definition so that it can accept any keyword args that were passed. Did you define it that way?

1 Like

I just find it’s because I didn’t update itensor package after you added this feature :upside_down_face: And it works perfectly now! Thanks for your reply!

And just to double check my understanding:

  1. For each sweep, |psi| shows the RAM used by the MPS and |PH| shows the RAM used by the Hamiltonian involved in this sweep, is this correct?

  2. And in my test results (attached), I find |psi| and |PH| are roughly same for each sweep if maxlinkdim is fixed, so does it mean RAM is not accumulated during sweeping? Then if say in one sweep total RAM~|psi|+|PH|=A+B, then I could set max RAM for the whole calculation to be also A+B, and it won’t be out of memory? (becasue I calculate in a cluster and need specify how much RAM I need for each job)

Thank you so much!

  1. yes |psi| and |PH| show the RAM used by the MPS and the Hamiltonian (projected into the current MPS basis)

  2. you’re correct that the amount of RAM used by these object is tied the maxdim. So if the maxdim stays fixed the RAM will stay fixed. Also yes the RAM is not accumulated during sweeping (unless the Julia garbage collector is slow to collect unused memory from previous sweeps). But there isn’t a way to set max RAM directly at the moment (i.e. no way to precisely specify how much RAM DMRG will use): the only thing you can do is control the maximum bond dimension either directly through the maxdim parameter or indirectly through the cutoff parameter.

I see! Thanks for your explination:)

So if I want to get an idea of what’s the largest maxdim I can set in a calculation (since my cluster has RAM constraint), the best way would be to try different maxdim and observe the value of |psi|+|PM| in a sweep? For example, if the largest RAM I can have is 128G, then I should figure out at what maxdim, |psi|+|PM| at each sweep starts approaching 128G and that maybe the largest maxdim I can have (up to some fine tuning considering apart from psi and PM there’s other usage of RAM in DMRG) .

Yes this is a good idea. The ram usage of the Hamiltonian is the larger of the two, and it scales as the bond dimension squared. So you could make a plot of the ram usage for smaller maxdim and fit it to a quadratic to estimate the ram usage at larger maxdims.

I see! And I just notice this, as you mentioned, in your example output ram of H is indeed much larger than that of psi, but in my output (attached above), I just notice my psi seems to occupy much larger ram than H, does this mean I may have something error in my code or this situation is also possible?

I think your code is ok. My statement is just a “rule of thumb” but there can be exceptions to it.

1 Like

I see! Thank you for your help:)!