gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.apex: init at 25.09 (#516546)

Browse Source
This commit is contained in:
Connor Baker
2026-05-07 15:03:59 +00:00
committed by GitHub
parent 97067f8f34 ae21b3713d
commit d4742ef7de
3 changed files with 390 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/apex/default.nix
+187
View File
@@ -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;
};
})
pkgs/development/python-modules/apex/fix-cudnn-frontend-compat.patch
+201
View File
@@ -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<cudnn_frontend::Operation const*, 9> 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);
pkgs/top-level/python-packages.nix
+2
View File
@@ -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 { };