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

…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 95.45455% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/operators/infinitepepo.jl	0.00%	2 Missing ⚠️
src/PEPSKit.jl	85.71%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/algorithms/ctmrg/ctmrg.jl	`92.30% <100.00%> (+0.95%)`	⬆️
src/algorithms/ctmrg/gaugefix.jl	`94.87% <100.00%> (+0.13%)`	⬆️
src/algorithms/toolbox.jl	`96.72% <100.00%> (+0.05%)`	⬆️
src/environments/ctmrg_environments.jl	`68.24% <100.00%> (+1.34%)`	⬆️
src/states/infinitepeps.jl	`67.77% <100.00%> (+1.49%)`	⬆️
src/utility/diffable_threads.jl	`100.00% <100.00%> (ø)`
src/utility/util.jl	`54.21% <100.00%> (-3.09%)`	⬇️
src/PEPSKit.jl	`87.50% <85.71%> (-12.50%)`	⬇️
src/operators/infinitepepo.jl	`18.51% <0.00%> (-0.35%)`	⬇️

[skip ci]

pbrehmer · 2024-10-30T10:28:49Z

I noticed that the default ntasks is always set to one since during pre-compilation Threads.nthreads always returns one. That's why I defaulted to ntasks=4 since on most machines (if Julia is started with the corresponding number of threads) this will multi-thread CTMRG in the four spatial directions. What do you think about that?

Other than that this seems mergeable, thanks a lot for pinpointing and circumventing the error! It is really weird how this Zygote problem never came up before...

src/PEPSKit.jl

src/utility/diffable_threads.jl

…diffable-threads

… definition

pbrehmer · 2024-11-01T17:11:12Z

Whoops somehow the tests are starting to fail which was not the case before the merge, need to investigate what's happening...

So somehow there are some segmentation faults appearing in the Julia LTS tests probably related to the multi-threaded reverse rule. This doesn't seem to be a problem on the latest Julia channel. What is weird is that even when I run the tests locally on Julia 1.10 they won't fail on my machine. I find this quite confusing.

@lkdvos What do you think is the best way to proceed here? Should we restrict the number of threads for the LTS tests to one?

lkdvos · 2024-11-01T18:00:12Z

I'm definitely not a big fan of fixing problems by just turning off the tests, so let's maybe try an see if we can reproduce this anyways? Given that it happens on all platforms, presumably it is a Julia specific problem, but I would still love to either know what causes it, so we can avoid it, or just fix it in the first place...

I guess the reason it only started turning up now is that we were not testing on julia 1.10 before, and instead had 1.9 an 1.11

pbrehmer · 2024-11-01T19:09:39Z

I'm definitely not a big fan of fixing problems by just turning off the tests, so let's maybe try an see if we can reproduce this anyways?

I definitely agree that turning off tests is not the way to go. However, I haven't managed to reproduce it in any way on Ubuntu with Julia 1.10.6 (and also other Julia versions) - the tests just run through without erroring or causing a segmentation fault, regardless of the number of threads.

Given that it happens on all platforms, presumably it is a Julia specific problem

Weirdly, the examples run through on the LTS on macOS (whereas they do not on Ubuntu and Windows). So there might also be some interaction with the OS and multi-threading?

Unfortunately, I don't have access to macOS or Windows machines so I don't really know how to approach this or narrow this down. Also, I don't have that much more time to spend on this (and probably neither do you), so that's why I was looking for a quick solution :-) Anyways, if you come up with ideas to try out, let me know and I can revisit this next week.

pbrehmer · 2024-11-04T09:58:01Z

Just a quick thought: Could it be that the OhMyThreads multi-threading is interfering with OpenBLAS multi-threading? If the CI runners have only 4 cores in total then choosing JULIA_NUM_THREADS=4 might lead to problems if OpenBLAS also tries to claim more than one thread.

lkdvos · 2024-11-04T11:17:14Z

This is unlikely to be the case (but of course without an actual solution I can't exclude anything), because of Julia's task-based parallelism, and this basically being handled by the operating system. In principle there is nothing stopping you from running 20 tasks on a single threaded setup, the OS will switch between these tasks as usual, and everything should still work the same. Of course, there will only ever be one task running in that case, while you still pay the overhead of the tasks and context switching, so typically this hinders performance a bit, but it shouldn't break anything.

I'll try to spend some more time to investigate this week, I definitely don't feel comfortable releasing something that breaks on the LTS on all platforms, so if we can't figure this out we might have to restrict to v1.11, which I'd rather avoid

pbrehmer · 2024-11-04T15:09:58Z

Oh well, you're right. I just tried to come up with some reason why the tests run through on all my Ubuntu machines on 1.10.6 and why they don't on the CI runners. Also the fact that the macOS examples do run through on 1.10.6 is very weird. This seems to happen deterministically (I reran the CI earlier) but it is somehow platform-specific.

Are the tests running through for you locally on an LTS release? I'm currently stumped on this one.

lkdvos · 2024-11-04T16:00:30Z

If this runs through, I'll re-enable the tests for all OS, and then it is definitely my fault. I think I might have screwed up something with the caching of precompiled data, which might have not included the julia version in the hash, thus possibly leading to these weird kinds of errors. (Just guessing, still need to check)
If this is the case, I'll fix and merge?

pbrehmer · 2024-11-04T16:35:43Z

Alright, I see no more segmentation faults, thanks for taking a look and fixing it! This looks mergeable now :-)

