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

[simonbyrne]

What's the status of this? I would love to make use of these in ReGen.

[CharlelieLrt]

@simonbyrne it's currently blocked by #790 and under review, but that will be merged in the coming days.
AFAIK the current implementation of ReGen does not support these optimizations and there will be some work required to do so.