From 1fdc859a69b2b5760dfc897e6b1ef952e1ee75d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Wed, 10 Dec 2025 12:12:46 +0000 Subject: [PATCH] Overlay Torch in the package set Before this change, we wouldn't overlay Torch in the package set. The main reason was that evaluation was too expensive -- even more so when we also had C++98 and C++11 variants (which doubled the number of Torch versions). However, this does not really work well with dependencies (some dependencies like `einops` will rely on Torch) or extra test dependencies that a user adds to a flake. So overlay Torch so that all of nixpkgs sees the version that is defined in the build set. In this change, the `cxx11Abi` is also removed and the C++11 ABI is set as the default. Upstream does not build for the C++98 ABI anymore, so it is time to remove the support before it bitrots. --- examples/relu-specific-torch/flake.nix | 1 - lib/build-sets.nix | 185 +++++++----------- lib/build-variants.nix | 4 +- lib/build.nix | 1 - lib/version-utils.nix | 1 - pkgs/python-modules/torch/binary/generic.nix | 4 - .../torch/binary/torch-versions.json | 16 -- .../torch/source/2_8/default.nix | 5 +- .../torch/source/2_9/default.nix | 5 +- pkgs/python-modules/torch/variant.nix | 5 +- versions.nix | 17 -- 11 files changed, 71 insertions(+), 173 deletions(-) diff --git a/examples/relu-specific-torch/flake.nix b/examples/relu-specific-torch/flake.nix index 5b01e602..bc21f5d4 100644 --- a/examples/relu-specific-torch/flake.nix +++ b/examples/relu-specific-torch/flake.nix @@ -17,7 +17,6 @@ { torchVersion = "2.9"; cudaVersion = "12.8"; - cxx11Abi = true; systems = [ "x86_64-linux" "aarch64-linux" diff --git a/lib/build-sets.nix b/lib/build-sets.nix index dc4df1ea..01ea458f 100644 --- a/lib/build-sets.nix +++ b/lib/build-sets.nix @@ -25,26 +25,22 @@ let in builtins.map (buildConfig: buildConfig // { backend = backend buildConfig; }) systemBuildConfigs; - cudaVersions = - let - withCuda = builtins.filter (torchVersion: torchVersion ? cudaVersion) torchVersions; - in - lib.unique (builtins.map (torchVersion: torchVersion.cudaVersion) withCuda); - - rocmVersions = - let - withRocm = builtins.filter (torchVersion: torchVersion ? rocmVersion) torchVersions; - in - lib.unique (builtins.map (torchVersion: torchVersion.rocmVersion) withRocm); - - xpuVersions = - let - withXpu = builtins.filter (torchVersion: torchVersion ? xpuVersion) torchVersions; - in - lib.unique (builtins.map (torchVersion: torchVersion.xpuVersion) withXpu); - flattenVersion = version: lib.replaceStrings [ "." ] [ "_" ] (lib.versions.pad 2 version); + overlayForTorchVersion = torchVersion: sourceBuild: self: super: { + pythonPackagesExtensions = super.pythonPackagesExtensions ++ [ + ( + python-self: python-super: with python-self; { + torch = + if sourceBuild then + python-self."torch_${flattenVersion torchVersion}" + else + python-self."torch-bin_${flattenVersion torchVersion}"; + } + ) + ]; + }; + # An overlay that overides CUDA to the given version. overlayForCudaVersion = cudaVersion: self: super: { cudaPackages = super."cudaPackages_${flattenVersion cudaVersion}"; @@ -57,6 +53,38 @@ let overlayForXpuVersion = xpuVersion: self: super: { xpuPackages = super."xpuPackages_${flattenVersion xpuVersion}"; }; + + backendConfig = { + cpu = { + allowUnfree = true; + }; + + cuda = { + allowUnfree = true; + cudaSupport = true; + }; + + metal = { + allowUnfree = true; + metalSupport = true; + }; + + rocm = { + allowUnfree = true; + rocmSupport = true; + }; + + xpu = { + allowUnfree = true; + xpuSupport = true; + }; + }; + + xpuConfig = { + allowUnfree = true; + xpuSupport = true; + }; + # Construct the nixpkgs package set for the given versions. mkBuildSet = buildConfig@{ @@ -67,34 +95,38 @@ let rocmVersion ? null, xpuVersion ? null, torchVersion, - cxx11Abi, system, bundleBuild ? false, sourceBuild ? false, }: let - pkgs = + backendOverlay = if buildConfig.backend == "cpu" then - pkgsForCpu + [ ] else if buildConfig.backend == "cuda" then - pkgsByCudaVer.${cudaVersion} + [ (overlayForCudaVersion buildConfig.cudaVersion) ] else if buildConfig.backend == "rocm" then - pkgsByRocmVer.${rocmVersion} + [ (overlayForRocmVersion buildConfig.rocmVersion) ] else if buildConfig.backend == "metal" then - pkgsForMetal + [ ] else if buildConfig.backend == "xpu" then - pkgsByXpuVer.${xpuVersion} + [ (overlayForXpuVersion buildConfig.xpuVersion) ] else throw "No compute framework set in Torch version"; - torch = - if sourceBuild then - pkgs.python3.pkgs."torch_${flattenVersion torchVersion}".override { - inherit cxx11Abi; - } - else - pkgs.python3.pkgs."torch-bin_${flattenVersion torchVersion}".override { - inherit cxx11Abi; - }; + config = + backendConfig.${buildConfig.backend} or (throw "No backend config for ${buildConfig.backend}"); + + pkgs = import nixpkgs { + inherit config system; + overlays = [ + overlay + ] + ++ backendOverlay + ++ [ (overlayForTorchVersion torchVersion sourceBuild) ]; + }; + + torch = pkgs.python3.pkgs.torch; + extension = pkgs.callPackage ./torch-extension { inherit torch; }; in { @@ -106,90 +138,5 @@ let bundleBuild ; }; - pkgsForXpuVersions = - xpuVersions: - builtins.listToAttrs ( - map (xpuVersion: { - name = xpuVersion; - value = import nixpkgs { - inherit system; - config = { - allowUnfree = true; - xpuSupport = true; - }; - overlays = [ - overlay - (overlayForXpuVersion xpuVersion) - ]; - }; - }) xpuVersions - ); - pkgsByXpuVer = pkgsForXpuVersions xpuVersions; - - pkgsForMetal = import nixpkgs { - inherit system; - config = { - allowUnfree = true; - metalSupport = true; - }; - overlays = [ - overlay - ]; - }; - - pkgsForCpu = import nixpkgs { - inherit system; - config = { - allowUnfree = true; - }; - overlays = [ - overlay - ]; - }; - - # Instantiate nixpkgs for the given CUDA versions. Returns - # an attribute set like `{ "12.4" = ; ... }`. - pkgsForCudaVersions = - cudaVersions: - builtins.listToAttrs ( - map (cudaVersion: { - name = cudaVersion; - value = import nixpkgs { - inherit system; - config = { - allowUnfree = true; - cudaSupport = true; - }; - overlays = [ - overlay - (overlayForCudaVersion cudaVersion) - ]; - }; - }) cudaVersions - ); - - pkgsByCudaVer = pkgsForCudaVersions cudaVersions; - - pkgsForRocmVersions = - rocmVersions: - builtins.listToAttrs ( - map (rocmVersion: { - name = rocmVersion; - value = import nixpkgs { - inherit system; - config = { - allowUnfree = true; - rocmSupport = true; - }; - overlays = [ - overlay - (overlayForRocmVersion rocmVersion) - ]; - }; - }) rocmVersions - ); - - pkgsByRocmVer = pkgsForRocmVersions rocmVersions; - in map mkBuildSet (buildConfigs system) diff --git a/lib/build-variants.nix b/lib/build-variants.nix index 7201211a..f7dfb777 100644 --- a/lib/build-variants.nix +++ b/lib/build-variants.nix @@ -9,7 +9,7 @@ rec { buildName = let - inherit (import ./version-utils.nix { inherit lib; }) abiString flattenVersion; + inherit (import ./version-utils.nix { inherit lib; }) flattenVersion; computeString = version: if backend version == "cpu" then @@ -29,7 +29,7 @@ rec { if version.system == "aarch64-darwin" then "torch${flattenVersion version.torchVersion}-${computeString version}-${version.system}" else - "torch${flattenVersion version.torchVersion}-${abiString version.cxx11Abi}-${computeString version}-${version.system}"; + "torch${flattenVersion version.torchVersion}-cxx11-${computeString version}-${version.system}"; # Build variants included in bundle builds. buildVariants = diff --git a/lib/build.nix b/lib/build.nix index c00ea764..8d8a4dbe 100644 --- a/lib/build.nix +++ b/lib/build.nix @@ -12,7 +12,6 @@ }: let - abi = torch: if torch.passthru.cxx11Abi then "cxx11" else "cxx98"; supportedCudaCapabilities = builtins.fromJSON ( builtins.readFile ../build2cmake/src/cuda_supported_archs.json ); diff --git a/lib/version-utils.nix b/lib/version-utils.nix index 8474ab78..6df34f98 100644 --- a/lib/version-utils.nix +++ b/lib/version-utils.nix @@ -6,5 +6,4 @@ in { flattenVersion = version: lib.replaceStrings [ "." ] [ "" ] (versions.majorMinor (versions.pad 2 version)); - abiString = cxx11Abi: if cxx11Abi then "cxx11" else "cxx98"; } diff --git a/pkgs/python-modules/torch/binary/generic.nix b/pkgs/python-modules/torch/binary/generic.nix index 83999272..98209383 100644 --- a/pkgs/python-modules/torch/binary/generic.nix +++ b/pkgs/python-modules/torch/binary/generic.nix @@ -43,8 +43,6 @@ url, hash, version, - # Remove, needed for compat. - cxx11Abi ? true, effectiveStdenv ? if cudaSupport then cudaPackages.backendStdenv else stdenv, }: @@ -322,7 +320,6 @@ buildPythonPackage { inherit cudaSupport cudaPackages - cxx11Abi rocmSupport rocmPackages xpuSupport @@ -333,7 +330,6 @@ buildPythonPackage { rocmArchs = if rocmSupport then supportedTorchRocmArchs else [ ]; } // (callPackage ../variant.nix { - inherit cxx11Abi; torchVersion = version; }); diff --git a/pkgs/python-modules/torch/binary/torch-versions.json b/pkgs/python-modules/torch/binary/torch-versions.json index 70167a72..d65a0560 100644 --- a/pkgs/python-modules/torch/binary/torch-versions.json +++ b/pkgs/python-modules/torch/binary/torch-versions.json @@ -2,98 +2,82 @@ { "torchVersion": "2.8.0", "cudaVersion": "12.6", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.8.0", "cudaVersion": "12.8", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.8.0", "cudaVersion": "12.9", - "cxx11Abi": true, "systems": ["x86_64-linux", "aarch64-linux"] }, { "torchVersion": "2.8.0", "rocmVersion": "6.3", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.8.0", "rocmVersion": "6.4", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.8.0", - "cxx11Abi": true, "metal": true, "systems": ["aarch64-darwin"] }, { "torchVersion": "2.8.0", - "cxx11Abi": true, "cpu": true, "systems": ["aarch64-linux", "x86_64-linux"] }, { "torchVersion": "2.8.0", "xpuVersion": "2025.1.3", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.9.0", "cudaVersion": "12.6", - "cxx11Abi": true, "systems": ["x86_64-linux", "aarch64-linux"] }, { "torchVersion": "2.9.0", "cudaVersion": "12.8", - "cxx11Abi": true, "systems": ["x86_64-linux", "aarch64-linux"] }, { "torchVersion": "2.9.0", "cudaVersion": "13.0", - "cxx11Abi": true, "systems": ["x86_64-linux", "aarch64-linux"] }, { "torchVersion": "2.9.0", "rocmVersion": "6.3", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.9.0", "rocmVersion": "6.4", - "cxx11Abi": true, "systems": ["x86_64-linux"] }, { "torchVersion": "2.9.0", - "cxx11Abi": true, "metal": true, "systems": ["aarch64-darwin"] }, { "torchVersion": "2.9.0", - "cxx11Abi": true, "cpu": true, "systems": ["aarch64-linux", "x86_64-linux"] }, { "torchVersion": "2.9.0", "xpuVersion": "2025.2.1", - "cxx11Abi": true, "systems": ["x86_64-linux"] } ] diff --git a/pkgs/python-modules/torch/source/2_8/default.nix b/pkgs/python-modules/torch/source/2_8/default.nix index 7b64209c..d62e89a9 100644 --- a/pkgs/python-modules/torch/source/2_8/default.nix +++ b/pkgs/python-modules/torch/source/2_8/default.nix @@ -26,7 +26,6 @@ mpi, nvtx, buildDocs ? false, - cxx11Abi ? true, # tests.cudaAvailable: callPackage, @@ -420,7 +419,7 @@ buildPythonPackage rec { CMAKE_GENERATOR = "Ninja"; # Whether to use C++11 ABI (or earlier). - _GLIBCXX_USE_CXX11_ABI = setBool cxx11Abi; + _GLIBCXX_USE_CXX11_ABI = setBool true; # Unlike MKL, oneDNN (née MKLDNN) is FOSS, so we enable support for # it by default. PyTorch currently uses its own vendored version @@ -728,7 +727,6 @@ buildPythonPackage rec { inherit cudaSupport cudaPackages - cxx11Abi rocmSupport rocmPackages xpuSupport @@ -743,7 +741,6 @@ buildPythonPackage rec { tests = callPackage ./tests.nix { }; } // (callPackage ../variant.nix { - inherit cxx11Abi; torchVersion = version; }); diff --git a/pkgs/python-modules/torch/source/2_9/default.nix b/pkgs/python-modules/torch/source/2_9/default.nix index 6b27b60a..b661ae16 100644 --- a/pkgs/python-modules/torch/source/2_9/default.nix +++ b/pkgs/python-modules/torch/source/2_9/default.nix @@ -27,7 +27,6 @@ mpi, nvtx, buildDocs ? false, - cxx11Abi ? true, # tests.cudaAvailable: callPackage, @@ -403,7 +402,7 @@ buildPythonPackage rec { CMAKE_GENERATOR = "Ninja"; # Whether to use C++11 ABI (or earlier). - _GLIBCXX_USE_CXX11_ABI = setBool cxx11Abi; + _GLIBCXX_USE_CXX11_ABI = setBool true; # Unlike MKL, oneDNN (née MKLDNN) is FOSS, so we enable support for # it by default. PyTorch currently uses its own vendored version @@ -711,7 +710,6 @@ buildPythonPackage rec { inherit cudaSupport cudaPackages - cxx11Abi rocmSupport rocmPackages xpuSupport @@ -726,7 +724,6 @@ buildPythonPackage rec { tests = callPackage ./tests.nix { }; } // (callPackage ../variant.nix { - inherit cxx11Abi; torchVersion = version; }); diff --git a/pkgs/python-modules/torch/variant.nix b/pkgs/python-modules/torch/variant.nix index c783e434..880af20c 100644 --- a/pkgs/python-modules/torch/variant.nix +++ b/pkgs/python-modules/torch/variant.nix @@ -12,15 +12,12 @@ lib, stdenv, - cxx11Abi, - torchVersion, }: let flattenVersion = version: lib.replaceStrings [ "." ] [ "" ] (lib.versions.majorMinor (lib.versions.pad 2 version)); - abiString = cxx11Abi: if cxx11Abi then "cxx11" else "cxx98"; backend = if cudaSupport then "cuda" @@ -49,7 +46,7 @@ in if stdenv.hostPlatform.system == "aarch64-darwin" then "torch${flattenVersion (lib.versions.majorMinor torchVersion)}-${computeString}-${stdenv.hostPlatform.system}" else - "torch${flattenVersion (lib.versions.majorMinor torchVersion)}-${abiString cxx11Abi}-${computeString}-${stdenv.hostPlatform.system}"; + "torch${flattenVersion (lib.versions.majorMinor torchVersion)}-cxx11-${computeString}-${stdenv.hostPlatform.system}"; noarchVariant = "torch-${backend}"; } diff --git a/versions.nix b/versions.nix index 8b842d7c..cbd937c4 100644 --- a/versions.nix +++ b/versions.nix @@ -2,14 +2,12 @@ { torchVersion = "2.8"; xpuVersion = "2025.1.3"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; bundleBuild = true; } { torchVersion = "2.8"; cudaVersion = "12.6"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; @@ -18,7 +16,6 @@ { torchVersion = "2.8"; cudaVersion = "12.8"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; @@ -27,7 +24,6 @@ { torchVersion = "2.8"; cudaVersion = "12.9"; - cxx11Abi = true; systems = [ "x86_64-linux" "aarch64-linux" @@ -37,27 +33,23 @@ { torchVersion = "2.8"; rocmVersion = "6.3.4"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; bundleBuild = true; } { torchVersion = "2.8"; rocmVersion = "6.4.2"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; bundleBuild = true; } { torchVersion = "2.8"; - cxx11Abi = true; metal = true; systems = [ "aarch64-darwin" ]; bundleBuild = true; } { torchVersion = "2.8"; - cxx11Abi = true; cpu = true; systems = [ "aarch64-darwin" @@ -69,14 +61,12 @@ { torchVersion = "2.9"; xpuVersion = "2025.2.1"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; bundleBuild = true; } { torchVersion = "2.9"; cudaVersion = "12.6"; - cxx11Abi = true; systems = [ "x86_64-linux" "aarch64-linux" @@ -86,7 +76,6 @@ { torchVersion = "2.9"; cudaVersion = "12.8"; - cxx11Abi = true; systems = [ "x86_64-linux" "aarch64-linux" @@ -96,7 +85,6 @@ { torchVersion = "2.9"; cudaVersion = "13.0"; - cxx11Abi = true; systems = [ "x86_64-linux" "aarch64-linux" @@ -106,27 +94,23 @@ { torchVersion = "2.9"; rocmVersion = "6.3.4"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; bundleBuild = true; } { torchVersion = "2.9"; rocmVersion = "6.4.2"; - cxx11Abi = true; systems = [ "x86_64-linux" ]; bundleBuild = true; } { torchVersion = "2.9"; - cxx11Abi = true; metal = true; systems = [ "aarch64-darwin" ]; bundleBuild = true; } { torchVersion = "2.9"; - cxx11Abi = true; cpu = true; systems = [ "aarch64-darwin" @@ -140,7 +124,6 @@ #{ # torchVersion = "2.8"; # cudaVersion = "12.4"; - # cxx11Abi = true; # systems = [ # "x86_64-linux" # "aarch64-linux"