lkdvos · 2024-11-04T16:41:04Z

Let's just wait a second until the tests actually all turn green (seems like they are not running now for some reason that I don't understand)

lkdvos

If everything now turns green, good to go for me

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

src/utility/diffable_threads.jl Outdated Show resolved Hide resolved

src/utility/diffable_threads.jl Outdated Show resolved Hide resolved

src/utility/diffable_threads.jl Outdated Show resolved Hide resolved

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

lkdvos and others added 3 commits October 29, 2024 20:37

formatter

ec24e4b

[skip ci]

Fix set_threading_kwargs and default ntasks

a031f2d

Formatter

5a858d1

pbrehmer and others added 2 commits October 30, 2024 11:48

Fix Defaults docstring

48dd0ef

Proposal for scheduler

fbee392

Fix formatting

029f47b

lkdvos reviewed Oct 31, 2024

View reviewed changes

src/PEPSKit.jl Outdated Show resolved Hide resolved

Add NotGiven as default arg

f7814d7

lkdvos reviewed Nov 1, 2024

View reviewed changes

src/utility/diffable_threads.jl Outdated Show resolved Hide resolved

pbrehmer added 2 commits November 1, 2024 17:40

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

b6fcea2

…diffable-threads

Export set_scheduler! directly from Defaults and fix duplicate method…

3c41ddb

… definition

pbrehmer requested a review from lkdvos November 1, 2024 17:07

lkdvos added 2 commits November 4, 2024 10:53

Change actions to stabilize tests (possibly)

70823a2

reenable tests on other OS

5672629

Fix codecov secret

6b2380d

lkdvos force-pushed the pb-diffable-threads branch 2 times, most recently from 9cebff2 to a9c0b82 Compare November 4, 2024 16:52

attempt to fix tests

3ea4833

lkdvos force-pushed the pb-diffable-threads branch from a9c0b82 to 3ea4833 Compare November 4, 2024 16:52

Update formatter action

4844ea0

lkdvos force-pushed the pb-diffable-threads branch from 46edd69 to 4844ea0 Compare November 4, 2024 17:51

enable utility tests

89f91e6

lkdvos approved these changes Nov 4, 2024

View reviewed changes

lkdvos enabled auto-merge (squash) November 4, 2024 18:12

lkdvos merged commit 3a9eb40 into master Nov 4, 2024
27 checks passed

pbrehmer deleted the pb-diffable-threads branch November 4, 2024 19:20

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

pbrehmer commented Oct 29, 2024

lkdvos commented Oct 29, 2024

codecov bot commented Oct 30, 2024 •

edited

Loading

pbrehmer commented Oct 30, 2024

pbrehmer commented Nov 1, 2024 •

edited

Loading

lkdvos commented Nov 1, 2024

pbrehmer commented Nov 1, 2024

pbrehmer commented Nov 4, 2024

lkdvos commented Nov 4, 2024

pbrehmer commented Nov 4, 2024

lkdvos commented Nov 4, 2024

pbrehmer commented Nov 4, 2024

lkdvos commented Nov 4, 2024

lkdvos left a comment

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

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

pbrehmer commented Oct 29, 2024

lkdvos commented Oct 29, 2024

codecov bot commented Oct 30, 2024 • edited Loading

Codecov Report

pbrehmer commented Oct 30, 2024

pbrehmer commented Nov 1, 2024 • edited Loading

lkdvos commented Nov 1, 2024

pbrehmer commented Nov 1, 2024

pbrehmer commented Nov 4, 2024

lkdvos commented Nov 4, 2024

pbrehmer commented Nov 4, 2024

lkdvos commented Nov 4, 2024

pbrehmer commented Nov 4, 2024

lkdvos commented Nov 4, 2024

lkdvos left a comment

Choose a reason for hiding this comment

pbrehmer commented Sep 27, 2024 •

edited

Loading

codecov bot commented Oct 30, 2024 •

edited

Loading

pbrehmer commented Nov 1, 2024 •

edited

Loading