157 lines
4 KiB
Nix
157 lines
4 KiB
Nix
{
|
|
lib,
|
|
stdenv,
|
|
python,
|
|
buildPythonPackage,
|
|
pythonOlder,
|
|
pythonAtLeast,
|
|
fetchurl,
|
|
|
|
# nativeBuildInputs
|
|
addDriverRunpath,
|
|
autoAddDriverRunpath,
|
|
autoPatchelfHook,
|
|
|
|
# buildInputs
|
|
cudaPackages,
|
|
|
|
# dependencies
|
|
filelock,
|
|
future,
|
|
jinja2,
|
|
networkx,
|
|
numpy,
|
|
pyyaml,
|
|
requests,
|
|
setuptools,
|
|
sympy,
|
|
typing-extensions,
|
|
triton,
|
|
|
|
callPackage,
|
|
}:
|
|
|
|
let
|
|
pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
|
|
srcs = import ./binary-hashes.nix version;
|
|
unsupported = throw "Unsupported system";
|
|
version = "2.4.1";
|
|
in
|
|
buildPythonPackage {
|
|
inherit version;
|
|
|
|
pname = "torch";
|
|
# Don't forget to update torch to the same version.
|
|
|
|
format = "wheel";
|
|
|
|
disabled = (pythonOlder "3.8") || (pythonAtLeast "3.13");
|
|
|
|
src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
|
|
|
|
nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [
|
|
addDriverRunpath
|
|
autoAddDriverRunpath
|
|
autoPatchelfHook
|
|
];
|
|
|
|
buildInputs = lib.optionals stdenv.hostPlatform.isLinux (
|
|
with cudaPackages;
|
|
[
|
|
# $out/${sitePackages}/nvfuser/_C*.so wants libnvToolsExt.so.1 but torch/lib only ships
|
|
# libnvToolsExt-$hash.so.1
|
|
cuda_nvtx
|
|
|
|
cuda_cudart
|
|
cuda_cupti
|
|
cuda_nvrtc
|
|
cudnn
|
|
libcublas
|
|
libcufft
|
|
libcurand
|
|
libcusolver
|
|
libcusparse
|
|
nccl
|
|
]
|
|
);
|
|
|
|
autoPatchelfIgnoreMissingDeps = lib.optionals stdenv.hostPlatform.isLinux [
|
|
# This is the hardware-dependent userspace driver that comes from
|
|
# nvidia_x11 package. It must be deployed at runtime in
|
|
# /run/opengl-driver/lib or pointed at by LD_LIBRARY_PATH variable, rather
|
|
# than pinned in runpath
|
|
"libcuda.so.1"
|
|
];
|
|
|
|
dependencies = [
|
|
filelock
|
|
future
|
|
jinja2
|
|
networkx
|
|
numpy
|
|
pyyaml
|
|
requests
|
|
setuptools
|
|
sympy
|
|
typing-extensions
|
|
] ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [ triton ];
|
|
|
|
postInstall = ''
|
|
# ONNX conversion
|
|
rm -rf $out/bin
|
|
'';
|
|
|
|
postFixup = lib.optionalString stdenv.hostPlatform.isLinux ''
|
|
addAutoPatchelfSearchPath "$out/${python.sitePackages}/torch/lib"
|
|
'';
|
|
|
|
# See https://github.com/NixOS/nixpkgs/issues/296179
|
|
#
|
|
# This is a quick hack to add `libnvrtc` to the runpath so that torch can find
|
|
# it when it is needed at runtime.
|
|
extraRunpaths = lib.optionals stdenv.hostPlatform.isLinux [
|
|
"${lib.getLib cudaPackages.cuda_nvrtc}/lib"
|
|
];
|
|
postPhases = lib.optionals stdenv.hostPlatform.isLinux [ "postPatchelfPhase" ];
|
|
postPatchelfPhase = ''
|
|
while IFS= read -r -d $'\0' elf ; do
|
|
for extra in $extraRunpaths ; do
|
|
echo patchelf "$elf" --add-rpath "$extra" >&2
|
|
patchelf "$elf" --add-rpath "$extra"
|
|
done
|
|
done < <(
|
|
find "''${!outputLib}" "$out" -type f -iname '*.so' -print0
|
|
)
|
|
'';
|
|
|
|
# The wheel-binary is not stripped to avoid the error of `ImportError: libtorch_cuda_cpp.so: ELF load command address/offset not properly aligned.`.
|
|
dontStrip = true;
|
|
|
|
pythonImportsCheck = [ "torch" ];
|
|
|
|
passthru.tests = callPackage ./tests.nix { };
|
|
|
|
meta = {
|
|
description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration";
|
|
homepage = "https://pytorch.org/";
|
|
changelog = "https://github.com/pytorch/pytorch/releases/tag/v${version}";
|
|
# Includes CUDA and Intel MKL, but redistributions of the binary are not limited.
|
|
# https://docs.nvidia.com/cuda/eula/index.html
|
|
# https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
|
|
# torch's license is BSD3.
|
|
# torch-bin used to vendor CUDA. It still links against CUDA and MKL.
|
|
license = with lib.licenses; [
|
|
bsd3
|
|
issl
|
|
unfreeRedistributable
|
|
];
|
|
sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
|
|
platforms = [
|
|
"aarch64-darwin"
|
|
"aarch64-linux"
|
|
"x86_64-linux"
|
|
];
|
|
hydraPlatforms = [ ]; # output size 3.2G on 1.11.0
|
|
maintainers = with lib.maintainers; [ junjihashimoto ];
|
|
};
|
|
}
|