Avoid CUDA context initialization during op compatibility checks at import

CONTRIBUTING.md

@@ -42,7 +42,9 @@ pytest --forked tests/unit/
 ```
 You can also provide the `-v` flag to `pytest` to see additional information about the
 tests. Note that [pytest-forked](https://github.com/pytest-dev/pytest-forked) and the
-`--forked` flag are required to test CUDA functionality in distributed tests.
+`--forked` flag are required to test CUDA functionality in distributed tests. Using
+`--forked` is safe because `import deepspeed` no longer initializes a CUDA context;
+earlier versions probed CUDA at import time, which poisoned `fork()`.
 
 You can also run:
 ```

deepspeed/__init__.py

@@ -3,6 +3,16 @@
 
 # DeepSpeed Team
 
+import os
+
+# By default, PyTorch's CUDA availability check (cudaGetDeviceCount/cuInit)
+# creates a CUDA context, which poisons fork()-based multiprocessing once
+# DeepSpeed probes op compatibility at import time. Opt into PyTorch's
+# NVML-based availability check so importing DeepSpeed never creates a CUDA
+# context, before importing torch or anything that may query CUDA.
+# setdefault() preserves an explicit user setting. See issue #7918.
+os.environ.setdefault("PYTORCH_NVML_BASED_CUDA_CHECK", "1")
+
 import argparse
 import sys
 import types

deepspeed/model_implementations/transformers/ds_transformer.py

@@ -14,7 +14,11 @@
 from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp
 from deepspeed.accelerator import get_accelerator
 import deepspeed
-if deepspeed.HAS_TRITON and get_accelerator().is_triton_supported():
+# Import the triton kernels whenever triton is installed. Previously this was also
+# gated on is_triton_supported(), which reads the GPU compute capability at import
+# time and thereby creates a CUDA context, breaking fork()-based multiprocessing
+# (issue #7918). Triton use is gated at runtime via self.config.use_triton below.
+if deepspeed.HAS_TRITON:
     from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP
     from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention
 

docs/contributing.md

