Distribute halos to ranks balancing particle count

[kyleaoman]

This is an initial attempt (thus: draft) to distribute halos across MPI ranks in a more balanced way. The current approach is a spatial partitioning of the simulation box where each halo is assigned to the "cell" containing its centre. In terms of load balancing, this assumes that the distribution is not clustered, which is inevitably not true (for a worst case consider a zoom snapshot).

Ideally we want a good balance of CPU load that still fits in the memory available to each rank. This PR has smaller ambitions: the goal is just to achieve a reasonable memory balance.

My initial proposal is to take the total particle count, divide that up amongst the available ranks, then iterate the halo list assigning halos to rank 0 until it has reached its share of the particles, then start assigning to rank 1, and so on until all halos are assigned.

Because halos have various particle counts we can't expect a perfect partition from this scheme. My initial approach is to set a tolerance, so each rank is allowed to have a bit (for now hard-coded to 15%) more particles than its nominal share. There is also a rule that the first halo assigned to a rank is allowed regardless of size, this way a single halo that is "too big" for any one rank is not an automatic failure. These rules don't guarantee a successful distribution, so some refinement will be wanted here, but I'd rather have some confidence that a simple scheme is working before making it more complex.

I've got as far as making sure the code compiles and running the test suite. I get a test failure:

11/12 Test #11: test_mergertree ..................***Failed   60.34 sec
Test failed: Subhalo[i].DescendantTrackId == Subhalo[i].TrackId at /cosma/home/durham/dc-oman1/code/HBT-HERONS/unit_tests/test_mergertree.cpp: 63
Test failed: Subhalo[i].DescendantTrackId == Subhalo[i].TrackId at /cosma/home/durham/dc-oman1/code/HBT-HERONS/unit_tests/test_mergertree.cpp: 63
Test failed: Subhalo[i].DescendantTrackId == Subhalo[i].TrackId at /cosma/home/durham/dc-oman1/code/HBT-HERONS/unit_tests/test_mergertree.cpp: 63
Test failed: Subhalo[i].DescendantTrackId == Subhalo[i].TrackId at /cosma/home/durham/dc-oman1/code/HBT-HERONS/unit_tests/test_mergertree.cpp: 63
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------

That also fails on branch non_consecutive_outputs that I've based off of though, so perhaps unrelated to my changes, but some reassurance would be welcome. Would also be good to know whether the test suite tests anything that would pick up issues in the halo assignment, or whether I should think of it as just a sanity check that I haven't completely broken things.

[VictorForouhar]

Improving the domain decomposition is definitely a good idea, thanks for kickstarting the discussion!

For the failed test you get, I believe it is indeed independent from your changes. See #102.

[kyleaoman]

Was mulling it over on my way home and I think that a better starting point would be a sorted list of halos by particle number, then assigning them in descending order in a round robin to the ranks. Once a rank is "full" it can be skipped in the round robin. That should end up quite close to being memory-balanced in most cases and won't do the worst possible job of cpu balance.

[VictorForouhar]

I agree. This way the biggest haloes are guaranteed to be in separate ranks. This is similar to what subfind in gadget4 does.

[VictorForouhar]

Something else that will help with testing this implementation: we are planning to add (debug-enabled) timers to each step in the unbinding of individual subhaloes and haloes, and which MPI rank is analysing them. This will help identify bottlenecks, and we plan on adding it sooner rather than later.

[kyleaoman]

Hi Victor - could I bring this back onto the radar? I've rebased on current main and I think I've proposed a sane round robin-based scheme now. If this looks sensible to you, could you offer some guidance on testing and what else would need to be done to make this a complete PR?

[VictorForouhar]

Sure, I will check your changes. I will probably want to run on a few other simulations on my side to check the performance hit/gain of your changes.

[kyleaoman]

"Other simulations" implies that I've tried running this on a simulation already 😅

[VictorForouhar]

Ah! Well, I meant in io vs unbinding limited simulations. We can also check if it indeed works. The test I was thinking about before is to check how expensive is to do the round-robin communication at each snapshot. Perhaps some big haloes will continuously change ranks because of random ordering and hence introduce unnecessary communication overhead.

Just speaking out loud here, it could be that your implementation prevents this (still need to check the code)

[kyleaoman]

Ah ok, I don't have anything that will try to keep groups on the same ranks. It's easy to do (just start by putting each group on its previous rank as the default option), but since the default will be set by the rank ordering of halos at the time they appear this probably messes up the general concept of the round robin. There are probably tweaks that can help with reducing comms though, like you could do the assignment then swap ranks so that the biggest halo on each rank stays on the same rank as it was at the previous snapshot. First thing is to get a baseline comparison of the existing implementation vs this naive round robin, I guess.

