⛰️[EPIC]: CorrDiff integration

[mnabian]

Placeholder for tracking the issues related to integrating the ResDiff model

Integration Tasks

Preview Give feedback

1. Use Modulus Docker image for ResDiff (ResDiff integration: Add dependencies #266)
2. Compare the results generated with stochastic and deterministic samplers (ResDiff integration: Enabling deterministic sampling #260)
3. Use Modulus logging utilities (ResDiff integration: Switch to Modulus Logging #251)
4. Use Modulus distributed utilities (ResDiff Integration: Switch to Modulus Distributed Util #247)
5. Use Modulus checkpointing (switch to modulus checkpointing for inference #282)
6. Use hydra configs (ResDiff Integration: Config Management #268)
7. Package and import the stochastic sampler (ResDiff integration: switch to pip install of edm sampler #250)
8. Add profiling markers (CorrDiff: Adding NVTX markers and removing redundant regression model inference #306)
9. Black formatting (ResDiff integration: Black formatting #257)
10. Linting (Resdiff Integration: Linting and clean up (Part I) #262)
11. Code cleanup (Resdiff Integration: Linting and clean up (Part I) #262, RessDiff Integration: Remove persistence, import models & loss functions from Modulus, move files to examples folder, more code cleanup #274)
12. Remove dependency to pickle due to security issues
13. Remove the persistance decorator (RessDiff Integration: Remove persistence, import models & loss functions from Modulus, move files to examples folder, more code cleanup #274)
14. Update README (ResDiff integration: Update README, enable FP16 for inference, update inference defaults #288)
15. Unit tests (ResDiff integration: Enable batch size > 1 & multi-GPU for inference, sequential model loading, bug fixes, more unit tests #281)
16. Docstrings (RessDiff Integration: Remove persistence, import models & loss functions from Modulus, move files to examples folder, more code cleanup #274)
17. Move files out of experimental (RessDiff Integration: Remove persistence, import models & loss functions from Modulus, move files to examples folder, more code cleanup #274)

Issues

Preview Give feedback

1. 'distributed.py' has to be removed but that breaks loading the existing checkpoints (RessDiff Integration: Remove persistence, import models & loss functions from Modulus, move files to examples folder, more code cleanup #274)

Performance Optimization

Preview Give feedback

1. Add support for CUDA Graphs (will not do)
2. Support for fp16 in inference (ResDiff integration: Update README, enable FP16 for inference, update inference defaults #288)
3. Enable multi-GPU inference (ResDiff integration: Enable batch size > 1 & multi-GPU for inference, sequential model loading, bug fixes, more unit tests #281)
4. Enable batch_size > 1 for inference (ResDiff integration: Enable batch size > 1 & multi-GPU for inference, sequential model loading, bug fixes, more unit tests #281)

Logging enhancement

Preview Give feedback

1. Add W&B logging (CorrDiff integration: Support wandb logging #316)

Dataset

Preview Give feedback

1. Get permission to publish the CWA dataset
2. Trim the Zarr dataset to include only the four target variables
3. Publish the dataset on NGC (CorrDiff integration: Adding link to the CWA dataset #330)

Validation

Preview Give feedback

1. Verify CorrDiff runs on multiple nodes

[jleinonen]

I'm gathering here the suggestions from our discussion:

Code structuring

• Split generate.py, train.py and training_loop.py into functions/classes that handle more well-defined tasks
• Reduce redundancy between the above files
• Clean up handling of config files
• Specify clearly in an interface what interface data loaders need to follow, to facilitate adding new loaders
• Same as above, but for models (e.g. U-Nets)
• Interoperability with diagnostic model code base (TBA)

Optimizations

• AMP support
• Support distributing generation over time steps as well as ensemble members

[mnabian]

Moving these suggested items to a new epic for CorrDiff enhancements: #353

[mnabian]

Epic Completed!