DM_cupy kernel crashing when too many pods are involved

[kahntm]

I am trying to reconstruct a scan with a bunch of diffraction patterns (27.6k) and multiple modes.
While 1,2,3 and 4 modes all work. I can not run the reconstruction with 5 or more modes.
Reducing the cropping size from initially 128 to 64, 32 or even 16 does NOT allow to use more modes.

The problem does scale with the number of GPUs I am using on the machine:

• using 2 GPUs, as described I can run using 1,2,3,4 modes ... and 5+ fails
• using 1 GPU, I can run using

The reconstructions do load, create pods, but fail once they should perform the very first iteration.
The error appears after writing the initial state (iteration 0) to a .ptyr file:

================================================================================
WARNING ptypy - Unable to import optimised cufft version - using cufft with separte callbacks instead
---------------------------------- Autosaving ----------------------------------
Generating copies of probe, object and parameters and runtime
Saving to ----.ptyr

Traceback (most recent call last):
Traceback (most recent call last):
File "----", line 269, in
P.run()
~~~~~^^
File "----/ptypy/ptypy/core/ptycho.py", line 784, in run
self.run(engine=engine)
~~~~~~~~^^^^^^^^^^^^^^^
File "----/ptypy/ptypy/core/ptycho.py", line 713, in run
engine.iterate()
~~~~~~~~~~~~~~^^
File "----/ptypy/ptypy/engines/base.py", line 233, in iterate
self.error = self.engine_iterate(niter_contiguous)
~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py", line 239, in engine_iterate
PROP.fw(aux, aux)
~~~~~~~^^^^^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/kernels.py", line 82, in _fw
self._fft1.ft(x, y)
~~~~~~~~~~~~~^^^^^^
File "----", line 269, in
P.run()
~~~~~^^
File ----/ptypy/ptypy/core/ptycho.py", line 784, in run
self.run(engine=engine)
~~~~~~~~^^^^^^^^^^^^^^^
File "----/ptypy/ptypy/core/ptycho.py", line 713, in run
engine.iterate()
~~~~~~~~~~~~~~^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/cufft.py", line 175, in _ft
self._prefilt(x, y)
~~~~~~~~~~~~~^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/cufft.py", line 152, in _prefilt
self.pre_fft_knl(grid=self.grid,
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
block=self.block,
^^^^^^^^^^^^^^^^^
...<3 lines>...
np.int32(self.arr_shape[0]),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
np.int32(self.arr_shape[1])))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "cupy/_core/raw.pyx", line 93, in cupy._core.raw.RawKernel.call
File "cupy/cuda/function.pyx", line 223, in cupy.cuda.function.Function.call
File "cupy/cuda/function.pyx", line 205, in cupy.cuda.function._launch
File "cupy_backends/cuda/api/driver.pyx", line 273, in cupy_backends.cuda.api.driver.launchKernel
File "cupy_backends/cuda/api/driver.pyx", line 63, in cupy_backends.cuda.api.driver.check_status
cupy_backends.cuda.api.driver.CUDADriverError: CUDA_ERROR_INVALID_VALUE: invalid argument
File "----/ptypy/ptypy/engines/base.py", line 233, in iterate
self.error = self.engine_iterate(niter_contiguous)
~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py", line 239, in engine_iterate
PROP.fw(aux, aux)
~~~~~~~^^^^^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/kernels.py", line 82, in _fw
self._fft1.ft(x, y)
~~~~~~~~~~~~~^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/cufft.py", line 175, in _ft
self._prefilt(x, y)
~~~~~~~~~~~~~^^^^^^
File "----/ptypy/ptypy/accelerate/cuda_cupy/cufft.py", line 152, in _prefilt
self.pre_fft_knl(grid=self.grid,
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
block=self.block,
^^^^^^^^^^^^^^^^^
...<3 lines>...
np.int32(self.arr_shape[0]),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
np.int32(self.arr_shape[1])))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "cupy/_core/raw.pyx", line 93, in cupy._core.raw.RawKernel.call
File "cupy/cuda/function.pyx", line 223, in cupy.cuda.function.Function.call
File "cupy/cuda/function.pyx", line 205, in cupy.cuda.function._launch
File "cupy_backends/cuda/api/driver.pyx", line 273, in cupy_backends.cuda.api.driver.launchKernel
File "cupy_backends/cuda/api/driver.pyx", line 63, in cupy_backends.cuda.api.driver.check_status
cupy_backends.cuda.api.driver.CUDADriverError: CUDA_ERROR_INVALID_VALUE: invalid argument

