Running HSSM through slurm scheduler

[mrinmayik]

Hello,

I'm trying to run HSSM on a high performance computing cluster through a SLURM scheduler. The job starts, and I get the following output:

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [v_Intercept, v_percent, v_displacement, v_trialtype, v_Group, v_1|participant_id_sigma, v_1|participant_id_offset, v_1|RunTrial_sigma, v_1|RunTrial_offset, a_Intercept, a_Group, a_1|participant_id_sigma, a_1|participant_id_offset, z_Intercept, z_Group, z_1|participant_id_sigma, z_1|participant_id_offset, z_1|RunTrial_sigma, z_1|RunTrial_offset, t_Intercept, t_Group, t_1|participant_id_sigma, t_1|participant_id_offset, t_1|RunTrial_sigma, t_1|RunTrial_offset]

It's been running for ~10 hours now, and there's no output. When I run the top command in the core it's running on, a python process periodically appears and disappears (here's a video). And it's not using anywhere near 100% of CPU. Is this normal?

When I run the command on my local machine, it runs through in ~5 minutes (and uses 100% of the core). But I need to run it on the cluster so I can increase the number of chains and draws.

This is the sample command I'm using:

DDModel_Organised.sample(
    sampler="mcmc",  # type of sampler to choose, 'nuts_numpyro', 'nuts_blackjax' of default pymc nuts sampler
    cores=NumThreads,  # how many cores to use
    chains=4,  # how many chains to run
    draws=500,  # number of draws from the markov chain
    tune=50,  # number of burn-in samples
    idata_kwargs=dict(log_likelihood=True),  # return log likelihood
    progressbar=False
)

Thank you!
Mrinmayi

[digicosmos86]

Hi Mrinmayi,

There might be a few things that you can try, which would be helpful for us to diagnose the issue, since we cannot reproduce it locally without your specific data.

1. Can you try running a very simple model (like the one in our documentation)? Try sampling with multiple cores as well so we can rule out if the issue is specific to the Slurm environment.
2. Maybe enable the progress bar? Given the complexity of your model it might actually take a lot of time. The progress bar can be helpful in monitoring if sampling is actually making progress.
3. Try the jax sampler and see if it makes any difference.

BTW what is the HSSM version that you are using, and what are the PyMC and bambi versions in the environment in which you are running HSSM on the cluster? Are they consistent with the versions of the packages in your local environment?

Please let us know if any of this helps.

Thanks!
Paul

[mrinmayik]

Thanks a lot for your quick reply!

The version of pymc, bambi, hssm and arviz are the same on my local machine and on the cluster.

Try the jax sampler and see if it makes any difference.

I would do this by specifying the sampler to be nuts_blackjax, correct? jax doesn't seem to be playing nicely with my versions of pymc, I think? I got a bunch of errors (the latest one being window_adaptation() missing 1 required positional argument: 'logdensity_fn') so I'm sticking to mcmc for now.

Can you try running a very simple model (like the one in our documentation)?

I tried running the simplest model:

DDModel_Organised = hssm.HSSM(
    data = TestGoodData,
    model = "ddm")

in 4 chains and 500 draws using 1 core, and it actually ran all the way through in about 10 minutes!

Maybe enable the progress bar?

I had enabled the progress bar before, but it was stuck at 0 draws and 0 chains, making me think the progress bar just wasn't updating in the slurm output. Turns out it was really stuck at the first draw even after 10 hours! With the simpler model the progress bar was being updated, so I'll leave it on now. Thanks!

Try sampling with multiple cores as well so we can rule out if the issue is specific to the Slurm environment.

I had set the number of cores to be equal to the number of chains based on this discussion. Would it help to increase the number of cores beyond that?

Importantly, I tried to make the models increasingly complex in terms of the number of parameters that were being estimated, while monitoring the CPU usage when running these jobs. Here is what I noticed:

