e3sm/oscar/spatial-parallelism#4

Draft

mahf708 wants to merge 46 commits intoe3sm/old-mainfrom

e3sm/oscar/spatial-parallelism

Commits on Nov 15, 2025

Adding the comm file from Makani and making the necessary changes to run ACE using PhysicsNemo. It works, but it does not utilize spatial parallelism yet.

odiazib
authored and
mahf708
committed
Implement a split of the dataset for spatial parallelism and create a unit test that divides the dataset into four parts, subsequently comparing the results with the original dataset.

odiazib
authored and
mahf708
committed
Adding the necessary files from Makani for spatial parallelism.

odiazib
authored and
mahf708
committed
Adding the necessary files from Makani for spatial parallelism.

odiazib
authored and
mahf708
committed
Adding spatial parallelism to the model, layers, and FFT. Testing this implementation using unit tests based on those developed by Makani.

odiazib
authored and
mahf708
committed
Adding NVIDIA PhysicsNemo

odiazib
authored and
mahf708
committed
Move the block of code from conf to the xarray class initialization.

odiazib
authored and
mahf708
committed
Reintroduce the logic to run the case in serial and on the CPU.

odiazib
authored and
mahf708
committed
Split domain for spatial parallelism.

odiazib
authored and
mahf708
committed
Moving code to distributed class.

odiazib
authored and
mahf708
committed
Fixing the xarray test with spatial parallelism.

odiazib
authored and
mahf708
committed
Spatial distributed model test is working, but tolerance error is 1e-3.

odiazib
authored and
mahf708
committed
Adding init_gradient_reduction_hooks to the model and testing training with spatial parallelism. The unit tests ran, but I have not checked for correctness.

odiazib
authored and
mahf708
committed
Fixing the dataset reader to make training work for the E3SM case.

odiazib
authored and
mahf708
committed
Getting spatial parallelism input parameters from the CLI.

odiazib
authored and
mahf708
committed
Saving and loading checkpoints when a model uses spatial parallelism.

odiazib
authored and
mahf708
committed
Moving init_gradient_reduction_hooks to the Distributed class.

odiazib
authored and
mahf708
committed
Fix initialization and checkpoint handling in distribute class

odiazib
authored and
mahf708
committed
Only use comm.cleanup() if spatial parallelism is enabled.

odiazib
authored and
mahf708
committed
Removing old code.

odiazib
authored and
mahf708
committed
Adding recommendations for PR review and deleting old code.

odiazib
authored and
mahf708
committed
Adding recommendations for PR review

odiazib
authored and
mahf708
committed
Apply suggestion from @mcgibbon

authored and
mahf708
committed
Removing 'comm' and moving the routine to the distribute class.

odiazib
authored and
mahf708
committed
Moving logic to distribute class.

odiazib
authored and
mahf708
committed
Removing 'comm' from the model implementation. This change makes the training slower by 10 seconds for each epoch.

odiazib
authored and
mahf708
committed
Removing old code.

odiazib
authored and
mahf708
committed
Cleaning up code based on PR review.

odiazib
authored and
mahf708
committed
Build the NeMo version of SFNO if spatial parallelism is on.

odiazib
authored and
mahf708
committed
Adding review recommendations.

odiazib
authored and
mahf708
committed
Fixing unit test for xarray sp

odiazib
authored and
mahf708
committed
Fixing test for sfnonet with spatial parallelism.

odiazib
authored and
mahf708
committed
routine to create a directory.

odiazib
authored and
mahf708
committed
Saving test

odiazib
authored and
mahf708
committed
unit test for loss function.

odiazib
authored and
mahf708
committed

Commits on Nov 17, 2025

The ERA5 dataset uses latitude and longitude instead of lon and lat. Thus, we must add logic to handle this case. I also moved this part of the code that reshapes the dataset to the distribution cl…
odiazib
committed

Commits on Nov 18, 2025

Adding code back.
odiazib
committed

Commits on Nov 19, 2025

Commits on Dec 1, 2025

tmp cleanup/modularization

mahf708
authored and
odiazib
committed
Cleaning up.
odiazib
committed

Commits on Dec 10, 2025

save unit test.
odiazib
committed