@@ -40,7 +40,9 @@ pytest --forked tests/unit/
 ```
 You can also provide the `-v` flag to `pytest` to see additional information about the
 tests. Note that [pytest-forked](https://github.com/pytest-dev/pytest-forked) and the
-`--forked` flag are required to test CUDA functionality in distributed tests.
+`--forked` flag are required to test CUDA functionality in distributed tests. Using
+`--forked` is safe because `import deepspeed` no longer initializes a CUDA context;
+earlier versions probed CUDA at import time, which poisoned `fork()`.
 
 ### Model Tests
 Model tests require four GPUs and training data downloaded for

op_builder/builder.py

@@ -614,6 +614,24 @@ def jit_load(self, verbose=True):
 
 class CUDAOpBuilder(OpBuilder):
 
+    def cuda_capability_major(self):
+        """Compute-capability major of CUDA device 0, or ``None`` when it cannot
+        be read without side effects.
+
+        ``torch.cuda.get_device_properties`` calls ``torch.cuda._lazy_init()``,
+        which creates a CUDA context. Doing that merely to check op compatibility
+        at ``import deepspeed`` time would poison ``fork()``-based multiprocessing,
+        because a forked child cannot reuse the parent's context (issue #7918).
+        We therefore probe only when a context already exists and we are not
+        inside such a forked child; otherwise the caller skips the
+        compute-capability check and defers it to build/load time.
+        """
+        if not torch.cuda.is_initialized():
+            return None
+        if hasattr(torch.cuda, '_is_in_bad_fork') and torch.cuda._is_in_bad_fork():
+            return None
+        return torch.cuda.get_device_properties(0).major
+
     def compute_capability_args(self, cross_compile_archs=None):
         """
         Returns nvcc compute capability compile flags.

op_builder/evoformer_attn.py

@@ -87,8 +87,8 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split(".")[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major  #ignore-cuda
-                if cuda_capability < 7:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 7:
                     if verbose:
                         self.warning("Please use a GPU with compute capability >= 7.0")
                     cuda_okay = False

op_builder/fp_quantizer.py

@@ -36,12 +36,12 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major  #ignore-cuda
-                if cuda_capability < 8:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 8:
                     if verbose:
                         self.warning("NVIDIA Inference is only supported on Ampere and newer architectures")
                     cuda_okay = False
-                if cuda_capability >= 8:
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

op_builder/inference_core_ops.py

@@ -32,12 +32,12 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major  #ignore-cuda
-                if cuda_capability < 6:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 6:
                     if verbose:
                         self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
                     cuda_okay = False
-                if cuda_capability >= 8:
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

op_builder/inference_cutlass_builder.py

@@ -31,12 +31,12 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major  #ignore-cuda
-                if cuda_capability < 6:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 6:
                     if verbose:
                         self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
                     cuda_okay = False
-                if cuda_capability >= 8:
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

op_builder/ragged_ops.py

@@ -32,12 +32,12 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major  #ignore-cuda
-                if cuda_capability < 6:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 6:
                     if verbose:
                         self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
                     cuda_okay = False
-                if cuda_capability >= 8:
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

op_builder/ragged_utils.py

@@ -32,12 +32,12 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():  #ignore-cuda
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major  #ignore-cuda
-                if cuda_capability < 6:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 6:
                     if verbose:
                         self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
                     cuda_okay = False
-                if cuda_capability >= 8:
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

op_builder/spatial_inference.py

@@ -31,8 +31,8 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major
-                if cuda_capability >= 8:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

op_builder/transformer_inference.py

@@ -31,12 +31,12 @@ def is_compatible(self, verbose=False):
             if not self.is_rocm_pytorch() and torch.cuda.is_available():
                 sys_cuda_major, _ = installed_cuda_version()
                 torch_cuda_major = int(torch.version.cuda.split('.')[0])
-                cuda_capability = torch.cuda.get_device_properties(0).major
-                if cuda_capability < 6:
+                cuda_capability = self.cuda_capability_major()
+                if cuda_capability is not None and cuda_capability < 6:
                     if verbose:
                         self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
                     cuda_okay = False
-                if cuda_capability >= 8:
+                if cuda_capability is not None and cuda_capability >= 8:
                     if torch_cuda_major < 11 or sys_cuda_major < 11:
                         if verbose:
                             self.warning("On Ampere and higher architectures please use CUDA 11+")

tests/unit/ops/test_op_builder.py

@@ -4,6 +4,8 @@
 # DeepSpeed Team
 
 import os
+import sys
+import subprocess
 import importlib.util
 from pathlib import Path
 from unittest.mock import MagicMock, patch
@@ -216,3 +218,83 @@ def test_non_jit_branch_canonical_dedupe_mixed_ptx_combinations():
             args = builder.compute_capability_args()
             assert os.environ["TORCH_CUDA_ARCH_LIST"] == expected_arch_list, arch_input
         assert args == expected_args, arch_input
+
+
+def test_cuda_capability_major_skips_probe_when_context_not_initialized():
+    # Probing device properties forces a lazy CUDA-context init, which creates a
+    # CUDA context. Doing that while checking op compatibility at "import deepspeed"
+    # time poisons fork()-based multiprocessing (issue #7918): a forked child cannot
+    # reuse the parent's context. With no context yet, the probe must be skipped.
+    builder = make_builder()
+    with patch.object(CUDA_API, "is_initialized", return_value=False):
+        with patch.object(
+                CUDA_API, "get_device_properties",
+                side_effect=AssertionError("must not initialize CUDA / poison fork")) as get_device_properties:
+            assert builder.cuda_capability_major() is None
+    get_device_properties.assert_not_called()
+
+
+def test_cuda_capability_major_probes_when_context_already_initialized():
+    # When a CUDA context already exists (e.g. at op load time), probing is safe
+    # and must report the real compute-capability major.
+    builder = make_builder()
+    device_properties = MagicMock(major=8)
+    with patch.object(CUDA_API, "is_initialized", return_value=True):
+        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=False):
+            with patch.object(CUDA_API, "get_device_properties",
+                              return_value=device_properties) as get_device_properties:
+                assert builder.cuda_capability_major() == 8
+    get_device_properties.assert_called_once_with(0)
+
+
+def test_cuda_capability_major_skips_probe_in_bad_fork():
+    # Inside a forked child that inherited an initialized-but-invalid context,
+    # probing would raise "Cannot re-initialize CUDA in forked subprocess", so it
+    # must be skipped there as well.
+    builder = make_builder()
+    with patch.object(CUDA_API, "is_initialized", return_value=True):
+        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=True):
+            with patch.object(CUDA_API,
+                              "get_device_properties",
+                              side_effect=AssertionError("must not probe in a forked child")) as get_device_properties:
+                assert builder.cuda_capability_major() is None
+    get_device_properties.assert_not_called()
+
+
+def test_forked_child_can_use_cuda_after_importing_deepspeed():
+    # Core contract of issue #7918: after the parent process runs
+    # ``import deepspeed``, a forked child must still be able to initialize and
+    # use CUDA. If import created a CUDA context in the parent, the child fails
+    # with "Cannot re-initialize CUDA in forked subprocess". Everything runs in a
+    # dedicated subprocess so a poisoned parent cannot leak into the pytest worker
+    # or other tests.
+    program = "\n".join([
+        "import os, sys",
+        "import torch",
+        "import deepspeed  # must not create a CUDA context in the parent",
+        # device_count() is NVML-based and never initializes a context, so it is
+        # a fork-safe way to check for a GPU before forking.
+        "if torch.cuda.device_count() == 0:",  #ignore-cuda
+        "    print('NO_CUDA'); sys.exit(0)",
+        "pid = os.fork()",
+        "if pid == 0:",
+        "    try:",
+        "        torch.ones(1, device='cuda')",
+        "        os._exit(0)",
+        "    except Exception as exc:",
+        "        sys.stderr.write(repr(exc))",
+        "        os._exit(1)",
+        "_, status = os.waitpid(pid, 0)",
+        "sys.exit(os.waitstatus_to_exitcode(status))",
+    ])
+    env = os.environ.copy()
+    repo_root = str(Path(__file__).resolve().parents[3])
+    env["PYTHONPATH"] = repo_root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
+    result = subprocess.run([sys.executable, "-c", program], capture_output=True, text=True, env=env, timeout=300)
+    if result.returncode != 0 and ("No module named 'deepspeed'" in result.stderr
+                                   or "No module named 'torch'" in result.stderr):
+        pytest.skip("deepspeed/torch not importable in a subprocess in this environment")
+    if result.stdout.strip() == "NO_CUDA":
+        pytest.skip("no CUDA device available")
+    assert result.returncode == 0, ("forked child could not use CUDA after 'import deepspeed' "
+                                    "(a CUDA context was created during import, issue #7918):\n" + result.stderr)