GPU parallel and other accelaration methods

[Schuyler-Zixiao-Shi]

Hi Ming and the ptychi team,

I’ve been experimenting with ptychi over the past month and wanted to share a few thoughts on potential acceleration strategies. Yi mentioned that some of these ideas have already been explored (or are actively under development), so I’m mainly posting this to consolidate the discussion and provide some context from our side.

On our end, most of our production runs are on A100 GPUs. In practice, many of our A100s are partitioned into 20 GB MIG instances, so one scenario we are particularly interested in is whether multiple 20 GB instances on the same physical A100 could be used cooperatively for a single ~16 GB dataset.

For reference, a typical dataset for us is 512 × 512 scan positions. At that scale, a single batch size of 512 (or larger) is often feasible. Conceptually, this seems like a workload that could be split across multiple 20 GB GPUs if inter-GPU communication within the same node is efficient. Extending this across multiple A100s would be even more attractive, although I suspect inter-node or inter-GPU communication overhead may quickly become the limiting factor.

In addition, I was curious about the cumulative impact of smaller FFT-related optimizations (e.g. batched FFTs, plan reuse, fixed grid sizes, avoiding unnecessary out-of-place copies). Yi mentioned that some of these have already been tested, and that the speedups were not very significant in isolation. It would still be interesting to understand whether a combination of these techniques might yield a more noticeable improvement.

Thanks a lot for all the work that’s gone into ptychi, and for sharing your experiences so far.

Take care,
Schuyler

[mdw771]

Hi Schuyler,

Thank you for sharing your thoughts and experience. I haven't used a MIG-segmented GPU before, but if the MIG instances are abstracted as if they are multiple physical GPUs from CUDA's perspective, Pty-Chi's AD and LSQML reconstructors can utilize multiple of them. AD uses PyTorch's DDP and LSQML uses a finer-grained distribution and synchronization scheme which is also based on torch.distributed; both work with multiple GPUs. See here for instructions. Note that both are data parallelism, meaning we only divide and distribute diffraction patterns into different GPUs, while object, probe and other buffers are copied on each GPU. However, if the memory consumed by your diffraction patterns is much larger than that used by the rest (which is the case in most scenarios), you can still expect near-linear memory reduction. On the other hand, speed improvement is more moderate when we tested on multiple physical GPUs as the communication overhead tends to outweigh the speedup at some point. You are welcomed to try it in your MIG scenario, but a quick search I did indicates that there isn't a more efficient communication interface between MIG instances, and such communications might be even slower than the cross-PCIe interface between physical GPUs.

We also have support for the JIT capability of PyTorch which is under test right now. It speeds up some routines dominated by element-wise operations and improved the overall speed by 22% in one of the experiment. FFT-side optimization like plan reuse is also something we have thought of, but given our bandwidth limit I don't have a clear timeline in mind.

Ming