Single device multi gpu support

[SamSandwich07]

We thought Hyrax already support distributed training across multiple GPUs on one node (computer)... apparently not, surprise!

This PR aims to put that functionality back in by using ignite.distributed's DistributedDataParallel to use as many GPUs as are available on a single device.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[SamSandwich07]

All PR checks finally pass!!! Hooray!!! Requesting a review on the current state of this. Following things are what stand out to me as causes for concern:

1. To pass the GitHub checks, I had to make Hyrax use only 1 device if it was not on cuda which means distributed cpu and mps support is not available. I tried using torch.multiprocessing.cpu_count() but that didn't work on GitHub or locally because the backend has to be gloo and not nccl. I changed that and still ran into issues but I haven't really investigated this deeply.

2. More pressing than issue 1: there is no good system to copy over attributes from the user's model to the DDP model (see pytorch_ignite.py. The short version of this issue is that in the user's train_batch function, they may use several attributes that they have defined, including optimizer, scheduler, grad_clip, etc. We don't necessarily know what these will be named, but in order for the current implementation to work there has to be some way to copy over these attributes to the DDP model or automatically pull them from the underlying model if not found on the DDP model.

   I tried to modify the __getattr__ function to attempt this second method already but it didn't work (weird python semantics :/). For the first way, we could allow the user to simply specify the attributes somewhere or attempt to dig in and figure out how to extract used attributes from a function. I welcome feedback on these and/or alternative ideas.

[codecov]

Codecov Report

❌ Patch coverage is 72.22222% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 65.78%. Comparing base (f3dde70) to head (fec1f6d).

Files with missing lines	Patch %	Lines
src/hyrax/pytorch_ignite.py	56.25%	7 Missing ⚠️
src/hyrax/verbs/train.py	83.33%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #974      +/-   ##
==========================================
- Coverage   65.79%   65.78%   -0.02%     
==========================================
  Files          85       85              
  Lines        8178     8204      +26     
==========================================
+ Hits         5381     5397      +16     
- Misses       2797     2807      +10     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[SamSandwich07]

More info regarding issue 1 from previous comment: even if the backend is update to use gloo for cpu/mps, we run into pickling errors because of nested functions. This happens not only in train.py (which is an easy fix), but even in the model class (including something as simple as hyrax_loopback)

[mtauraso]

Overall this looks pretty good as a first pass.

Everything i'm commenting on could be addressed as follow up PRs

[drewoldag]

I like what I'm seeing. The comments that are already here deserve to be addressed. But this is definitely moving the right direction. I would like to know what happens when you run train with multiple gpus when there is a train/validate split in the dataset.