Implement multi-threading using OhMyThreads and make it differentiable #70

pbrehmer · 2024-09-27T14:38:58Z

Here we'll replace the @fwdthreads macro with tmap and foreach calls. Additionally, we will code up reverse rules such that the backwards pass also runs in parallel.

pbrehmer · 2024-09-30T15:36:11Z

@lkdvos I could use some help on this one... I still find writing rrules very confusing sometimes, so could you maybe take a look at the rrule for dtforeach, once you have time? And while the dtmap rrule already seems to work, it might not be perfect yet :-)

Other than that we have to think about how to pass along the threading kwargs to the dtmap and dtforeach calls. One option is to store those kwargs inside the CTMRG struct but that feels a bit wrong. (Also there are some calls which wouldn't have access to a CTMRG instance.) Not sure what would be the best solution for that. Would global variables work where the user can set that somehow?

I'll also add some rrule tests in the end.

lkdvos · 2024-09-30T21:19:38Z

I think tforeach is going to be a bit hard, because it has no outputs, and is thus necessarily in-place... I have no clue how the zygote buffer magic works, so I can't really say I know how to deal with that either.

I should have some more time to think this through next week though!

For the global variables, I would maybe suggest ScopedVariables.jl instead, this is a little more flexible and shouldn't incur too much runtime costs

pbrehmer · 2024-10-01T07:53:47Z

I think tforeach is going to be a bit hard, because it has no outputs, and is thus necessarily in-place... I have no clue how the zygote buffer magic works, so I can't really say I know how to deal with that either.

How about we just stick to tmap then? The only slightly annoying thing is having to separate multiple return values at different indices but that should not incur too much overhead.

For the global variables, I would maybe suggest ScopedVariables.jl instead, this is a little more flexible and shouldn't incur too much runtime costs

Wasn't aware of ScopedValues.jl yet, that looks like a great solution. But I don't quite understand the necessity of a scoped value here since we never need to access the threading settings inside a multi-threaded map, right? In any case, we can probably just have a global scoped Dict with the threading settings that are passed on to the dtmap calls, and that can be mutated by some set_thread_settings function.

Anyways, I will give these things a go and then we can review next week, when you have time :)

…om :sequential mode)

…ction

lkdvos

I think I should have elaborated a bit more on what I had in mind with the scoped values: I would keep the threading strategies in scoped values, such that they can always be accessed as if they are global values. (this prevents them from bloating all of our algorithms)
The benefit of having them as a scoped value however means that users could still change them by calling the peps function from within a scope with a modified scoped value, thus changing the scheduler.

src/utility/diffable_threads.jl

lkdvos · 2024-10-05T10:20:12Z

src/utility/diffable_threads.jl

+        backevals = tmap(CartesianIndices(A); kwargs...) do idx
+            last(el_rrules[idx])(dy[idx])
+        end
+        df = ProjectTo(f)(sum(first, backevals))


Somehow it should be possible to do a tmapreduce and combine all of this, but I'll try and look into it a bit more

src/utility/diffable_threads.jl

…Defaults docstring

…iffable-threads

pbrehmer · 2024-10-29T16:07:45Z

I was trying to fix the Zygote error but with no luck, I really don't know how to handle the NoTangents (dA is an Array{NoTangent,3}) as they are being converted to Nothing. I'll copy the error here for future reference:

ERROR: MethodError: no method matching length(::Nothing)
The function `length` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  length(::Zygote.Grads, Any...; kwargs...)
   @ Zygote ~/.julia/packages/MacroTools/Cf2ok/src/examples/forward.jl:17
  length(::Combinatorics.Partition)
   @ Combinatorics ~/.julia/packages/Combinatorics/Udg6X/src/youngdiagrams.jl:8
  length(::Base.MethodSpecializations)
   @ Base reflection.jl:1317
  ...