[VictorForouhar]

Run a sanity check on a small 25m7 COLIBRE test box we have. Mass functions are the same as master:

I will now run on FLAMINGO DMO L1000m9 for a larger box test. That run will test how the changes affect I/O limited analysis. I will also see which of the larger COLIBRE boxes we can use, which will test an unbinding-limited analysis.

[kyleaoman]

Great, thanks! Been continuing to think about tweaks - currently halos are assigned in the order (for 4 ranks) 1 2 3 4 1 2 3 4 1 2 3 4 ... A variant would be 1 2 3 4 4 3 2 1 1 2 3 4 ... That would tend to "balance out" rank 1 getting the most massive halo by giving it the least massive in the 2nd round, etc. I think it makes sense to get a baseline comparison to start with but this could be a variant to try out, maybe.

[VictorForouhar]

I think the round-robin assignment will need to be written in a more careful manner, as it is currently done on a local MPI rank level. This means that when assigning MPI ranks based on particle count, there is no information about what assignment other MPI ranks have made.

This was not a problem with the spatial assignment.

[kyleaoman]

Is every rank not holding all halos (e.g. InHaloSizes is not the same on every rank)? If they all get the same input to this function then I think that the result is deterministic and every rank will evaluate all of the assignments for all ranks.

[VictorForouhar]

No, I do not think so. What happens varies depending on the simulation format.

For SWIFT, each MPI rank reads a given section of a snapshot group. It then uses the FoF membership of the particles it loaded to create FoF fragments because they are not guaranteed to have been read by a single MPI only.

These fragments are subsequently grouped and joined in a given MPI rank by ExchangeAndMerge. The final MPI rank where they will be placed and analysed is done by ParticleSnapshot_t::ExchangeHalos.

The end result is that vector<Halo_T> InHalos contains distinct FoFs per rank, and since we build InHaloSizes from that vector, it is not universal across MPI ranks.

Of course, this has an easy solution: make InHaloSizes locally and all gather.

[VictorForouhar]

Relating to the comment above, and more to with the formatting of the code before this PR was created, I find the fact that there are several Exchanges and DecideTargetProcessor confusing. I will try to make them more distinct in an upcoming commit.

[kyleaoman]

Ah ok - indeed gathering first is the solution. Feel free to push to this branch if that's useful at any point.

[VictorForouhar]

Actually, I misread an if condition, so the second step in my last comment does not trigger for swiftsimio. Hence, the changes you made have had no effect on the domain decomposition. Will move them to the ExchangeAndMerge function call, which is the only relevant part for domain decomposition.

[VictorForouhar]

The basic assignment of tasks to haloes is done. I used @kyleaoman's round robin assignment code within the relevant files for swift and apostle. Below is a bar plot showing the particle load across ranks for the COLIBRE test box (25m7).

The maximum particle difference between all ranks is just 16. Of course, using the number of particles as cost estimator is a simplification. A halo with 10 subgroups of 100 particles (1000 particles total in the halo) each is more expensive to analyse than 10 haloes of 100 particles each.

Note that I still need to solve a few issues that cause crashes with MPI communication of the haloes themselves (I modified that part of the code to make it more readable, but now needs verification).

[robjmcgibbon]

I've had a look through the changes, and have a few thoughts before I implement anything.

My understanding is that the current implementation discards all spatial information. The FoFs are sorted globally by particle count and assigned to ranks with no regard for position. This achieves good particle balance, but since the assignment is recomputed from scratch each snapshot, large FoFs can end up on completely different ranks between snapshots. An alternative that would require relatively minor changes to what's already here: start with the spatial COM-based assignment from master (assign each FoF to the rank corresponding to its spatial cell), then greedily rebalance. After the initial assignment, iteratively move the largest FoF from the most overloaded rank to the most underloaded rank until all ranks are within some tolerance of the mean.

I think there is a better option, but it would require considerably more work. The FoF catalogues contain the size and COM of every group, and are small enough that all ranks can read or receive them before touching the particle data. This means you could compute a load-balanced spatial partition before doing any I/O. Each rank could then read only the snapshot cells covering its assigned region. This would eliminate the post-I/O redistribution exchange (almost) entirely, since FoFs would land on the correct rank directly from I/O. The partition would be roughly stable between snapshots because most FoFs don't cross partition boundaries from one output to the next. This is probably out of scope for the current PR though.

