Add custom build backend to support build args

This implements a custom build backend, inspired by [the solution used
by Pillow](python-pillow/Pillow#7171).

With this change, argument flags now need to be specified as key-value
pairs instead of key-only as `--config-settings`. Every argument can
either be left unspecified or have a value that is either "enabled" or
"disabled".

For example:

```
--config-settings --global-option=--cpp_ext
--global-option --cpp_ext
```

becomes

```
--config-settings cpp_ext=enabled
```

The setuptools version requirement of 40.8.0 was also adopted from
Pillow. It is the version that introduced the `build_meta:__legacy__`
backend.

README.md

@@ -129,10 +129,7 @@ CUDA and C++ extensions via
 ```bash
 git clone https://github.com/NVIDIA/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 "--global-option=--cpp_ext" --config-settings "--global-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" ./
+pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--cpp_ext=enable" --config-settings "--cuda_ext=enable" ./
 ```
 
 APEX also supports a Python-only build via
@@ -148,7 +145,7 @@ A Python-only build omits:
 
 
 ### [Experimental] Windows
-`pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--global-option=--cpp_ext" --config-settings "--global-option=--cuda_ext" .` may work if you were able to build Pytorch from source
+`pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--cpp_ext=enable" --config-settings "--cuda_ext=enable" .` 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.  
 If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.
 

_custom_build/backend.py

@@ -0,0 +1,66 @@
+# Inspired by https://github.com/python-pillow/Pillow/pull/7171.
+
+import sys
+
+from setuptools.build_meta import *  # noqa: F401, F403
+
+backend_class = build_wheel.__self__.__class__
+
+
+class _CustomBuildMetaBackend(backend_class):
+    def run_setup(self, setup_script="setup.py"):
+        if self.config_settings:
+
+            def config_has(key, value):
+                settings = self.config_settings.get(key)
+                if settings:
+                    if not isinstance(settings, list):
+                        settings = [settings]
+                    return value in settings
+
+            flags = []
+            for ext in (
+                "cpp_ext",
+                "cuda_ext",
+                "distributed_adam",
+                "distributed_lamb",
+                "permutation_search",
+                "bnp",
+                "xentropy",
+                "focal_loss",
+                "index_mul_2d",
+                "deprecated_fused_adam",
+                "deprecated_fused_lamb",
+                "fast_layer_norm",
+                "fmha",
+                "fast_multihead_attn",
+                "transducer",
+                "cudnn_gbn",
+                "peer_memory",
+                "nccl_p2p",
+                "fast_bottleneck",
+                "fused_conv_bias_relu",
+            ):
+                if ext not in self.config_settings:
+                    continue
+
+                if config_has(ext, "enable"):
+                    flags.append("--" + ext)
+                elif not config_has(ext, "disable"):
+                    raise ValueError(
+                        f'unknown argument value for {ext}; must be either '
+                        f'"enable" or "disable"'
+                    )
+
+            if flags:
+                sys.argv = sys.argv[:1] + ["build_ext"] + flags + sys.argv[1:]
+        return super().run_setup(setup_script)
+
+    def build_wheel(
+        self, wheel_directory, config_settings=None, metadata_directory=None
+    ):
+        self.config_settings = config_settings
+        return super().build_wheel(wheel_directory, config_settings, metadata_directory)
+
+
+build_wheel = _CustomBuildMetaBackend().build_wheel

pyproject.toml

@@ -1,6 +1,7 @@
 [build-system]
 requires = [
-    "setuptools",
+    "setuptools >= 40.8.0",
     "wheel",
 ]
-build-backend = "setuptools.build_meta"
+build-backend = "backend"  # custom setuptools.build_meta
+backend-path = ["_custom_build"]