Paralellized the cpus>1, starting_points=1 situation (CumulativeModel)

[rickdott]

Parallelized estim_probs() when no parallelization occurs elsewhere (cpus >1, starting_points = 1).

I think the current implementation would also work with kfold != 1, since that sets cpus=1 on the EventModel.fit() call, which enters the serial execution.

This is however now the fourth 'level' of parallelization:

• folds
• starting points
• groups
• trials

Since all levels end up running estim_probs() I propose we unify parallelization to only run estim_probs() in parallel, and thus loop serially over folds, starting_points, and groups. Theoretically parallelizing at a higher level would be optimal, but practically I think it would often result in processes with unequal finishing times (due to different group sizes for example).

I tested 2, 4, 10, 20 participant datasets (~800 trials --> ~21400 trials) and 1, 2, 4, 8 cores:

Take the RAM results with a grain of salt, I don't think it's actually using less RAM, its counting the shared memory conservatively. Visual inspection of the system monitor showed that it was equal-ish.

[rickdott]

I think the module-level functions are ugly, but they do need to be instantiated/read before anything uses them. If sometime in the future we create workers in the __init__() function this needs to be above EventModel

[GWeindel]

The speed-up is impressive and definitely needed, also RAM usage management is important as the previous parallelization attempt did not manage that very well.

Because of that last point and the other estimation method that come in with their own parallelization (see #324) we need 1) to move this to the incoming estimator module (#244) so that this chunking sits for the EM estimator and 2) to scrap the other parallelization options to only retain that one that. While the former parallel loops had the advantage of being in higher level functions they, as you say, do not use the parallelization optimally because of waiting times. I also think this is more coherent from a user perspective (cpus > 1 always works the same).

How should we continue? I can help but only later in the month.

[rickdott]

I'll need to get used to the new estimator anyway, so I can add it as a PR to be merged after #244 ? I'll unify the parallelization approaches.

[GWeindel]

Awesome! The PR is merged so you can rebase/update

[rickdott]

That was easier to get used to than I thought! Props to @kiante-fernandez (maybe you also want to have a look, since this now also touches your code).

I think I removed all other parallelization pathways.

[GWeindel]

There is still parallelization for the eliminative and cumulative methods I wonder whether we should keep them or not. The data passed to workers is the same so there's no problems like in the group parallel loop, on the other end it might again make sense for consistency to parallelize always at the same level also.

Could you try timing eliminative with workers prior to EventModel call vs after? I can do later also of course.

[kiante-fernandez]

That was easier to get used to than I thought! Props to @kiante-fernandez (maybe you also want to have a look, since this now also touches your code).

I think I removed all other parallelization pathways.

I will take a look at this tomorrow. My biggest concern is understanding how the parallelization will play with JAX. Which we will need to prevent out of the box deadlocks as we prep to get the the pymc estimators running.

[GWeindel]

That was easier to get used to than I thought! Props to @kiante-fernandez (maybe you also want to have a look, since this now also touches your code).
I think I removed all other parallelization pathways.

I will take a look at this tomorrow. My biggest concern is understanding how the parallelization will play with JAX. Which we will need to prevent out of the box deadlocks as we prep to get the the pymc estimators running.

I guess with sampling methods you can't parallelize that way so you can go for the chain level but that's not necessarily the best way to use large sets of cores. I think we need to have parallelization depending on the estimator?

[rickdott]

I found the remaining parallelization over folds that I missed. Replacing it with within-fold parallelization is a bit inefficient as folds are a split over trials, and parallelization further splits over trials again. With normal usecases (high amounts of trials, relatively low amount of folds) it is still a very good speedup though (theoretically).

[rickdott]

Regarding how it would work with JAX/pymc: I'm not familiar with that, but presumably we can have the estimator call EM with cpus=1 if necessary? This should always be a serial execution, which is useful if something up the chain (samplers?) is itself parallel

[rickdott]

Oops, found a mistake that made the entire parallelization not work. I overreacted to the kfold change, it parallelizes fine even at small set sizes :)

[rickdott]

Last commit speeds up PCA by ~1.5x, should grow linearly with trial count

[kiante-fernandez]

Thanks for reworking this. Creating the pool in fit rather than inside _estim_probs_groups seems right. Before, it was being built and torn down once per EM iteration, plus once per step halving, so a single fit was forking hundreds of times. Once per fit is what we want, and dropping the dead cpus and _n_chunks attributes along the way is nice.

I think the main thing left is that the pool is created in EventModel.fit for whatever estimator is passed, but only EM reaches chunked_estim_probs. That means a sampler forks workers it never uses, and it also means the paths that go through the estimator directly, or through transform, get no pool but still pay by having the chunking overhead. Moving the pool into EMEstimator.fit would settle both issues, since it would only exist where it is used.

