Stormcast refactor with domain parallel + DiT support, testing, Pydantic configs

[jleinonen]

PhysicsNeMo Pull Request

Refactors the StormCast training code as a follow-up to #1379.

Major additions:

1. Support for domain parallel training, controlled by training.domain_parallel_size. The domain parallelism is handled by ParallelManager in utils/parallel.py. FSDP is now used for parallelism regardless of domain parallel size.
2. Support for DiT for diffusion models, with standard configuration with NATTEN loaded by the setup function on nn.py.
3. Configuration parsing using Pydantic. This eliminates most of the manual validation and default value checking as those are now handled by the Pydantic validation. Hydra is still used for loading configs but afterwards they are passed into Pydantic.
4. A test suite using pytest. Includes tests for both single- and multi-GPU setups. A mock dataset that produces random data is included for testing. Multi-GPU tests can be run as torchrun --standalone --nproc-per-node=<number_of_gpus> --no-python pytest test_training.py.

More minor changes:

1. Logging is now passed through utils/logging.py which handles both Tensorboard and WandB logging depending on which are selected.
2. utils/apex.py enables using Apex LayerNorm with ShardTensor.
3. Split training-time plotting into multiple plots as this allows one to view them as an animation, which makes it easier to see if the model is learning temporal dynamics properly.

TODO / pending:

1. To support domain parallelism, utils/loss.py and utils/sampler.py contain versions of EDMLoss and deterministic_sampler modified to work with ShardTensor. TBD if we want to upstream the support to PhysicsNeMo library.
2. README.md still to be updated. Initial updated README.md is added.
3. inference.py still to be properly checked for domain parallel support.

Description

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.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[pzharrington]

/blossom-ci

[pzharrington]

Nice work! 🚀