[Bug]: Issue with `resplit_` for a (quite) large array

[mrfh92]

What happened?

MPI-error when resplitting a (quite) large array

Code snippet triggering the error

import heat as ht

X = ht.random.randn(6505535, 363, 7, 7, split=1, dtype=ht.float32)

print("data created")

X.resplit_(0)

print("done")

with

#!/bin/bash

#SBATCH --clusters=MYCLUSTER
#SBATCH --partition=CPUPARTITION
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=40
#SBATCH --time=00:10:00
#SBATCH --account=MYACCOUNT

source ~/modules_heat.sh

export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MKL_NUM_THREADS=$SLURM_CPUS_PER_TASK
mpirun python test.py

Interestingly, with 8 MPI-procs per node and 20 cores per task there was no problem anymore.

Error message or erroneous outcome

`mpi4py.MPI.Exception: MPI_ERR_TRUNCATE: message truncated` in `  File "src/mpi4py/MPI.src/Comm.pyx", line 1080, in mpi4py.MPI.Comm.Alltoallw
    comm.Alltoallw(`

Version

main

Python version

Python 3.10.10 [GCC 12.2.0]

PyTorch version

2.6.0+cu126

MPI version

Open MPI 4.1.5, mpi4py 4.0.1

[mrfh92]

Likely, it could be something where we forgot some index measured in int32

[JuanPedroGHM]

That error usually indicates that the buffers were not properly constructed (too small, wrong address, or displacements). Do you have a longer stack trace?

[mrfh92]

Similar:

import heat as ht 

M, N = 2**31, 2

x = ht.random.randn(M, N, dtype=ht.float32, split=1) 
print("Hi, here is the original x:", x.__repr__())

x.resplit_(0)
print("Hi, here is the resplit x:", x.__repr__())

yields

Hi, here is the original x: <DNDarray(MPI-rank: 1, Shape: (2147483648, 2), Split: 1, Local Shape: (2147483648, 1), Device: cpu:0, Dtype: float32)>
Hi, here is the original x: <DNDarray(MPI-rank: 0, Shape: (2147483648, 2), Split: 1, Local Shape: (2147483648, 1), Device: cpu:0, Dtype: float32)>
Traceback (most recent call last):
  File "/home/XXX/my_dev_stuff/resplit/test_resplit.py", line 8, in <module>
Traceback (most recent call last):
  File "/home/XXX/my_dev_stuff/resplit/test_resplit.py", line 8, in <module>
    x.resplit_(0)
  File "/home/XXX/heat/heat/core/dndarray.py", line 1561, in resplit_
    self._axis2axisResplit(
  File "/home/XXX/heat/heat/core/manipulations.py", line 3841, in <lambda>
    target_tile: _axis2axisResplit(
  File "/home/XXX/heat/heat/core/manipulations.py", line 3825, in _axis2axisResplit
    comm.Alltoallw(
  File "/home/XXX/heat/heat/core/communication.py", line 1590, in Alltoallw
    subarray_type = send_datatype.Create_subarray(
  File "src/mpi4py/MPI.src/Datatype.pyx", line 304, in mpi4py.MPI.Datatype.Create_subarray
mpi4py.MPI.Exception: MPI_ERR_ARG: invalid argument of some other kind
    x.resplit_(0)
  File "/home/XXX/heat/heat/core/dndarray.py", line 1561, in resplit_
    self._axis2axisResplit(
  File "/home/XXX/heat/heat/core/manipulations.py", line 3841, in <lambda>
    target_tile: _axis2axisResplit(
  File "/home/XXX/heat/heat/core/manipulations.py", line 3825, in _axis2axisResplit
    comm.Alltoallw(
  File "/home/XXX/heat/heat/core/communication.py", line 1590, in Alltoallw
    subarray_type = send_datatype.Create_subarray(
  File "src/mpi4py/MPI.src/Datatype.pyx", line 304, in mpi4py.MPI.Datatype.Create_subarray
mpi4py.MPI.Exception: MPI_ERR_ARG: invalid argument of some other kind
/opt/miniforge/lib/python3.10/tempfile.py:860: ResourceWarning: Implicitly cleaning up <TemporaryDirectory '/tmp/tmp8kdbd7ti'>
  _warnings.warn(warn_message, ResourceWarning)
/opt/miniforge/lib/python3.10/tempfile.py:860: ResourceWarning: Implicitly cleaning up <TemporaryDirectory '/tmp/tmpyve52kk4'>
  _warnings.warn(warn_message, ResourceWarning)
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------
--------------------------------------------------------------------------
mpirun detected that one or more processes exited with non-zero status, thus causing
the job to be terminated. The first process to do so was:

  Process name: [[28830,1],1]
  Exit code:    1
--------------------------------------------------------------------------

[github-actions]

Branch bugs/1932-_Bug_Issue_with_resplit__for_a_quite_large_array created!

[mrfh92]

@JuanPedroGHM Do you know whether

subarray_type = send_datatype.Create_subarray(
                        lshape, subsizes, substarts, order=MPI.ORDER_C
                    ).Commit()

accepts inputs that are more than maxint32?

[JuanPedroGHM]

It does not, we need to find which of the arguments is larger than int32 max and find a workaround using other custom datatypes.

[mrfh92]

ok, then this explains the error because for lshape, subsizes, substarts we have:

[2147483648, 1] [1073741824, 1] [0, 0]

on both processes.

[JuanPedroGHM]

Yeah, the subtype construction is clearly wrong. Looks like I missed some edge cases on the subarray construction.