Pangu Improvements

[dallasfoster]

Modulus Pull Request

Description

This PR adds the following features/changes to the Pangu model and training script:

1. Configurable number of constant, surface, and atmosphere variables in the model.
2. Configurable number of upsampled and downsampled transformer blocks.
3. Gradient checkpointing support in the Pangu processor (encoder/decoder) layers.
4. Improved training script with improved static capture support, multistep rollout, validation function, and weighted loss function.

Checklist

• [ ] I am familiar with the Contributing Guidelines.
• [x] 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.

[dallasfoster]

/blossom-ci

[dallasfoster]

/blossom-ci

[dallasfoster]

/blossom-ci

[dallasfoster]

/blossom-ci

[dallasfoster]

/blossom-ci

[dallasfoster]

Depends on #660

[pzharrington]

I added some minor suggestions, overall this looks great though. I found the LambdaLR scheme with the custom hydra resolver a bit convoluted, maybe a ConstantLR would be simpler and more readable to achieve the same effect. However it is nice to have the example in there if someone wants to do more custom scheduling.