Add support for multi-GPU

[allevitan]

I think that this is going to be really important, long-term. If there's no way to do multi-GPU, joint ptycho-tomo will be so much less practical.

I'd like to make a feature branch and test out a few patterns for getting it to work.

[yoshikisd]

I've opened up a draft PR #17 that contains a working implementation of multi-GPU support for CDTools (at least on the machine I use). This is a work in progress as I'd like to experiment with a few items listed in the PR. But if you have any comments/suggestions/edits, I'd be more than happy to hear them!

[allevitan]

This is an awesome development! I'll definitely be following along with the progress of the PR and be involved in any discussions about how to make the code patterns for multi-GPU feel natural/mesh with the rest of the codebase

[yoshikisd]

Thanks Abe! Regarding your original comment, you made a point about about how joint ptycho-tomo may not be practical if multi-GPU support is unavailable. Is this mostly due to the large size of the model that must be distributed across several GPU, or the need to accelerate calculations using several GPUs (where the model is capable of living in the memory of one GPU)?

[allevitan]

I was more worried about the computational time issue. When not done with a joint model, all the projections get reconstructed separately, usually that gets broadcast across many GPUs.

However, model size is a problem as well - just storing the volume for e.g. a 2000 pixel tomogram with 64-bit complex floats takes 64 GB. It's not the immediate limiter, though, afaik.

[yoshikisd]

Something I'd like to keep here for memory retention is the fact that someone has looked into a bug related to the use of DDP and complex number types (pytorch/pytorch#158753 (comment)), which seems to have been addressed yesterday in (pytorch/pytorch#166863). The point of the issue is that DDP-wrapping causes the imaginary component of the gradient to be zeroed.

It's not clear if this is exactly the cause of the issue I saw in #17 (comment) because I didn't check the imaginary part of the gradient then. However, if the new commit does solve this issue, then the multi-GPU implementation could be simplified further. I might try experimenting on the nightly version of PyTorch.

That said, the all-reduce implementation in #55 seems to work fine, and another person seems to be using it on NERSC. It would be great if we could get a review on it at some point but I understand everyone's busy.