[REDUX] Refactor Apex build process to use the PyTorch JIT extension flow

.gitignore

@@ -148,3 +148,7 @@ cython_debug/
 *.hip
 *_hip.*
 *hip* 
+
+
+#file temporarily created for build process
+apex/git_version_info_installed.py

MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include apex/contrib/csrc *
+recursive-include apex/csrc *

Makefile

@@ -0,0 +1,17 @@
+PYTHON = python3
+PIP = $(PYTHON) -m pip
+
+clean: # This will remove ALL build folders.
+	@test -d build/ && echo "Deleting build folder" || true
+	@test -d build/ && rm -r build/ || true
+	@test -d dist/ && echo "Deleting dist folder" || true
+	@test -d dist/ && rm -r dist/ || true
+	@test -d apex.egg-info/ && echo "Deleting apex.egg-info folder" || true
+	@test -d apex.egg-info/ && rm -r apex.egg-info/ || true
+
+	$(PYTHON) scripts/clean.py # remove the apex extensions installed at torch extensions folder 
+
+aiter:
+	$(PIP) uninstall -y aiter
+	cd third_party/aiter && $(PIP) install . --no-build-isolation --no-deps
+

README.md

@@ -100,24 +100,21 @@ Note that we recommend restoring the model using the same `opt_level`. Also note
 # Installation
 
 ## Containers
-ROCm pytorch containers are available from https://hub.docker.com/r/rocm/pytorch.
+ROCm pytorch containers contain apex package and these are available from https://hub.docker.com/r/rocm/pytorch.
 
 ## From Source
 
-To install Apex from source, we recommend using the nightly Pytorch obtainable from https://github.com/rocm/pytorch.
+Torch must be installed before installing apex. We recommend using the nightly Pytorch obtainable from https://github.com/rocm/pytorch. The latest stable release obtainable from https://pytorch.org should also work.
 
-The latest stable release obtainable from https://pytorch.org should also work.
-
-## ROCm
 Apex on ROCm supports both python only build and extension build.
 Note: Pytorch version recommended is >=1.5 for extension build.
 
-### To install using python only build use the following command in apex folder:
+### The following command will install all the extensions, which will be built and linked at runtime using [PyTorch's JIT (just-in-time) loader](https://pytorch.org/docs/stable/cpp_extension.html):
+This requires ninja to be installed
 ```
-python setup.py install
+pip install . --no-build-isolation
 ```
 
-=======
 ### Supported Versions
 | ``APEX Version`` | ``APEX branch`` | ``Torch Version`` |
 |------------------|-----------------|-------------------|
