Sfno fix (#239)

* Add warning if jsbeutifier not installed, set default for h5 in inference, fix import

* copy pytorch patches instead of using monkeypatching

* Update README.md to include patching doc

---------

Co-authored-by: Mohammad Amin Nabian <[email protected]>

Author: daviddpruitt

modulus/experimental/sfno/README.md

@@ -11,6 +11,8 @@ This is a research code built for massively parallel training of SFNO for weathe
 
 ## Getting started
 
+**For distributed training or inference, run `patch_pytorch.sh` in advance. This will patch the pytorch distributed utilities to support complex values.**
+
 ## Installing optional dependencies
 
 Install the optional dependencies by running

modulus/experimental/sfno/convert_legacy_to_flexible.py

@@ -28,6 +28,7 @@
 )
 from modulus.experimental.sfno.utils import logging_utils
 
+import torch.distributed as dist
 
 from modulus.experimental.sfno.networks.models import get_model
 
@@ -36,10 +37,6 @@
 from modulus.experimental.sfno.utils.trainer import Trainer
 from modulus.experimental.sfno.utils.YParams import ParamsBase
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 class CheckpointSaver(Trainer):
     """

modulus/experimental/sfno/inference/inferencer.py

@@ -36,10 +36,8 @@
 # distributed computing stuff
 from modulus.experimental.sfno.utils import comm
 from modulus.experimental.sfno.utils import visualize
+import torch.distributed as dist
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 class Inferencer(Trainer):
     """

modulus/experimental/sfno/mpu/helpers.py

@@ -14,14 +14,12 @@
 
 import torch
 import torch.nn.functional as F
+import torch.distributed as dist
 
 from modulus.experimental.sfno.utils import comm
 
 from torch._utils import _flatten_dense_tensors
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 def get_memory_format(tensor):
     if tensor.is_contiguous(memory_format=torch.channels_last):

modulus/experimental/sfno/mpu/layers.py

@@ -16,6 +16,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.distributed as dist
 from torch.cuda.amp import custom_fwd, custom_bwd
 from modulus.experimental.sfno.utils import comm
 
@@ -28,10 +29,6 @@
 from modulus.experimental.sfno.mpu.helpers import pad_helper
 from modulus.experimental.sfno.mpu.helpers import truncate_helper 
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 class distributed_transpose_w(torch.autograd.Function):
 

modulus/experimental/sfno/mpu/mappings.py

@@ -18,6 +18,7 @@
 import torch
 from torch.nn.parallel import DistributedDataParallel  
 from modulus.experimental.sfno.utils import comm
+import torch.distributed as dist
 
 # torch utils
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
@@ -28,9 +29,6 @@
 from modulus.experimental.sfno.mpu.helpers import _split
 from modulus.experimental.sfno.mpu.helpers import _gather
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 # generalized
 class _CopyToParallelRegion(torch.autograd.Function):

modulus/experimental/sfno/networks/helpers.py

@@ -15,10 +15,7 @@
 import torch
 
 from utils import comm
-
-# imprt patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
+import torch.distributed as dist
 
 def count_parameters(model, device):
     with torch.no_grad():

modulus/experimental/sfno/patch_pytorch.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+cp third_party/torch/distributed/utils.py /usr/local/lib/python3.10/dist-packages/torch/distributed/
+echo "Patching complete"

modulus/experimental/sfno/perf_tests/distributed/comm_test.py

@@ -20,15 +20,13 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.distributed as dist
 from torch.cuda import amp
 
 sys.path.append(os.path.join("/opt", "ERA5_wind"))
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from modulus.experimental.sfno.utils import comm
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 # profile stuff
 from ctypes import cdll

modulus/experimental/sfno/perf_tests/distributed/dist_fft.py

@@ -20,6 +20,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.distributed as dist
 from torch.cuda import amp
 
 sys.path.append(os.path.join("/opt", "ERA5_wind"))
@@ -31,10 +32,6 @@
 from modulus.experimental.sfno.mpu.mappings import gather_from_parallel_region, scatter_to_parallel_region, reduce_from_parallel_region
 from modulus.experimental.sfno.mpu.layers import DistributedRealFFT2, DistributedInverseRealFFT2
 
-# imprt patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 def main(args, verify):
     # parameters