Otherwise nothing here is blocking from my side. The rest of my comments are smaller. And the one on pca.py is only about scope rather than the change itself.

One thing this notes for stuff down the pipe is that the sampler imports JAX at the module level, and JAX warns that forking after it initializes can deadlock: for reference I have seen that it can throw something like os.fork() is incompatible with multithreaded code, and JAX is multithreaded (see _at_fork in jax/_src/xla_bridge.py, and jax-ml/jax#1805). Note that in the little tests I ran Note I had only this pop up sometimes but it is not something I want users hitting by default.
I will need to make sure the MCMC path calls EM with cpus=1. Pinning the start method in one place would make that robust, rather than something each caller has to remember.

[kiante-fernandez]

See other comments but holding the pool on the model is what forces __getstate__, and that changes pickling for every consumer of EventModel, not just the workers. There is also no matching __setstate__, so a worker's model ends up with no _pool attribute at all rather than None.

[kiante-fernandez]

The pool is created for whatever estimator is passed, but only EM reaches chunked_estim_probs. fit(..., estimator=MCMCEstimator(...), cpus=4) forks four workers that are never used, and since importing the sampler pulls JAX in at module level, that is also the fork-after-JAX case JAX warns about. Creating it inside EMEstimator.fit instead would mean it only exists for the estimator that wants it.

[GWeindel]

Yes parallelization should be estimator specific

[kiante-fernandez]

Also mp.Pool takes the platform default start method (?), so a fork on Linux and spawn on macOS and Windows. Might be worth pinning explicitly with mp.get_context(...), ideally through one helper shared with the four other mp.Pool sites in hmp/io, so the choice is made in one place.

[kiante-fernandez]

This builds a full-length index array, and estim_probs distinguishes a mask from indices by length, so it reads this back as a boolean mask and drops trial 0. I think that means that no caller hits.
Passing subset_epochs straight through would avoid having two normalization rules for one concept. For ref see tests/gen_data/generate_reference.py:86

[kiante-fernandez]

self is pickled once per chunk on every call. Now that the pool is long lived, the model could go into _WORKER_DATA alongside pattern_data in _init_worker, leaving only the parameter arrays and the chunk indices on the wire.

[kiante-fernandez]

transform(cpus=4) never creates a pool, so self._pool is None here while cpus > 1, and we split the trials, run each chunk, then pad and reassemble, which is strictly more work than a single estim_probs call and no faster (?). The same applies to anything calling EMEstimator.fit directly rather than through EventModel.fit, which is what the sampler's EM seeding path does.

[kiante-fernandez]

this can go too since its only caller gets killed in the PR?

[GWeindel]

yes

[GWeindel]

This looks good so far but indeed parallelization must be estimator specific.

Regarding deadlocks the previous buggy implementations (my bad) always lead me there using the default method with linux. If we can use the parallel refactoring in this PR to have a clear forking/method strategy that would be great. I guess @kiante-fernandez comments could be the way to go but I'm a bit lost when it comes to multiprocessing TBH.

So I think this PR still needs to address those two points (+ EM_star). The other ones would make sense in this PR, but also be the object of other PRs + issues to keep track if needed.

[GWeindel]

Yes parallelization should be estimator specific

[GWeindel]

yes

[rickdott]

Thanks for your feedback both! I considered and/or implemented all points.

Sorting trials by duration speedup
See the comment I added to Kiante's review comment

multiprocessing context
I am developing on ubuntu, which (with my versions of packages/Python) defaults to fork
(See https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods for more information)
I added this function to utils.py:

def _get_mp_context():
    available_methods = mp.get_all_start_methods()
    for method in ["fork", "forkserver", "spawn"]:
        if method in available_methods:
            return mp.get_context(method)

Which is used everywhere mp.Pool was used.
With a small dataset, fork was much faster than spawn due to how spawn pickles/sends data to processes. Spawn did not profit from multiprocessing at all.

I think this is mostly some 'base' costs of spawn and that this will lessen with bigger data. I'll run a bigger test and comment on this PR when it is done. For now I think this is a sensible default though.

[rickdott]

Results of the test, spawn indeed has a 'base' time added to it, but still grows nicely with dataset size.

[rickdott]

Sorry to put another PCA commit here instead of in a new PR, just did it as the previous one was in here as well. The exact same output as before, just 4.5x faster (with 13k trials from 45s to 9s for just the PCA). Should be even more with bigger datasets.

[GWeindel]

Sorry I did not finish my review, wanted to continue on the phone but I can't resume it. It all looks good from what I could see so far not sure the structure is estimator specific enough but if you're confident enough you can merge. We want that parallelization stable by 1.0.0 and we can rework on that before if needed

[kiante-fernandez]

I found a problem with the chunked likelihood that I think needs addressing before this merges. The likelihood returned by a parallel fit depends on the number of cpus.

The cause is in how each chunk computes its pmf. estim_probs takes max_duration from the trials in its own subset, and distribution_pdf normalizes the pmf over that window with p = p / np.sum(p). So a chunk that does not contain the longest trial normalizes over a shorter window than the serial computation, and its trial likelihoods come out different. With my test data below, about 22% of the pmf mass sits beyond the short chunk's window, so the effect is not just some float point noise.

Here is a minimal reproducer against the current branch head:

import numpy as np
import xarray as xr

from hmp.models import EventModel
from hmp.patterndata import PatternData
from hmp.patterns import HalfSine

# 10 short trials (60 samples) then 10 long (200): np.array_split with 2
# chunks puts all short trials in chunk 0, whose local max_duration is 60.
durations_vals = np.array([60] * 10 + [200] * 10)
boundaries = durations_vals.cumsum()
starts = np.roll(boundaries, 1)
starts[0] = 0

durations = xr.DataArray(durations_vals, dims=("trial",),
                         coords={"trial": np.arange(len(durations_vals))})
pattern = HalfSine()
cross_corr = np.random.default_rng(0).normal(0, 0.3, size=(boundaries[-1], 3))

data = PatternData(durations=durations, starts=starts, ends=boundaries - 1,
                   sfreq=100.0, pattern=pattern, template=pattern.template,
                   cross_corr=cross_corr)

ref = None
for cpus in [1, 2, 3, 4, 5]:
    m = EventModel(n_events=2)
    m.fit(data, verbose=False, cpus=cpus)
    chunk_maxes = [int(durations_vals[c].max())
                   for c in np.array_split(np.arange(len(durations_vals)), cpus)]
    if ref is None:
        ref = m
    dtp = np.nanmax(np.abs(m.time_pars - ref.time_pars))
    print(f"cpus={cpus}  lkh={m.lkhs:.10f}  chunk max_durs={chunk_maxes}  "
          f"max|time_pars - cpus1|={dtp:.6f}")

Output on my machine:

cpus=1  lkh=-115.6418511154  chunk max_durs=[200]  max|time_pars - cpus1|=0.000000
cpus=2  lkh=-107.5727835367  chunk max_durs=[60, 200]  max|time_pars - cpus1|=0.203896
cpus=3  lkh=-109.9952627951  chunk max_durs=[60, 200, 200]  max|time_pars - cpus1|=0.203896
cpus=4  lkh=-107.5727835367  chunk max_durs=[60, 60, 200, 200]  max|time_pars - cpus1|=0.203896
cpus=5  lkh=-109.1877697089  chunk max_durs=[60, 60, 200, 200, 200]  max|time_pars - cpus1|=0.203896

Four different likelihoods from five cpu settings on identical data, and the fitted time_pars differ from the serial fit. You can see the likelihood track the chunk layout: it shifts by exactly how many short trials landed in a chunk whose local max is 60, which is why cpus=2 and cpus=4 agree with each other. The EM stopping rule takes in this likelihood, so tolerance checks, selection across starting points, and any model comparison downstream all become functions of cpu count.

test_multiprocessing_matches_serial does not catch this because its data is two trials of the same duration (71 samples), so every chunk's local max equals the global max and the comparison passes trivially. Making the durations in that test heterogeneous would turn it into a proper guard for this.

The good news is that the chunking itself seems fine. I think the fix is to compute max_duration once over the full subset in chunked_estim_probs and pass it into each per chunk estim_probs call. Test on your end and let me know what you think

[rickdott]

Great catch! This fix outputs:

cpus=1  lkh=-115.6418511154  chunk max_durs=[200]  max|time_pars - cpus1|=0.000000
cpus=2  lkh=-115.6418511154  chunk max_durs=[60, 200]  max|time_pars - cpus1|=0.000000
cpus=3  lkh=-115.6418511154  chunk max_durs=[60, 200, 200]  max|time_pars - cpus1|=0.000000
cpus=4  lkh=-115.6418511154  chunk max_durs=[60, 60, 200, 200]  max|time_pars - cpus1|=0.000000
cpus=5  lkh=-115.6418511154  chunk max_durs=[60, 60, 200, 200, 200]  max|time_pars - cpus1|=0.000000

Which is looking better.

I'll give the estimator-specificity some more thinking before merging. I don't really like passing pool through so many functions, but having it as a class member makes pickling/multiprocessing more difficult.