Sfno fix (NVIDIA#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

modulus/experimental/sfno/perf_tests/distributed/dist_fft3d.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"))
@@ -30,10 +31,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.fft3d import RealFFT3, InverseRealFFT3, DistributedRealFFT3, DistributedInverseRealFFT3
 
-# imprt patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 def main(args, verify):
     # parameters

modulus/experimental/sfno/perf_tests/distributed/dist_ifft.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

modulus/experimental/sfno/perf_tests/distributed/dist_ifft3d.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"))
@@ -30,10 +31,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.fft3d import RealFFT3, InverseRealFFT3, DistributedRealFFT3, DistributedInverseRealFFT3
 
-# imprt patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 def main(args, verify):
     # parameters

modulus/experimental/sfno/perf_tests/primitives/comp_mult.py

@@ -20,15 +20,13 @@
 from torch.cuda import amp
 import time
 import apex
+import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
 
 sys.path.append(os.path.join("/opt", "ERA5_wind"))
 
 from modulus.experimental.sfno.mpu.layers import compl_mul_add_fwd, compl_mul_add_fwd_c
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 class ComplexMult(nn.Module):
     def __init__(self, num_blocks, block_size, hidden_size_factor, use_complex_kernels=True):

modulus/experimental/sfno/perf_tests/sfno/shtfilter.py

@@ -21,6 +21,8 @@
 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", "makani"))
@@ -31,10 +33,6 @@
 from torch_harmonics import RealSHT as RealSphericalHarmonicTransform
 from torch_harmonics import InverseRealSHT as InverseRealSphericalHarmonicTransform
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 # profile stuff
 from ctypes import cdll
 libcudart = cdll.LoadLibrary('libcudart.so')

modulus/experimental/sfno/utils/comm.py

@@ -18,14 +18,11 @@
 from modulus.experimental.sfno.utils.logging_utils import disable_logging
 import math
 import torch
+import torch.distributed as dist
 import datetime as dt
 from typing import Union
 import numpy as np
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 # dummy placeholders
 _COMM_LIST = []
 _COMM_NAMES = {}

modulus/experimental/sfno/utils/dataloader.py

@@ -19,12 +19,9 @@
 from torch.utils.data import DataLoader
 
 # distributed stuff
+import torch.distributed as dist
 from modulus.experimental.sfno.utils import comm
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 def init_distributed_io(params):
     # set up sharding

modulus/experimental/sfno/utils/dataloaders/data_loader_dali_2d.py

@@ -23,6 +23,7 @@
 #import cv2
 
 # distributed stuff
+import torch.distributed as dist
 from modulus.experimental.sfno.utils import comm
 
 # DALI stuff
@@ -35,10 +36,6 @@
 import modulus.experimental.sfno.utils.dataloaders.dali_es_helper_2d as esh
 from modulus.experimental.sfno.utils.grids import GridConverter
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 
 class ERA5DaliESDataloader(object):
 

modulus/experimental/sfno/utils/distributed_patch.py

@@ -18,12 +18,9 @@
 # distributed computing stuff
 from modulus.experimental.sfno.utils import comm
 from modulus.experimental.sfno.utils.metrics.functions import GeometricL1, GeometricRMSE, GeometricACC, Quadrature
+import torch.distributed as dist
 from modulus.experimental.sfno.mpu.mappings import gather_from_parallel_region
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
-
 class MetricsHandler():
     """
     Handler object which takes care of computation of metrics. Keeps buffers for the computation of 

modulus/experimental/sfno/utils/metric.py

@@ -40,6 +40,7 @@
 # distributed computing stuff
 from modulus.experimental.sfno.utils import comm
 from modulus.experimental.sfno.utils import visualize
+import torch.distributed as dist
 
 # for the manipulation of state dict
 from collections import OrderedDict
@@ -51,9 +52,6 @@
 from modulus.experimental.sfno.third_party.torch.optim.adam import Adam as CustomAdam
 from modulus.experimental.sfno.third_party.torch.optim.adamw import AdamW as CustomAdamW
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 class Trainer():
     """

modulus/experimental/sfno/utils/trainer.py

@@ -39,6 +39,7 @@
 # distributed computing stuff
 from modulus.experimental.sfno.utils import comm
 from modulus.experimental.sfno.utils import visualize
+import torch.distributed as dist
 
 # for the manipulation of state dict
 from collections import OrderedDict
@@ -50,9 +51,6 @@
 from modulus.experimental.sfno.third_party.torch.optim.adam import Adam as CustomAdam
 from modulus.experimental.sfno.third_party.torch.optim.adamw import AdamW as CustomAdamW
 
-# import patched distributed
-from modulus.experimental.sfno.utils.distributed_patch import dist_patch
-dist = dist_patch()
 
 # profile stuff
 from ctypes import cdll