python3Packages.torchao: 0.13.0 -> 0.14.1 (#454442)

pkgs/development/python-modules/torchao/default.nix

@@ -3,10 +3,18 @@
   stdenv,
   buildPythonPackage,
   fetchFromGitHub,
+  replaceVars,
 
   # build-system
   setuptools,
 
+  # nativeBuildInputs
+  cmake,
+
+  # buildInputs
+  cpuinfo,
+  llvmPackages,
+
   # dependencies
   torch,
 
@@ -24,20 +32,43 @@
 
 buildPythonPackage rec {
   pname = "ao";
-  version = "0.13.0";
+  version = "0.14.1";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "pytorch";
     repo = "ao";
     tag = "v${version}";
-    hash = "sha256-R9H4+KkKuOzsunM3A5LT8upH1TfkHrD+BZerToCHwjo=";
+    hash = "sha256-L9Eoul7Nar/+gS44+hA8JbfxCgkMH5xAMCleggAZn7c=";
   };
 
+  patches = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    ./use-system-cpuinfo.patch
+    (replaceVars ./use-llvm-openmp.patch {
+      inherit (llvmPackages) openmp;
+    })
+  ];
+
   build-system = [
     setuptools
   ];
 
+  nativeBuildInputs = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    cmake
+  ];
+  dontUseCmakeConfigure = true;
+
+  buildInputs = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    cpuinfo
+  ];
+
+  propagatedBuildInputs = lib.optionals stdenv.hostPlatform.isDarwin [
+    # Otherwise, torch will fail to include `omp.h`:
+    # torch._inductor.exc.InductorError: CppCompileError: C++ compile error
+    # OpenMP support not found.
+    llvmPackages.openmp
+  ];
+
   dependencies = [
     torch
   ];
@@ -90,6 +121,36 @@ buildPythonPackage rec {
     "test_save_load_int4woqtensors_2_cpu"
     "test_save_load_int8woqtensors_0_cpu"
     "test_save_load_int8woqtensors_1_cpu"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isDarwin) [
+    # AssertionError: Scalars are not equal!
+    "test_comm"
+    "test_fsdp2"
+    "test_fsdp2_correctness"
+    "test_precompute_bitnet_scale"
+    "test_qlora_fsdp2"
+    "test_uneven_shard"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    # RuntimeError: No packed_weights_format was selected
+    "TestIntxOpaqueTensor"
+    "test_accuracy_kleidiai"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [
+    # Flaky: [gw0] node down: keyboard-interrupt
+    "test_int8_weight_only_quant_with_freeze_0_cpu"
+    "test_int8_weight_only_quant_with_freeze_1_cpu"
+    "test_int8_weight_only_quant_with_freeze_2_cpu"
+  ];
+
+  disabledTestPaths = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+    # Require unpackaged 'coremltools'
+    "test/prototype/test_groupwise_lowbit_weight_lut_quantizer.py"
+
+    # AttributeError: '_OpNamespace' 'mkldnn' object has no attribute '_is_mkldnn_acl_supported'
+    "test/quantization/pt2e/test_arm_inductor_quantizer.py"
+    "test/quantization/pt2e/test_x86inductor_fusion.py"
+    "test/quantization/pt2e/test_x86inductor_quantizer.py"
   ];
 
   meta = {

pkgs/development/python-modules/torchao/use-llvm-openmp.patch

@@ -0,0 +1,13 @@
+diff --git a/torchao/csrc/cpu/shared_kernels/Utils.cmake b/torchao/csrc/cpu/shared_kernels/Utils.cmake
+index be7004784..0e1a2ed0e 100644
+--- a/torchao/csrc/cpu/shared_kernels/Utils.cmake
++++ b/torchao/csrc/cpu/shared_kernels/Utils.cmake
+@@ -21,7 +21,7 @@ function(target_link_torchao_parallel_backend target_name torchao_parallel_backe
+         target_link_libraries(${target_name} PRIVATE "${TORCH_LIBRARIES}")
+ 
+         target_compile_definitions(${target_name} PRIVATE TORCHAO_PARALLEL_ATEN=1 AT_PARALLEL_OPENMP=1 INTRA_OP_PARALLEL=1)
+-        target_link_libraries(${target_name} PRIVATE ${TORCH_INSTALL_PREFIX}/lib/libomp${CMAKE_SHARED_LIBRARY_SUFFIX})
++        target_link_libraries(${target_name} PRIVATE @openmp@/lib/libomp${CMAKE_SHARED_LIBRARY_SUFFIX})
+ 
+     elseif(TORCHAO_PARALLEL_BACKEND_TOUPPER STREQUAL "EXECUTORCH")
+         message(STATUS "Building with TORCHAO_PARALLEL_BACKEND=TORCHAO_PARALLEL_EXECUTORCH")

pkgs/development/python-modules/torchao/use-system-cpuinfo.patch

@@ -0,0 +1,22 @@
+diff --git a/torchao/csrc/cpu/CMakeLists.txt b/torchao/csrc/cpu/CMakeLists.txt
+index aaea27ec7..53dee80fd 100644
+--- a/torchao/csrc/cpu/CMakeLists.txt
++++ b/torchao/csrc/cpu/CMakeLists.txt
+@@ -49,16 +49,7 @@ if (NOT TARGET cpuinfo)
+     add_compile_options(-Wno-unused-function -Wno-unused-variable)
+ 
+     # set(CMAKE_POLICY_VERSION_MINIMUM 3.5)
+-    include(FetchContent)
+-    set(CPUINFO_BUILD_UNIT_TESTS OFF CACHE BOOL "" FORCE)
+-    set(CPUINFO_BUILD_MOCK_TESTS OFF CACHE BOOL "" FORCE)
+-    set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
+-    FetchContent_Declare(cpuinfo
+-        GIT_REPOSITORY https://github.com/pytorch/cpuinfo.git
+-        GIT_TAG c61fe919607bbc534d7a5a5707bdd7041e72c5ff
+-    )
+-    FetchContent_MakeAvailable(
+-        cpuinfo)
++    find_package(cpuinfo REQUIRED)
+ 
+     cmake_policy(POP)
+ endif()