Draft: Resolve "Basic DDP example"

Closes #23

changed milestone to %Demonstrate beam-parallel optimization of neural network parameters and incident wavefunction with simulated dataset. Show near linear scaling up to 30 beams using 5 Summit nodes

added In Progress label

assigned to @28t

I think this may be a good time to learn how to use PyTorch parametrizations: https://pytorch.org/tutorials/intermediate/parametrizations.html We would have probe_f parametrized by probe_f_real which is NxNx2 and the parametrization would just call torch.view_as_complex(self.probe_f_real). Then we could use probe_f as normal without having to remember to manually do the parametrization every time. Note that I have not used this feature yet though!

Note that parametrizations are only available in 1.9+.

@4jh I am still hitting the RuntimeError: Input tensor data type is not supported for NCCL process group: ComplexDouble Even if I put an exit on at the first line of training_step and remove all complex from the training script still give that error. Also printing the parameters before going to training step, I can see no complex there.

Probably it needs a change beyond the lightning script I think, also not quite get training step works, I think if I exit before go into I shouldn't see a nccl error. Also DDP transfer 1 MB by default hard-coded in, but if no complex on lightning script why there is a nccl error?

I assume some other module is returning a complex during the training, so even if no complex in the lighting script we are still going to hit it, but exiting training step I assume I am only doing forward, so no nccl there I assume

@4jh I started a simple pure pytorch example to debug the DDP errors. Here is my first try: https://code.ornl.gov/ai-ptychography/ptychopath/-/blob/23-basic-ddp-example/examples/train.py, https://code.ornl.gov/ai-ptychography/ptychopath/-/blob/23-basic-ddp-example/examples/runtrain.sh and it is working for single GPU (just runs I mean no convergence test), and going to mutli I am getting ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).

The error was obvious, didn't properly set device to local rank. So not if I use DDP I am getting the complex error from nccl. Interesting if I do manually average gradients https://code.ornl.gov/ai-ptychography/ptychopath/-/blob/23-basic-ddp-example/examples/train.py#L104 I am not getting that error.

added 1 commit

• fc8dc138 - adding pure torch example to better debug the DDP errors

Compare with previous version

added 1 commit

• 3d42796e - fixing device bug in the train.py file

Compare with previous version

added 1 commit

• d9c01612 - correcting torch model

Compare with previous version

If you need commits from another branch it is better to use git cherry-pick, which maintains the commit messages associated with commits and ensures you also get the files that are included in the commits you cherry-pick like tests/PTO3.xtl which is present in !13 so you shouldn't need to hardcode a path like this which will break CI.

changed this line in version 12 of the diff

@xf9 I'm just curious: why do you need the work-in-progress stuff from !13 in this MR?

Hi Jacob, I don't need it now, and I will delete it sometime today. I was thinking about using your class function to read the xtl file from Mark's dataset and generate ground truth scattering potential for comparison. However, it seems that you haven't implemented a function that can generate potential from crystal object. I also tried pyms make potential function. However, that doesn't support xtl files. It can only read file from xyz or p1 files. I don't need make potential function right now, but will need it to generate ground truth for convergence plot in the future.