(/conda_envs/env_ptypy) user@node:path$

Is there some kind of minimum size a pod blocks on GPU memory?
Or is there a maximum number of PODs that can be addressed on a GPU?

I know that running on a node with even more GPUs or across nodes would solve this issue. But it s frustrating that just running smaller cropping does not at all seem to reduce the problem / allow me to run scans with more pods on the same hardware.

In case it matters... With the scan I have at hand I can run 55.200 pods per GPU, but 69 000 PODs per GPU fails. That smells like 2**16 = 65536 is the limit.

[kahntm]

The symptoms are identical for using 2xA40 and 2xL40S

[kahntm]

Further digging makes it seems like the limit is the CUDA maximum grid dimensions size limit.

[daurer]

@kahntm thanks for reporting this. I managed to reproduce this with a simple moonflower example using

p.frames_per_block = 100000
p.scans.MF.data.num_frames = 1340
p.scans.MF.coherence = u.Param(num_probe_modes=70)

this creates 70490 pods which then creates an exit array that has shape (70490,...,...) which seems to cause issues for CUDA.
While I am further investigating the issue, a simple fix is to reduce the frames_per_block such that frames_per_block x num_probe_modes < 60k

[daurer]

And yes, indeed. The limit is imposed by the maximum nr. of blocks per grid in CUDA which is 2^31-1 / 65,535 / 65,535 (for X/Y/Z).
For the batched_multiply.cu kernel (and possibly others) we are ligning up the batchsize with the X dimension - therefore a batchsize larger than 65,535 will run into this limit.

[kahntm]

I guess using the first axis as list index would limit us to diffraction patterns of 256 x 256 px.

[kahntm]

Something about the "frames_per_block x num_probe_modes < 60k" is not working right:

I have a scan with 40k diffraction patterns and 6 modes (240k PODs overall) that I run on 2 GPUs (mpirun -n 2 ...).
When I set p.scan.frames_per_block to 10000 (like you suggested and as it seems logical to me), I get processes that create 120k PODs each. Only when I set frames_per_block to 5000 I get the 60k pods per process that I expected when actually using 10k frames per block and 6 modes.
But in either case ... the reconstruction fails exactly as described above.

When I start it on a single GPU (mpirun -n 1 ...), then 10k frames_per_block and 6 modes actually results in processes that crate 60k pods each. And now the actual reconstruction never even starts... I am stuck with the "End of scan reached" message just before the first engine should start.
Update... the reconstruction proceeded to fail just like those before... even though everywhere the list of PODs is 60k entries long.

[kahntm]

To summarize:
Somehow using multiple GPUs messes up the way the number of PODs created each process is calculated.
Even with <60 PODs per process the reconstructions still fail as described above.
When I limit the number of diffraction patterns (or modes) used ... then the reconstruction runs.

[bjoernenders]

One thing to note is that still every process creates all PODs. So to reach 10000 diff frames in the first block for each rank , every rank needs to process 20000 frames (setting half of them inactive). That makes it 120 PODs each, with two blocks of 60000 depth, one on each rank. For frames_per_block to 5000, all of that is halved.

[bjoernenders]

@kahntm Could you maybe paste your recon script in here?

[daurer]

@kahntm can you confirm which scanmodel you are using? Should be BlockFull and not Full.

The issue here (as Bjoern already mentioned) is actually not so much the nr. of PODs per process but rather the size of each block that is transferred to the GPU for calculations. If using the Full model that block size is always going to be the full size of your problem (nr. of frames times nr. of modes divided by nr. of processes). Only with the BlockFull model the nr. of frames per block parameter is recognised and the maximum size of a block is going to be frames_per_block times num_probe_modes.

I have managed to reproduce similar to your numbers mentioned above and it works on 2 GPUs with BlockFull but it fails with Full.

[kahntm]

I confirm... it was 'Full' instead of 'BlockFull'.
When using the right model, 10k frames per block and 6 modes is running (not failing on the first iteration).

[bjoernenders]

BlockFull works best if there are at least a few blocks so it can mask the transfer latency from host to device. I would go with something like >5 blocks

[kahntm]

Just to confirm... the number of blocks I end up with is the total number of frames divided by the frames_per_block. And the number of PODs per block is frames_per_block x number of probe modes. And the later should be < 2*16

[bjoernenders]

You got it!