From 4e59252a2a6afaaa8b6d712e4c572ceaac1ded2b Mon Sep 17 00:00:00 2001 From: Zhong Jianxin Date: Wed, 8 Mar 2023 15:22:11 +0800 Subject: [PATCH 1/2] python3Packages.torch: Enable Accelerate.framework on darwin to increase performance --- pkgs/development/python-modules/torch/default.nix | 4 ++-- pkgs/top-level/python-packages.nix | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index b00146b573fe..7a399f54512a 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -9,7 +9,7 @@ # Build inputs numactl, - CoreServices, libobjc, + Accelerate, CoreServices, libobjc, # Propagated build inputs numpy, pyyaml, cffi, click, typing-extensions, @@ -275,7 +275,7 @@ in buildPythonPackage rec { ++ lib.optionals rocmSupport [ openmp ] ++ lib.optionals (cudaSupport || rocmSupport) [ magma ] ++ lib.optionals stdenv.isLinux [ numactl ] - ++ lib.optionals stdenv.isDarwin [ CoreServices libobjc ]; + ++ lib.optionals stdenv.isDarwin [ Accelerate CoreServices libobjc ]; propagatedBuildInputs = [ cffi diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 21e75c132962..9bf6a2e922dc 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11675,7 +11675,7 @@ self: super: with self; { torch = callPackage ../development/python-modules/torch { cudaSupport = pkgs.config.cudaSupport or false; - inherit (pkgs.darwin.apple_sdk.frameworks) CoreServices; + inherit (pkgs.darwin.apple_sdk.frameworks) Accelerate CoreServices; inherit (pkgs.darwin) libobjc; inherit (pkgs.llvmPackages_rocm) openmp; }; From 0d758c2aa116f2eb17ce5cde62c3b1e361a64661 Mon Sep 17 00:00:00 2001 From: Zhong Jianxin Date: Thu, 9 Mar 2023 21:16:21 +0800 Subject: [PATCH 2/2] python3Packages.torch: Disable MKLDNN on aarch64-darwin It negatively impacts performance, this is also what official pytorch build does. In my test in a Macbook Pro M1 2020, using [pianotrans][1] to transcribe [cut_liszt.opus][2], transcription time: Enable mkldnn and disable Accelerate.framework: ~88s Disable mkldnn and disable Accelerate.framework: ~21s Disable mkldnn and enable Accelerate.framework: ~9s The final result is close to using torch-bin. See also https://github.com/NixOS/nixpkgs/pull/219104 [1]: https://github.com/azuwis/pianotrans [2]: https://github.com/azuwis/pianotrans/raw/master/test/cut_liszt.opus --- pkgs/development/python-modules/torch/default.nix | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index 7a399f54512a..d889dfc60201 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -1,6 +1,6 @@ { stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python, cudaSupport ? false, cudaPackages, magma, - mklDnnSupport ? true, useSystemNccl ? true, + useSystemNccl ? true, MPISupport ? false, mpi, buildDocs ? false, @@ -17,6 +17,10 @@ # Unit tests hypothesis, psutil, + # Disable MKLDNN on aarch64-darwin, it negatively impacts performance, + # this is also what official pytorch build does + mklDnnSupport ? !(stdenv.isDarwin && stdenv.isAarch64), + # virtual pkg that consistently instantiates blas across nixpkgs # See https://github.com/NixOS/nixpkgs/pull/83888 blas,