@@ -140,26 +137,73 @@ ubuntu|pytorch|apex|release/1.0.0|06c33eee43f7a22f3ed7d9c3e5be0ddd757dc345|https
 centos|pytorch|apex|release/1.0.0|06c33eee43f7a22f3ed7d9c3e5be0ddd757dc345|https://github.com/ROCmSoftwarePlatform/apex
 ```
 
-### To install using extensions enabled use the following command in apex folder:
+### To pre-build and install all the supported extensions while installing apex, use the following command in apex folder:
+```
+APEX_BUILD_CPP_OPS=1 APEX_BUILD_CUDA_OPS=1 pip install . --no-build-isolation
 ```
-# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...
-pip install -v --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
-# otherwise
-python setup.py install --cpp_ext --cuda_ext
 
+It is also possible to pre-build and install specific extensions by using the following command in apex folder:
+```
+APEX_BUILD_<OP_NAME>=1 pip install . --no-build-isolation
 ```
-Note that using --cuda_ext flag to install Apex will also enable all the extensions supported on ROCm including "--distributed_adam", "--distributed_lamb", "--bnp", "--xentropy", "--deprecated_fused_adam", "--deprecated_fused_lamb", and "--fast_multihead_attn".
+The following extensions are supported:
+| extension | environment to build specific extension | install option |
+|-----------|-----------|-----------|
+| amp_C | APEX_BUILD_AMP_C=1 | APEX_BUILD_CUDA_OPS=1 |
+| apex_C | APEX_BUILD_APEX_C=1 | APEX_BUILD_CPP_OPS=1 |
+| bnp | APEX_BUILD_BNP=1 | APEX_BUILD_CUDA_OPS=1 |
+| distributed_adam_cuda | APEX_BUILD_DISTRIBUTED_ADAM=1 | APEX_BUILD_CUDA_OPS=1 |
+| distributed_lamb_cuda | APEX_BUILD_DISTRIBUTED_LAMB=1 | APEX_BUILD_CUDA_OPS=1 |
+| fast_multihead_attn | APEX_BUILD_FAST_MULTIHEAD_ATTN=1 | APEX_BUILD_CUDA_OPS=1 |
+| focal_loss_cuda | APEX_BUILD_FOCAL_LOSS=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_adam_cuda | APEX_BUILD_FUSED_ADAM=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_bias_swiglu | APEX_BUILD_FUSED_BIAS_SWIGLU=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_dense_cuda | APEX_BUILD_FUSED_DENSE=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_index_mul_2d | APEX_BUILD_FUSED_INDEX_MUL_2D=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_lamb_cuda | APEX_BUILD_FUSED_LAMB=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_layer_norm_cuda | APEX_BUILD_FUSED_LAYER_NORM=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_rotary_positional_embedding | APEX_BUILD_FUSED_ROPE=1 | APEX_BUILD_CUDA_OPS=1 |
+| fused_weight_gradient_mlp_cuda | APEX_BUILD_FUSED_WEIGHT_GRADIENT_MLP=1 | APEX_BUILD_CUDA_OPS=1 |
+| generic_scaled_masked_softmax_cuda | APEX_BUILD_GENERIC_SCALED_MASKED_SOFTMAX_CUDA=1 | APEX_BUILD_CUDA_OPS=1 |
+| mlp_cuda | APEX_BUILD_MLP=1 | APEX_BUILD_CUDA_OPS=1 |
+| _apex_nccl_allocator | APEX_BUILD_NCCL_ALLOCATOR=1 | APEX_BUILD_CUDA_OPS=1 |
+| nccl_p2p_cuda | APEX_BUILD_NCCL_P2P=1 | APEX_BUILD_CUDA_OPS=1 |
+| peer_memory_cuda | APEX_BUILD_PEER_MEMORY=1 | APEX_BUILD_CUDA_OPS=1 |
+| scaled_masked_softmax_cuda | APEX_BUILD_SCALED_MASKED_SOFTMAX_CUDA=1 | APEX_BUILD_CUDA_OPS=1 |
+| scaled_softmax_cuda | APEX_BUILD_SCALED_SOFTMAX_CUDA=1 | APEX_BUILD_CUDA_OPS=1 |
+| scaled_upper_triang_masked_softmax_cuda | APEX_BUILD_SCALED_UPPER_TRIANG_MASKED_SOFTMAX_CUDA=1 | APEX_BUILD_CUDA_OPS=1 |
+| syncbn | APEX_BUILD_SYNCBN=1 | APEX_BUILD_CUDA_OPS=1 |
+| transducer_joint_cuda | APEX_BUILD_TRANSDUCER_JOINT=1 | APEX_BUILD_CUDA_OPS=1 |
+| transducer_loss_cuda | APEX_BUILD_TRANSDUCER_LOSS=1 | APEX_BUILD_CUDA_OPS=1 |
+| xentropy_cuda | APEX_BUILD_XENTROPY=1 | APEX_BUILD_CUDA_OPS=1 |
+
+For example, to build FUSED_DENSE​ you can use the following command: 
+```
+APEX_BUILD_FUSED_DENSE​=1 pip install . --no-build-isolation
+```
+This will pre-build and install FUSED_DENSE​ module and rest of the modules are installed to be JIT built and loaded at runtime. 
+
 
-In addition, aiter backend can be built during apex installation by providing --aiter flag
+
+Aiter backend can be built and used for fused rope. To install aiter:
 ```
-# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...
-pip install -v --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" --config-settings "--build-option=--aiter" ./
-# otherwise
-python setup.py install --cpp_ext --cuda_ext --aiter
+make aiter
 ```
 
 To use aiter in fused rope, you can use the flag ```USE_ROCM_AITER_ROPE_BACKEND=1```.
 
+### To create a wheel and then install apex using the wheel, use the following command in apex folder:
+```
+python -m build --wheel --no-isolation (can use the same environment variables to build specific extensions, cpp extensions and cuda extensions)
+pip install dist/apex-*.whl​
+```
+
+### To uninstall apex and its extensions, use the following command in apex folder:
+```
+pip uninstall apex
+make clean
+```
+
 ### Enable hipblasLT on ROCm
 hipblasLT is supported only on mi300 (gfx942) only.  
 python setup.py automatically builds apex with hipblasLT support only if GPU device id is gfx942  
@@ -173,33 +217,22 @@ CUDA and C++ extensions via
 ```bash
 git clone https://github.com/rocm/apex
 cd apex
-# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...
-pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
-# otherwise
-pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
-```
-
-Apex also supports a Python-only build via
-```bash
-pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./
+pip install . --no-build-isolation
 ```
-A Python-only build omits:
-- Fused kernels required to use `apex.optimizers.FusedAdam`.
-- Fused kernels required to use `apex.normalization.FusedLayerNorm` and `apex.normalization.FusedRMSNorm`.
-- Fused kernels that improve the performance and numerical stability of `apex.parallel.SyncBatchNorm`.
-- Fused kernels that improve the performance of `apex.parallel.DistributedDataParallel` and `apex.amp`.
-`DistributedDataParallel`, `amp`, and `SyncBatchNorm` will still be usable, but they may be slower.
 
 
 ### [Experimental] Windows
-`pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .` may work if you were able to build Pytorch from source
-on your system. A Python-only build via `pip install -v --no-cache-dir .` is more likely to work.  
+`pip install . --no-build-isolation` may work if you were able to build Pytorch from source
+on your system. A Python-only build via `pip install --no-build-isolation -v --no-cache-dir .` is more likely to work.  
 If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.
 
-
 # Release notes
 
-# Release notes
+## release/1.10.0
+
+Build and installation related
+- Support JIT (just-in-time) load cpp and CUDA extensions
+
 ## release/1.9.0
 
 - No new features were added in this release cycle.

apex/contrib/test/run_rocm_extensions.py

@@ -2,25 +2,27 @@
 import sys
 
 
-test_dirs = ["groupbn", "fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "."] # "." for test_label_smoothing.py
+test_dirs = ["groupbn", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", ".", \
+              "optimizers", "clip_grad"] # "." for test_label_smoothing.py
 ROCM_BLACKLIST = [
     "layer_norm"
 ]
 
-runner = unittest.TextTestRunner(verbosity=2)
+if __name__ ==  '__main__':
+    runner = unittest.TextTestRunner(verbosity=2)
 
-errcode = 0
+    errcode = 0
 
-for test_dir in test_dirs:
-    if test_dir in ROCM_BLACKLIST:
-        continue
-    suite = unittest.TestLoader().discover(test_dir)
+    for test_dir in test_dirs:
+        if test_dir in ROCM_BLACKLIST:
+            continue
+        suite = unittest.TestLoader().discover(test_dir)
 
-    print("\nExecuting tests from " + test_dir)
+        print("\nExecuting tests from " + test_dir)
 
-    result = runner.run(suite)
+        result = runner.run(suite)
 
-    if not result.wasSuccessful():
-        errcode = 1
+        if not result.wasSuccessful():
+            errcode = 1
 
-sys.exit(errcode)
+    sys.exit(errcode)

apex/csrc

@@ -0,0 +1 @@
+../csrc

apex/git_version_info.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Portions of this code were adapted from DeepSpeed:
+# https://github.com/microsoft/DeepSpeed
+# Modified for ROCm Apex
+
+try:
+    #  This is populated by setup.py
+    from .git_version_info_installed import *  # noqa: F401 # type: ignore
+except ModuleNotFoundError:
+    import os
+    if os.path.isfile('version.txt'):
+        # Will be missing from checkouts that haven't been installed (e.g., readthedocs)
+        version = open('version.txt', 'r').read().strip()
+    else:
+        version = "0.0.0"
+    git_hash = '[none]'
+    git_branch = '[none]'
+
+    from .op_builder.all_ops import ALL_OPS
+    installed_ops = dict.fromkeys(ALL_OPS.keys(), False)
+    torch_info = {'version': "0.0", "cuda_version": "0.0", "hip_version": "0.0"}
+
+# compatible_ops list is recreated for each launch
+from .op_builder.all_ops import ALL_OPS
+
+compatible_ops = dict.fromkeys(ALL_OPS.keys(), False)
+for op_name, builder in ALL_OPS.items():
+    op_compatible = builder.is_compatible()
+    compatible_ops[op_name] = op_compatible
+    compatible_ops["apex_not_implemented"] = False

apex/op_builder

@@ -0,0 +1 @@
+../op_builder

compatibility/_apex_nccl_allocator.py

@@ -0,0 +1,37 @@
+import sys
+import importlib
+
+class _ApexNcclAllocatorModule:
+    def __init__(self):
+        self._loaded_module = None
+        self._loading = False
+
+    def _load_module(self):
+        if self._loaded_module is None and not self._loading:
+            self._loading = True
+            try:
+                apex_op_builder = importlib.import_module('apex.op_builder')
+                builder = getattr(apex_op_builder, 'NCCLAllocatorBuilder')
+                self._loaded_module = builder().load()
+            except Exception as e:
+                self._loading = False
+                raise ImportError(f"Failed to load _apex_nccl_allocator : {e}")
+            finally:
+                self._loading = False
+        return self._loaded_module
+
+    def __getattr__(self, name):
+        if name.startswith("_") and name != "__class__":
+            raise AttributeError(f"module _apex_nccl_allocator has no attribute '{name}'")
+        return getattr(self._load_module(), name)
+
+    def __dir__(self):
+        try:
+            return dir(self._load_module())
+        except:
+            return []
+
+    def __repr__(self):
+        return "<module '_apex_nccl_allocator'>"
+
+sys.modules[__name__] = _ApexNcclAllocatorModule()

compatibility/amp_C.py

@@ -0,0 +1,37 @@
+import sys
+import importlib
+
+class _AmpCModule:
+    def __init__(self):
+        self._loaded_module = None
+        self._loading = False
+
+    def _load_module(self):
+        if self._loaded_module is None and not self._loading:
+            self._loading = True
+            try:
+                apex_op_builder = importlib.import_module('apex.op_builder')
+                builder = getattr(apex_op_builder, 'AmpCBuilder')
+                self._loaded_module = builder().load()
+            except Exception as e:
+                self._loading = False
+                raise ImportError(f"Failed to load amp_C : {e}")
+            finally:
+                self._loading = False
+        return self._loaded_module
+
+    def __getattr__(self, name):
+        if name.startswith("_"):
+            raise AttributeError(f"module amp_C has no attribute '{name}'")
+        return getattr(self._load_module(), name)
+
+    def __dir__(self):
+        try:
+            return dir(self._load_module())
+        except:
+            return []
+
+    def __repr__(self):
+        return "<module 'amp_C'>"
+
+sys.modules[__name__] = _AmpCModule()

compatibility/apex_C.py

@@ -0,0 +1,37 @@
+import sys
+import importlib
+
+class _ApexCModule:
+    def __init__(self):
+        self._loaded_module = None
+        self._loading = False
+
+    def _load_module(self):
+        if self._loaded_module is None and not self._loading:
+            self._loading = True
+            try:
+                apex_op_builder = importlib.import_module('apex.op_builder')
+                builder = getattr(apex_op_builder, 'ApexCBuilder')
+                self._loaded_module = builder().load()
+            except Exception as e:
+                self._loading = False
+                raise ImportError(f"Failed to load apex_C : {e}")
+            finally:
+                self._loading = False
+        return self._loaded_module
+
+    def __getattr__(self, name):
+        if name.startswith("_"):
+            raise AttributeError(f"module apex_C has no attribute '{name}'")
+        return getattr(self._load_module(), name)
+
+    def __dir__(self):
+        try:
+            return dir(self._load_module())
+        except:
+            return []
+
+    def __repr__(self):
+        return "<module 'apex_C'>"
+
+sys.modules[__name__] = _ApexCModule()