Various Corrdiff optimizations for drastic increase of training efficiency#809

Merged

CharlelieLrt merged 21 commits intoNVIDIA:mainfrom

jialusui1102:dev-alicia

May 1, 2025

Commits on Apr 14, 2025

mult-gpu training supported corrdiff optimization

Alicia Sui
authored and
jialusui1102
committed

Commits on Apr 15, 2025

enable mixed precision for val

Alicia Sui
authored and
jialusui1102
committed

Commits on Apr 16, 2025

clean codebase for opt

Alicia Sui
authored and
jialusui1102
committed
add amp_mode aware model architecture
jialusui1102
committed
add None checking for params

Alicia Sui
authored and
jialusui1102
committed
revise datatype casting schema

Alicia Sui
authored and
jialusui1102
committed
Add test cases for corrdiff optimizations

LostnEkko
authored and
jialusui1102
committed
revised from_checkpoint, update tests and CHANGELOG
jialusui1102
committed
Lint and format code properly

LostnEkko
authored and
jialusui1102
committed
add multi-gpu optimization
jialusui1102
committed

Commits on Apr 24, 2025

rebase changes and update tests and configs
jialusui1102
committed

Commits on Apr 29, 2025

merge ResidualLoss and refactored layer and Unet init based on PR review
jialusui1102
committed