Assigning different accelerator count per host does not work

Assigning different accelerator count per host, where the total number of accelerators is not divisible by host count (for eg, 9 accelerators, 2 hosts, 5 & 4 accelerators on each host) does not work. If the accelerator count is divisible by host count, it works,

Specifically following command fails

mlpstorage training run --hosts newmlvm1:5,newmlvm2:4 --model unet3d --data-dir /mnt/training/ --params reader.read_threads=20 dataset.num_files_train=25000 dataset.num_subfolders_train=4 reader.odirect=true checkpoint.checkpoint_folder=/mnt/training --client-host-memory-in-gb 192 --num-accelerators 9 --accelerator-type h100 --checkpoint-folder /mnt/training

Error is

  File "/home/ubuntu/.venvs/myenv/lib/python3.12/site-packages/dlio_benchmark/main.py", line 459, in run_benchmark
    benchmark = DLIOBenchmark(cfg['workload'])
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/.venvs/myenv/lib/python3.12/site-packages/dlio_benchmark/main.py", line 153, in __init__
    self.stats = StatsCounter()
                 ^^^^^^^^^^^^^^
  File "/home/ubuntu/.venvs/myenv/lib/python3.12/site-packages/dlio_benchmark/utils/statscounter.py", line 105, in __init__
    self.MPI.comm().Reduce(host_memory, host_memory_agg, op=MPI.SUM, root=0)
  File "src/mpi4py/MPI.src/Comm.pyx", line 1100, in mpi4py.MPI.Comm.Reduce
  File "/home/ubuntu/.venvs/myenv/lib/python3.12/site-packages/dlio_benchmark/main.py", line 459, in run_benchmark
    benchmark = DLIOBenchmark(cfg['workload'])
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/.venvs/myenv/lib/python3.12/site-packages/dlio_benchmark/main.py", line 153, in __init__
    self.stats = StatsCounter()
                 ^^^^^^^^^^^^^^
  File "/home/ubuntu/.venvs/myenv/lib/python3.12/site-packages/dlio_benchmark/utils/statscounter.py", line 105, in __init__
    self.MPI.comm().Reduce(host_memory, host_memory_agg, op=MPI.SUM, root=0)
mpi4py.MPI.Exception: MPI_ERR_TRUNCATE: message truncated
  File "src/mpi4py/MPI.src/Comm.pyx", line 1100, in mpi4py.MPI.Comm.Reduce
mpi4py.MPI.Exception: MPI_ERR_TRUNCATE: message truncated

Following one works.

mlpstorage training run --hosts newmlvm1:4,newmlvm2:6 --model unet3d --data-dir /mnt/training/ --params reader.read_threads=20 dataset.num_files_train=25000 dataset.num_subfolders_train=4 reader.odirect=true checkpoint.checkpoint_folder=/mnt/training --client-host-memory-in-gb 192 --num-accelerators 10 --accelerator-type h100 --checkpoint-folder /mnt/training