Stacktrace:
  [1] productfunc(xs::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}, dy::Array{Nothing, 3})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/lib/array.jl:278
  [2] (::Zygote.var"#collect_product_pullback#744"{Base.Iterators.ProductIterator{Tuple{…}}})(dy::Array{Nothing, 3})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/lib/array.jl:300
  [3] (::Zygote.var"#2941#back#745"{Zygote.var"#collect_product_pullback#744"{…}})(Δ::Array{Nothing, 3})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
  [4] ctmrg_renormalize
    @ ~/repos/PEPSKit.jl/src/algorithms/ctmrg/ctmrg.jl:383 [inlined]
  [5] (::Zygote.Pullback{Tuple{…}, Any})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
  [6] ctmrg_iter
    @ ~/repos/PEPSKit.jl/src/algorithms/ctmrg/ctmrg.jl:163 [inlined]
  [7] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{CTMRGEnv{…}, Nothing})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
  [8] #221
    @ ~/repos/PEPSKit.jl/src/algorithms/ctmrg/ctmrg.jl:118 [inlined]
  [9] (::Zygote.Pullback{Tuple{…}, Any})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [10] #35
    @ ~/.julia/packages/LoggingExtras/cFgEq/src/verbosity.jl:117 [inlined]
 [11] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [12] (::Zygote.var"#ad_pullback#61"{Tuple{…}, Zygote.Pullback{…}})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/chainrules.jl:264
 [13] with_logger_pullback
    @ ~/.julia/packages/ChainRules/vdf7M/src/rulesets/Base/CoreLogging.jl:12 [inlined]
 [14] (::Zygote.ZBack{ChainRules.var"#with_logger_pullback#862"{…}})(dy::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/chainrules.jl:212
 [15] #withlevel#34
    @ ~/.julia/packages/LoggingExtras/cFgEq/src/verbosity.jl:113 [inlined]
 [16] (::Zygote.Pullback{Tuple{…}, Any})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [17] withlevel
    @ ~/.julia/packages/LoggingExtras/cFgEq/src/verbosity.jl:107 [inlined]
 [18] (::Zygote.Pullback{Tuple{…}, Any})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [19] withlevel
    @ ~/.julia/packages/LoggingExtras/cFgEq/src/verbosity.jl:107 [inlined]
 [20] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [21] leading_boundary
    @ ~/repos/PEPSKit.jl/src/algorithms/ctmrg/ctmrg.jl:115 [inlined]
 [22] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [23] #37
    @ ./REPL[24]:2 [inlined]
 [24] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface2.jl:0
 [25] (::Zygote.var"#78#79"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface.jl:91
 [26] withgradient(f::Function, args::InfinitePEPS{TrivialTensorMap{ComplexSpace, 1, 4, Matrix{ComplexF64}}})
    @ Zygote ~/.julia/packages/Zygote/NRp5C/src/compiler/interface.jl:213
 [27] top-level scope
    @ REPL[24]:1

lkdvos · 2024-10-29T21:08:45Z

I'm honestly not so sure how no one has ever run into this, but I seem to have been able to circumvent the issue by just not differentiating through the collect(Iterators.product())calls. It might have something to do with our tmapgradient returning an array of nothing instead of simply nothing, but I don't really have the time to investigate that and this seems to work.

codecov · 2024-10-30T00:04:23Z

Codecov Report

Attention: Patch coverage is 93.54839% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/utility/diffable_threads.jl	83.33%	3 Missing ⚠️
src/operators/infinitepepo.jl	0.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/PEPSKit.jl	`100.00% <ø> (ø)`
src/algorithms/ctmrg/ctmrg.jl	`92.39% <100.00%> (+0.99%)`	⬆️
src/algorithms/ctmrg/gaugefix.jl	`94.87% <100.00%> (+0.13%)`	⬆️
src/algorithms/toolbox.jl	`98.33% <100.00%> (+0.02%)`	⬆️
src/environments/ctmrg_environments.jl	`72.29% <100.00%> (+1.17%)`	⬆️
src/states/infinitepeps.jl	`67.77% <100.00%> (+1.49%)`	⬆️
src/utility/util.jl	`54.21% <100.00%> (-1.60%)`	⬇️
src/operators/infinitepepo.jl	`18.86% <0.00%> (ø)`
src/utility/diffable_threads.jl	`83.33% <83.33%> (ø)`

... and 1 file with indirect coverage changes

[skip ci]

pbrehmer added 5 commits September 27, 2024 13:59

Replace maps and loops with tmap, tforeach and tmapreduce

0fced36

Replace tmapreduce by tmap

bb1643b

Merge branch 'master' into pb-diffable-threads

a92c91f

Add rrule for tmap via dtmap wrapper

0d7b807

Set up tests for dtmap and dtforeach, reorganize test folder a bit

6e48268

pbrehmer added 2 commits October 1, 2024 10:39

Remove tforeach calls and replace by tmaps (remove parallelization fr…

a741762

…om :sequential mode)

Add global threading_kwargs scoped dict and set_threading_kwargs! fun…

7b858fb

…ction

lkdvos reviewed Oct 5, 2024

View reviewed changes

pbrehmer and others added 6 commits October 7, 2024 17:50

Remove ScopedValues and use threading kwargs global constant, update …

8b8bbf4

…Defaults docstring

Remove tmap from rrule

207ec1e

small improvements tmap rrule

5af77fc

Merge branch 'master' of github.com:quantumghent/PEPSKit.jl into pb-d…

c0bd923

…iffable-threads

Fix formatting

46f380b

Fix merge problems

491a464

Fix problem with AD through Iterators.product

a95570a

lkdvos force-pushed the pb-diffable-threads branch from cbc1fe7 to f3d7192 Compare October 29, 2024 21:57

Formatter

cab6c69

lkdvos force-pushed the pb-diffable-threads branch from f3d7192 to cab6c69 Compare October 29, 2024 23:12

formatter

ec24e4b

[skip ci]

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement multi-threading using OhMyThreads and make it differentiable #70

Implement multi-threading using OhMyThreads and make it differentiable #70

pbrehmer commented Sep 27, 2024 •

edited

Loading

pbrehmer commented Sep 30, 2024

lkdvos commented Sep 30, 2024

pbrehmer commented Oct 1, 2024

lkdvos left a comment

lkdvos Oct 5, 2024

pbrehmer commented Oct 29, 2024

lkdvos commented Oct 29, 2024

codecov bot commented Oct 30, 2024 •

edited

Loading

Implement multi-threading using OhMyThreads and make it differentiable #70

Are you sure you want to change the base?

Implement multi-threading using OhMyThreads and make it differentiable #70

Conversation

pbrehmer commented Sep 27, 2024 • edited Loading

pbrehmer commented Sep 30, 2024

lkdvos commented Sep 30, 2024

pbrehmer commented Oct 1, 2024

lkdvos left a comment

Choose a reason for hiding this comment

lkdvos Oct 5, 2024

Choose a reason for hiding this comment

pbrehmer commented Oct 29, 2024

lkdvos commented Oct 29, 2024

codecov bot commented Oct 30, 2024 • edited Loading

Codecov Report

pbrehmer commented Sep 27, 2024 •

edited

Loading

codecov bot commented Oct 30, 2024 •

edited

Loading