Various Corrdiff optimizations for drastic increase of training efficiency

[LostnEkko]

Various Corrdiff optimizations for drastic increase of training efficiency

Description

• Updated CorrDiff training code to support multiple patch iterations
  to amortize regression cost and usage of torch.compile
• Refactored modulus/models/diffusion/layers.py to optimize data type casting workflow,
  avoiding unnecessary casting under autocast mode
• Refactored Conv2d to enable fusion of conv2d with bias addition
• Refactored GroupNorm, UNetBlock, SongUNet, SongUNetPosEmbd to support usage of
  Apex GroupNorm, fusion of activation with GroupNorm, and AMP workflow.
• Updated SongUNetPosEmbd to avoid unnecessary HtoD Memcpy of pos_embd
• Updated from_checkpoint to accommodate usage of Apex GroupNorm
• Refactored CorrDiff NVTX annotation workflow to be configurable

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

[mnabian]

/blossom-ci