[kyleaoman]

I think that your "greedy rebalancing" probably helps quite a bit. A possible follow-on optimization would be to adjust the mapping between partition segments and MPI rank IDs until the minimum number of particles need to be moved. Will illustrate what I mean...

[kyleaoman]

[robjmcgibbon]

I've implemented the "greedy" rebalancing, here's what I found.

Testing

The small colibre box seems to work without issues. I then ran it on snapshots 56-63 (around z=4) of the L0050N1504/Hybrid . Here is part of the log from using master branch

17010373459 particles loaded from snapshot 61 (SnapshotIndex = 61)
Unbinding... Took 2252.91 seconds. Maximum imbalance across ranks was 3.42714.
Snapshot 61 (SnapshotIndex = 61) done. It took 4496.18 seconds.

17010375615 particles loaded from snapshot 62 (SnapshotIndex = 62)
Unbinding... Took 2348.13 seconds. Maximum imbalance across ranks was 3.41354.
Snapshot 62 (SnapshotIndex = 62) done. It took 4678.88 seconds.

17010379249 particles loaded from snapshot 63 (SnapshotIndex = 63)
Unbinding... Took 2611.45 seconds. Maximum imbalance across ranks was 3.59786.
Snapshot 63 (SnapshotIndex = 63) done. It took 5091.1 seconds.

And here is from this branch

17010373459 particles loaded from snapshot 61 (SnapshotIndex = 61)
Unbinding... Took 1257.21 seconds. Maximum imbalance across ranks was 1.86442.
Snapshot 61 (SnapshotIndex = 61) done. It took 3228.97 seconds.

17010375615 particles loaded from snapshot 62 (SnapshotIndex = 62)
Unbinding... Took 1339.11 seconds. Maximum imbalance across ranks was 1.8914.
Snapshot 62 (SnapshotIndex = 62) done. It took 3418.59 seconds.

17010379249 particles loaded from snapshot 63 (SnapshotIndex = 63)
Unbinding... Took 1438.31 seconds. Maximum imbalance across ranks was 1.9105.
Snapshot 63 (SnapshotIndex = 63) done. It took 3603.36 seconds.

The time spent unbinding is considerably reduced. The imbalance also gets worse as we go to lower redshifts, so we'll probably save even more time there. The bound mass function looks fine.

@kyleaoman could you please try the current version of this branch on your run?

Overview of how exchange_and_merge works (so I don't forget)

Each rank computes the particle count and centre of mass of each of its local fragments. These are sent to a temporary "accumulation rank" determined by hashing the HaloId, so all fragments of the same FoF arrive at the same rank and can be summed. Rank 0 then gathers these totals from all accumulation ranks and assigns each FoF a target rank using the spatial COM-based assignment from master, followed by greedy rebalancing (see below). Rank 0 scatters the assignments back to the accumulation ranks, again using the hash routing so each assignment lands on the rank holding that FoF's accumulated data.

At this point, each accumulation rank knows the target for the FoFs that hashed to it, but the ranks that actually own the fragments (and will need to send the particles) don't know yet. Therefore each rank sends fragment metadata to the same hash-based accumulation rank, which stamps each with its target rank and sends them back. Now every rank knows where to send each of its fragments.

Fragments are sorted by target rank and exchanged via MPI, first the halo metadata then the particles. After the exchange, each rank holds all fragments of its assigned FoFs and merges them by concatenating particle lists.

Greedy rebalancing

We start with the spatial cell assignment (used on master already). We then find the ranks with the most and least particles. We select the largest FoF from the overloaded rank that satisfies 2*size < gap (where gap is the difference in particles between the ranks). This condition ensures the move reduces the imbalance without making the previously underloaded rank the new maximum. We continue this loop until the most loaded rank is within 5% of the mean particle count, or when no single FoF is small enough to be moved.

Future changes

As mentioned above I think it would be better to use the FoF catalogues to determine the domain decomposition in advance of loading the particles. We could then make use of the fact the the swift snapshots are sorted to reduce the amount of exchanges required between ranks. However, unbinding dominates the runtime (at least for colibre), meaning the exchanges aren't a major problem, so I'm not planning on implementing this now.

The output from the L50m5h still shows an imbalance. I guess this is because the time to unbind is not linear with the number of particles. I think we should try this balancing for a bit to see how it works for low redshift on the big runs. Also a perfect cpu-time balance probably will not correspond to a perfect memory balance, and I'm not sure which would be better to optimise for.