From ae21b3713de84d897b42a22068de0281fd89d540 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Wed, 22 Apr 2026 11:10:56 +0000 Subject: [PATCH] python3Packages.apex: init at 25.09 --- .../python-modules/apex/default.nix | 187 ++++++++++++++++ .../apex/fix-cudnn-frontend-compat.patch | 201 ++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 3 files changed, 390 insertions(+) create mode 100644 pkgs/development/python-modules/apex/default.nix create mode 100644 pkgs/development/python-modules/apex/fix-cudnn-frontend-compat.patch diff --git a/pkgs/development/python-modules/apex/default.nix b/pkgs/development/python-modules/apex/default.nix new file mode 100644 index 000000000000..6524afed8197 --- /dev/null +++ b/pkgs/development/python-modules/apex/default.nix @@ -0,0 +1,187 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + torch, + + # buildInputs + pybind11, + + # nativeBuildInputs + writableTmpDirAsHomeHook, + + # dependencies + cxxfilt, + numpy, + packaging, + pytest, + pyyaml, + tqdm, + + # tests + onnxscript, + pytestCheckHook, + torchvision, + + apex, + + cudaPackages, + cudaSupport ? torch.cudaSupport, +}: + +buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { + pname = "apex"; + version = "25.09"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "nvidia"; + repo = "apex"; + tag = finalAttrs.version; + hash = "sha256-/WcFCDjNXWbCnWoprYYAUcLt9p1CqJLzPXcBkPn+ics="; + }; + + patches = [ + # Fix incompatibility with more recent versions of cudnn to de-vendor it: + # error: ‘throw_if’ is not a member of ‘cudnn_frontend’ + ./fix-cudnn-frontend-compat.patch + ]; + + # Don't use git submodules for cuda dependencies + postPatch = '' + substituteInPlace setup.py \ + --replace-fail \ + 'subprocess.run(["git", "submodule", "update", "--init", "apex/contrib/csrc/multihead_attn/cutlass"])' \ + "" \ + --replace-fail \ + 'subprocess.run(["git", "submodule", "update", "--init", "apex/contrib/csrc/cudnn-frontend/"])' \ + "" + ''; + + env = { + APEX_CPP_EXT = 1; + } + // lib.optionalAttrs cudaSupport { + CUDA_HOME = (lib.getBin cudaPackages.cuda_nvcc).outPath; + TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}"; + + # Even if APEX_ALL_CONTRIB_EXT is enabled, APEX_CUDA_EXT must be explicitly enable + APEX_CUDA_EXT = 1; + + # Enable all contrib extensions at once + # https://github.com/NVIDIA/apex/tree/25.09#custom-ccuda-extensions-and-install-options + APEX_ALL_CONTRIB_EXT = 1; + + NVCC_APPEND_FLAGS = lib.toString [ + # Make kernel compilation slightly more parallel + "--threads 2" + ]; + }; + + preBuild = '' + export APEX_PARALLEL_BUILD=$NIX_BUILD_CORES + ''; + + build-system = [ + setuptools + torch + ]; + + buildInputs = [ + pybind11 + ] + ++ lib.optionals cudaSupport ( + with cudaPackages; + [ + cuda_cudart # cuda_runtime.h + cuda_profiler_api # cuda_profiler_api.h + cudnn # cudnn.h + cudnn-frontend # cudnn_frontend.h + cutlass # cutlass/cutlass.h + libcublas # cublas_v2.h + libcufile # cufile.h + libcurand # curand_kernel.h + libcusolver # cusolverDn.h + libcusparse # cusparse.h + nccl # nccl.h + ] + ); + + nativeBuildInputs = [ + writableTmpDirAsHomeHook + ]; + + dependencies = [ + cxxfilt + numpy + packaging + pytest + pyyaml + tqdm + ]; + + pythonImportsCheck = [ + "apex" + "apex_C" + ] + ++ lib.optionals cudaSupport [ + "_apex_gpu_direct_storage" + "_apex_nccl_allocator" + "amp_C" + "apex_C" + "bnp" + "fmhalib" + "fused_layer_norm_cuda" + "nccl_p2p_cuda" + "syncbn" + ]; + + nativeCheckInputs = [ + onnxscript + pytestCheckHook + torchvision + ]; + preCheck = '' + rm -rf apex + '' + # Otherwise, test collection fails with: + # ModuleNotFoundError: No module named 'test_fused_optimizer' + + '' + rm tests/L0/run_optimizers/__init__.py + ''; + doCheck = false; + disabledTestPaths = [ + # Try to read the driver version from nvidia-smi (failing in the sandbox) + # TypeError: expected string or bytes-like object, got 'NoneType' + "tests/L0/run_transformer/" + + # apex.parallel was removed in https://github.com/NVIDIA/apex/pull/1896, but some tests still + # try to import it + "tests/distributed/DDP/ddp_race_condition_test.py" + "tests/distributed/synced_batchnorm/" + ]; + + disabledTests = [ + # RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. + # torch.onnx._internal.exporter._errors.TorchExportError: Failed to export the model with torch.export. + "test_layer_norm_export_cuda" + "test_rms_export_cuda" + ]; + + passthru.gpuCheck = apex.overridePythonAttrs { + requiredSystemFeatures = [ "cuda" ]; + doCheck = true; + }; + + meta = { + description = "Tools for easy mixed precision and distributed training in Pytorch"; + homepage = "https://github.com/nvidia/apex"; + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ GaetanLepage ]; + broken = !cudaSupport; + }; +}) diff --git a/pkgs/development/python-modules/apex/fix-cudnn-frontend-compat.patch b/pkgs/development/python-modules/apex/fix-cudnn-frontend-compat.patch new file mode 100644 index 000000000000..a9b6fdb050ee --- /dev/null +++ b/pkgs/development/python-modules/apex/fix-cudnn-frontend-compat.patch @@ -0,0 +1,201 @@ +diff --git a/apex/contrib/csrc/bottleneck/bottleneck.cpp b/apex/contrib/csrc/bottleneck/bottleneck.cpp +index e80607c..b85e3fe 100644 +--- a/apex/contrib/csrc/bottleneck/bottleneck.cpp ++++ b/apex/contrib/csrc/bottleneck/bottleneck.cpp +@@ -621,7 +621,7 @@ run_conv_scale_bias_add_activation(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -749,7 +749,7 @@ run_conv_scale_bias(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -877,7 +877,7 @@ run_dconv_drelu_dscale(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -983,7 +983,7 @@ run_dconv(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -1093,7 +1093,7 @@ run_dconv_add(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -2242,7 +2242,7 @@ run_conv_add_scale_bias_activation(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -2497,7 +2497,7 @@ run_conv_scale_bias_add_activation_mask(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } else { + + std::array ops = {&conv_op, &scale_op, &bias_op, &act_op, &genIndex_op, &lessThan_op, &greaterThan_op, &logicalOr_op, &selection_op}; +@@ -2532,7 +2532,7 @@ run_conv_scale_bias_add_activation_mask(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; +@@ -2679,7 +2679,7 @@ run_dconv_add_drelu_dscale(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -2898,7 +2898,7 @@ run_dconv_drelu_dscale_mask(int64_t* x_dim_padded, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +diff --git a/apex/contrib/csrc/conv_bias_relu/conv_bias_relu.cpp b/apex/contrib/csrc/conv_bias_relu/conv_bias_relu.cpp +index 5fbe537..14be67b 100644 +--- a/apex/contrib/csrc/conv_bias_relu/conv_bias_relu.cpp ++++ b/apex/contrib/csrc/conv_bias_relu/conv_bias_relu.cpp +@@ -345,7 +345,7 @@ run_conv_bias(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -562,7 +562,7 @@ run_conv_bias_mask_relu(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -781,7 +781,7 @@ run_conv_cscale_cbias_relu(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -961,7 +961,7 @@ run_conv_bias_relu(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -1107,7 +1107,7 @@ run_drelu_dscale(int64_t* dy_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -1241,7 +1241,7 @@ run_drelu_dbias(int64_t* dy_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -1418,7 +1418,7 @@ run_dconv_drelu_dbias(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -1548,7 +1548,7 @@ run_dconv(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +@@ -1638,7 +1638,7 @@ run_dbias(int64_t* x_dim, + DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe()); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + checkCudnnErr(status); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + } catch (cudnn_frontend::cudnnException e) { + std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl; + } +diff --git a/apex/contrib/csrc/cudnn_gbn/norm_sample.cpp b/apex/contrib/csrc/cudnn_gbn/norm_sample.cpp +index e145021..ea2eafc 100644 +--- a/apex/contrib/csrc/cudnn_gbn/norm_sample.cpp ++++ b/apex/contrib/csrc/cudnn_gbn/norm_sample.cpp +@@ -273,7 +273,7 @@ void execute_batch_norm_forward(cudnn_frontend::ExecutionPlan plan, + .build(); + //std::cout << "variantPack " << variantPack.describe() << std::endl; + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + + // Reset local communication buffer + cudaMemsetAsync(peer_devPtrs[rank_id], 0, peer_size*4, stream); +@@ -463,7 +463,7 @@ void execute_batch_norm_backward(cudnn_frontend::ExecutionPlan plan, + .build(); + cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc()); + +- cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error", status); ++ if (status != CUDNN_STATUS_SUCCESS) { throw cudnn_frontend::cudnnException("Plan execute error", status); } + + // Reset local communication buffer + cudaMemsetAsync(peer_devPtrs[rank_id], 0, peer_size*4, stream); diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 8198b154810f..212f2fcc4c86 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -858,6 +858,8 @@ self: super: with self; { apcaccess = callPackage ../development/python-modules/apcaccess { }; + apex = callPackage ../development/python-modules/apex { }; + apeye = callPackage ../development/python-modules/apeye { }; apeye-core = callPackage ../development/python-modules/apeye-core { };