From 55f56e92b7457d2414985346841da34049e3efba Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Mon, 6 Apr 2026 08:37:22 +0000 Subject: [PATCH] python3Packages.torchao: 0.16.0 -> 0.17.0 Diff: https://github.com/pytorch/ao/compare/v0.16.0...v0.17.0 Changelog: https://github.com/pytorch/ao/releases/tag/v0.17.0 --- .../python-modules/torchao/default.nix | 42 ++++++++++++++----- .../torchao/use-llvm-openmp.patch | 20 +++++---- 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/pkgs/development/python-modules/torchao/default.nix b/pkgs/development/python-modules/torchao/default.nix index 75fec3ec39b2..8f546dcd71a9 100644 --- a/pkgs/development/python-modules/torchao/default.nix +++ b/pkgs/development/python-modules/torchao/default.nix @@ -27,27 +27,37 @@ pytestCheckHook, parameterized, tabulate, + torchvision, transformers, unittest-xml-reporting, }: - +let + inherit (stdenv.hostPlatform) + isDarwin + isLinux + isAarch64 + isx86_64 + ; + isAarch64Darwin = isDarwin && isAarch64; + isAarch64Linux = isLinux && isAarch64; +in buildPythonPackage (finalAttrs: { pname = "torchao"; - version = "0.16.0"; + version = "0.17.0"; pyproject = true; src = fetchFromGitHub { owner = "pytorch"; repo = "ao"; tag = "v${finalAttrs.version}"; - hash = "sha256-FyBsIVb3zdKtA8Vqjt301bRrGIoyeqiOUADVFGxiRPY="; + hash = "sha256-Mry6jsZKkoC8dq3fYNsRyGbL4+S8ZYuHpkETNDy5qsg="; }; # AttributeError: 'typing.Union' object has no attribute '__module__' and no __dict__ for setting # new attributes. Did you mean: '__reduce__'? disabled = pythonAtLeast "3.14"; - patches = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [ + patches = lib.optionals isAarch64Darwin [ ./use-system-cpuinfo.patch (replaceVars ./use-llvm-openmp.patch { inherit (llvmPackages) openmp; @@ -58,16 +68,16 @@ buildPythonPackage (finalAttrs: { setuptools ]; - nativeBuildInputs = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [ + nativeBuildInputs = lib.optionals isAarch64Darwin [ cmake ]; dontUseCmakeConfigure = true; - buildInputs = lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [ + buildInputs = lib.optionals isAarch64Darwin [ cpuinfo ]; - propagatedBuildInputs = lib.optionals stdenv.hostPlatform.isDarwin [ + propagatedBuildInputs = lib.optionals isDarwin [ # Otherwise, torch will fail to include `omp.h`: # torch._inductor.exc.InductorError: CppCompileError: C++ compile error # OpenMP support not found. @@ -97,6 +107,7 @@ buildPythonPackage (finalAttrs: { pytest-xdist pytestCheckHook tabulate + torchvision transformers unittest-xml-reporting ]; @@ -160,7 +171,7 @@ buildPythonPackage (finalAttrs: { # execnet.gateway_base.DumpError: can't serialize "test_numerical_consistency_per_tensor" ] - ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ + ++ lib.optionals isAarch64Linux [ # AssertionError: tensor(False) is not true "test_quantize_per_token_cpu" @@ -177,7 +188,7 @@ buildPythonPackage (finalAttrs: { "test_save_load_int8woqtensors_0_cpu" "test_save_load_int8woqtensors_1_cpu" ] - ++ lib.optionals stdenv.hostPlatform.isDarwin [ + ++ lib.optionals isDarwin [ # AssertionError: Scalars are not equal! "test_comm" "test_fsdp2" @@ -237,6 +248,13 @@ buildPythonPackage (finalAttrs: { "test_optim_smoke_optim_name_AdamFp8_float32_device_mps" "test_subclass_slice_subclass0_shape1_device_mps" + # RuntimeError: Expected to find "triton_per_fused" but did not find it + "test_available_gpu_kernels_device_mps" + + # RuntimeError: quantized engine NoQEngine is not supported + "test_qat_mobilenet_v2" + "test_qat_resnet18" + # Crash (Trace/BPT trap: 5) "test_copy__mismatch_metadata_apply_quant0" "test_copy__mismatch_metadata_apply_quant1" @@ -278,7 +296,7 @@ buildPythonPackage (finalAttrs: { "test_workflow_e2e_numerics_config4" "test_workflow_e2e_numerics_config5" ] - ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [ + ++ lib.optionals (isDarwin && 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" @@ -296,7 +314,7 @@ buildPythonPackage (finalAttrs: { # ImportError: cannot import name 'fp8_blockwise_weight_dequant' from 'torchao.kernel.blockwise_quantization' "test/kernel/test_blockwise_triton.py" ] - ++ lib.optionals stdenv.hostPlatform.isDarwin [ + ++ lib.optionals isDarwin [ # Require unpackaged 'coremltools' "test/prototype/test_groupwise_lowbit_weight_lut_quantizer.py" @@ -316,6 +334,8 @@ buildPythonPackage (finalAttrs: { "test/test_low_bit_optim.py::TestQuantize::test_bf16_stochastic_round_dtensor_device_mps_compile_True" ]; + __darwinAllowLocalNetworking = true; + meta = { description = "PyTorch native quantization and sparsity for training and inference"; homepage = "https://github.com/pytorch/ao"; diff --git a/pkgs/development/python-modules/torchao/use-llvm-openmp.patch b/pkgs/development/python-modules/torchao/use-llvm-openmp.patch index b399b347da39..5d655a8f2202 100644 --- a/pkgs/development/python-modules/torchao/use-llvm-openmp.patch +++ b/pkgs/development/python-modules/torchao/use-llvm-openmp.patch @@ -1,13 +1,17 @@ diff --git a/torchao/csrc/cpu/shared_kernels/Utils.cmake b/torchao/csrc/cpu/shared_kernels/Utils.cmake -index be7004784..0e1a2ed0e 100644 +index 3ff024538..0763cab67 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}) +@@ -30,10 +30,10 @@ function(target_link_torchao_parallel_backend target_name torchao_parallel_backe + if(_TORCH_OMP_RUNTIME) + target_link_libraries(${target_name} PRIVATE "${_TORCH_OMP_RUNTIME}") + else() +- 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}) + endif() + else() +- 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}) + endif() elseif(TORCHAO_PARALLEL_BACKEND_TOUPPER STREQUAL "EXECUTORCH") - message(STATUS "Building with TORCHAO_PARALLEL_BACKEND=TORCHAO_PARALLEL_EXECUTORCH")