• When I ran the simplest model from the documentation with 1 core, the process was showing ~99% CPU usage
• When I ran a slightly more complex model where v varied by condition, trial and participant, and z varied by group, participant and trial, with 4 cores, the 4 processes show ~25% usage each, and the model ran all the way through.
• But when I tried to include estimation of t, the progress bar remains stuck at -| 0.00% [0/2200 00:00<? Sampling 4 chains, 0 divergences], and the 4 processes show ~0.1% CPU usage. This seems suspicious, correct? Why might this be happening?

Thanks again!
Mrinmayi

[digicosmos86]

Hi Mrinmayi,

Sorry for the delayed response. The jax sampler is actually the nuts_numpyro sampler. However, from what you describe, it seems that the reason why your sampling is stuck because of the regressions on parameters such as a and t. We are experiencing some common issues with this and it is our top priority to fix it right now. We will release something in the next few weeks to address this issue. For now, you might try the approx_differentiable likelihood and nuts_numpyro sampler. Mathematically it is slightly more robust than the analytical likelihood. However, it might still misbehave with more complex hierarchical models.

Thanks!
Paul

[mrinmayik]

Hi @digicosmos86

Thanks a lot for your response, and so sorry for the delay in mine!

I played around with the command a bit, and I think I've figured out in what situation the sample command is getting stuck.
This is command I am using:

DDModel_Organised = hssm.HSSM(
    data = DDMData,
    model = "ddm",
    hierarchical=True,
    include = [
        {
            "name": "v",
            "formula": "v ~ 1 + Group + displacement + trialtype + percent + (1|participant_id) + (1|RunTrial)" # + Group + displacement + trialtype 
        },
        {
            "name": "t",
            "formula": "t ~ 1 + Group + (1|participant_id) + (1|RunTrial)"
        },
        {
            "name": "a",
            "formula": "t ~ 1 + Group + (1|participant_id) + (1|RunTrial)"
        },
        {
            "name": "z",
            "formula": "t ~ 1 + Group + (1|participant_id) + (1|RunTrial)"
        }
    ]
)

DDModelOrganised_samples = DDModel_Organised.sample(
    sampler="mcmc",  # type of sampler to choose, 'nuts_numpyro', 'nuts_blackjax' of default pymc nuts sampler
    cores=4,  # how many cores to use
    chains=4,  # how many chains to run
    draws=15000,  # number of draws from the markov chain
    tune=2000,  # number of burn-in samples
    idata_kwargs=dict(log_likelihood=True),  # return log likelihood
    progressbar=True
)

Based on this discussion, I had been setting the number of cores to be equal to the number of chains. Additionally, since I was running this on a computing cluster, I was also able to provide it multiple cpus using the cpus-per-task (see here). I was giving it the same number of cpus as the cores I had specified in the sample command. But this is what causes the sampler to get stuck. When cores is set to 4, but cpus-per-task is set to 1, the job seems to be progressing. As before, when I monitor the CPU usage, I see 4 processes using 25% of the CPU. However, this seemed to slow down the sampling a lot. This job was running for 3 days, and it had only gotten through 9% of the draws (10132 of 102000 draws). Do you have any suggestions for why the job gets stuck when cpus-per-task is greater than 1? Also, does the computation slow down so much because it's running on a single core?

By the way, I also tried using the nuts_numpyro sampler based on your suggestion above, but I get the following error:
AttributeError: module 'jax.numpy' has no attribute 'Product'

Thanks, again, for all your help!
Mrinmayi

[digicosmos86]

HI @mrinmayik,

Sorry for the delay getting back to you. First, it seems that for your a ant t parameters, the regression target is set to t, which might cause hssm to not build the right model for you.

Sampling hierarchical models with multiple regressions can be challenging since different parameters have different ranges. We have made some improvements in hssm 0.2.0 and added some strategies for better sampling of hierarchical models. Please see here for more details.

Thanks, and please let me know if the new version solves your problems.
Paul