update

setup.py

@@ -177,16 +177,15 @@ def run(self):
         'xlrd>=2.0.1',
         'xlsxwriter>=1.3.8',
         'xmltodict>=0.12.0',
-        'types-requests',
-        'mpi4py>=3.1.0',
+        'types-requests'
     ],
     extras_require=(
         lambda x: {
             **x,
             'develop': x['dev'] + x['test'],
             'cpuworker': x['torch'],
-            'amdworker': x['torch'] + x['ort'] + x['amd'],
-            'nvworker': x['torch'] + x['ort'] + x['nvidia'],
+            'amdworker': x['torch'] + x['ort'] + x['amd'] + x['mpi'],
+            'nvworker': x['torch'] + x['ort'] + x['nvidia'] + x['mpi'],
         }
     )(
         {
@@ -220,6 +219,7 @@ def run(self):
             ],
             'nvidia': ['py3nvml>=0.2.6'],
             'amd': ['pyrsmi>=1.0.1'],
+            'mpi': ['mpi4py>=3.1.0'],
         }
     ),
     include_package_data=True,

superbench/benchmarks/model_benchmarks/megatron_gpt3.py

@@ -12,8 +12,6 @@
 from pathlib import Path
 import re
 
-from mpi4py import MPI
-
 from superbench.benchmarks import BenchmarkRegistry
 from superbench.benchmarks.context import Platform, Precision
 from superbench.benchmarks.model_benchmarks.model_base import ModelBenchmark
@@ -343,6 +341,7 @@ def _sync_result(self, data):
         Args:
             data (list): the data to be reduced.
         """
+        from mpi4py import MPI
         comm = MPI.COMM_WORLD
         data = np.array(data, dtype=np.float64)
         # Reduce the data to a single value on rank 0