achyuthan.s
5aef6d8833
## Summary `import deepspeed` initialized a CUDA context in the parent process, which permanently breaks `fork()`-based multiprocessing (`Cannot re-initialize CUDA in forked subprocess`). This makes importing DeepSpeed fork-safe. Fixes #7918. ## Root cause On a GPU box, `import deepspeed` reached **three** distinct calls that create a CUDA context, each gated differently (which is why a single patch kept missing one): 1. **`torch.cuda.is_available()`** — called during accelerator auto-detection (`real_accelerator.py`) and in every CUDA op builder's `is_compatible()`. By default it runs `cudaGetDeviceCount → cuInit`, creating a context. Per the [PyTorch docs](https://docs.pytorch.org/docs/stable/generated/torch.cuda.is_available.html) this is only avoided with `PYTORCH_NVML_BASED_CUDA_CHECK=1`. Note it does **not** set `torch.cuda.is_initialized()`, so an import-time `assert not is_initialized()` is a false-green. 2. **`torch.cuda.get_device_properties(0)`** — in the eight builders' `is_compatible()` (run at import by `git_version_info.py`); triggers `torch.cuda._lazy_init()`. 3. **`is_triton_supported()` → `torch.cuda.get_device_capability()`** — called at module import in `ds_transformer.py`, gated on `deepspeed.HAS_TRITON`. This only fires when **triton is installed**, so it was invisible in triton-less environments — but it was the first initializer on a real GPU node. ## Fix 1. `deepspeed/__init__.py` sets `os.environ.setdefault("PYTORCH_NVML_BASED_CUDA_CHECK", "1")` as the very first statement, so `torch.cuda.is_available()` uses the NVML-based check and never initializes a context. `setdefault()` preserves an explicit user setting. 2. `CUDAOpBuilder.cuda_capability_major()` (in `op_builder/builder.py`) reads compute capability only when a context already exists (`is_initialized()`) and we are not in a forked child (`_is_in_bad_fork()`, mirroring #7977); otherwise returns `None`. All eight builders route through it and skip the capability gate when probing is unsafe. 3. `ds_transformer.py` imports the triton kernels whenever triton is installed (`if deepspeed.HAS_TRITON:`) instead of also gating on `is_triton_supported()`. The capability probe is removed from import; actual triton use stays gated at runtime by `config.use_triton`, where CUDA is already initialized. ## Behavior / tradeoff - NVML-based availability is a slightly weaker assessment than the default runtime check and falls back to `cudaGetDeviceCount` if NVML is unavailable (documented PyTorch behavior); a non-issue on standard NVIDIA boxes. - Dropping the import-time capability gate means triton kernel modules are imported whenever triton is installed (even on pre-Ampere). Importing them has no CUDA side effects; their use is still gated by `config.use_triton`. ## Tests - Three unit tests for `cuda_capability_major()`'s decision tree (not-initialized → skip, initialized → probe, bad-fork → skip), mocked `torch.cuda`, no GPU required. - `test_forked_child_can_use_cuda_after_importing_deepspeed` — forks after `import deepspeed`, the child runs a real CUDA op, parent asserts success. ## Validation Verified on a CUDA GPU node (NVIDIA, torch 2.4.1+cu121). After `import deepspeed`: - `torch.cuda.is_initialized()` → `False` - a forked child runs `torch.ones(1, device="cuda")` successfully (exit 0) - instrumenting `torch.cuda._lazy_init` shows **0** distinct import-time CUDA-touch sites (down from the `ds_transformer.py:17` initializer + its downstream builder probe). ## Docs Updated `CONTRIBUTING.md` and `docs/contributing.md`: `--forked` is safe now that `import deepspeed` no longer initializes CUDA. cc @tjruwase @loadams @tohtana --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>
3.2 KiB
title, permalink
| title | permalink |
|---|---|
| Contributing | /contributing/ |
DeepSpeed welcomes your contributions!
Prerequisites
DeepSpeed uses pre-commit to ensure that formatting is
consistent across DeepSpeed. First, ensure that pre-commit is installed from either
installing DeepSpeed or pip install pre-commit. Next, the pre-commit hooks must be
installed once before commits can be made:
pre-commit install
Afterwards, our suite of formatting tests run automatically before each git commit. You
can also run these manually:
pre-commit run --all-files
If a formatting test fails, it will fix the modified code in place and abort
the git commit. After looking over the changes, you can git add <modified files>
and then repeat the previous git commit command.
Testing
DeepSpeed tracks two types of tests: unit tests and more costly model convergence tests.
The model convergence tests train
DeepSpeedExamples and measure
end-to-end convergence and related metrics. Unit tests are found in tests/unit/ and
the model convergence tests are found in tests/model/.
Unit Tests
PyTest is used to execute tests. PyTest can be
installed from PyPI via pip install pytest. Simply invoke pytest --forked to run the
unit tests:
pytest --forked tests/unit/
You can also provide the -v flag to pytest to see additional information about the
tests. Note that pytest-forked and the
--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 DeepSpeedExamples.
To execute model tests, first install DeepSpeed. The DeepSpeedExamples repository is cloned as part of this process. Next, execute the model test driver:
cd tests/model/
pytest run_sanity_check.py
Note that the --forked flag is not necessary for the model tests.
Contributor License Agreement
This project welcomes contributions and suggestions. Most contributions require you to agree to a Developer Certificate of Origin (DCO)[https://wiki.linuxfoundation.org/dco] stating that they agree to the terms published at https://developercertificate.org for that particular contribution.
DCOs are per-commit, so each commit needs to be signed off. These can be signed in
the commit by adding the -s flag. DCO enforcement can also be signed off in the PR
itself by clicking on the DCO enforcement check.
Code of Conduct
This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact opencode@microsoft.com with any additional questions or comments.