From 0d758c2aa116f2eb17ce5cde62c3b1e361a64661 Mon Sep 17 00:00:00 2001 From: Zhong Jianxin Date: Thu, 9 Mar 2023 21:16:21 +0800 Subject: [PATCH] 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,