MPI bug when multiple GPUs are used per calculation

[HershGupta404]

When running a repex job for 2 systems on separate nodes with 2 GPUs per system, there appears to be uneven mixing of states depending on which GPU number the replica is. All of the even-numbered replicas (0 GPU) appear to be mixing decently, but the odd-numbered replicas (1 GPU) appear to not be mixing properly. The input files are below. The output files are at /data/chodera/hgupta/repex_flatbottom_nal_md2/no_restraint/experiments. I've also attached a script that reads the total number of states visited for each replica, and a sample output for data at /data/chodera/hgupta/repex_flatbottom_nal_md2/no_restraint/experiments/experiment-neg-MD2.

input_files.zip

count_states.zip

environment.yml.zip

[HershGupta404]

I'm going to do a few more runs this week to check replicability. I'll post the results below.

[HershGupta404]

I ran same system above on shorter runs (500 iterations instead of 1500) with 1 GPU or 2 GPUs. It appears that the 2 GPU run is not mixing properly. See data in files attached.
state_count_sing.txt
state_count_mult.txt

[andrrizzi]

I can confirm I can reproduce this. I'll dig into what could be happening as soon as I have time.

[jchodera]

Great! This should be a high priority bug.

[HershGupta404]

I haven't had much time to work on this due to school, but one thing I've noticed is that approximately the same amount of mixing is reported as occurring from the log files between a single GPU and multiple GPU run.

[andrrizzi]

Thanks! Apologies I haven't been able to look into this yet, but other projects are keeping me busy. This is my next YANK-related thing to look at as soon as I can spare some time.

[andrrizzi]

Since we moved the multistate module to openmmtools, I've moved the discussion here: choderalab/openmmtools#449. I would still leave this open until we figure this out.

[jchodera]

We believe this is addressed (see choderalab/openmmtools#449 (comment)), though we are unclear what resolved the issue.