python312Packages.torchao: init at 0.11.0

2025-06-16 00:04:00 +02:00
parent c4901fb437
commit d768d6deac
2 changed files with 102 additions and 0 deletions
@@ -0,0 +1,100 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  torch,
+
+  # tests
+  bitsandbytes,
+  expecttest,
+  fire,
+  pytest-xdist,
+  pytestCheckHook,
+  parameterized,
+  tabulate,
+  transformers,
+  unittest-xml-reporting,
+}:
+
+buildPythonPackage rec {
+  pname = "ao";
+  version = "0.11.0";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "pytorch";
+    repo = "ao";
+    tag = "v${version}";
+    hash = "sha256-CNb9xaubOmIRanLq3TM4sBbszTcVK/WFpcq/sWpof44=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    torch
+  ];
+
+  env = {
+    USE_SYSTEM_LIBS = true;
+  };
+
+  # Otherwise, the tests are loading the python module from the source instead of the installed one
+  preCheck = ''
+    rm -rf torchao
+  '';
+
+  pythonImportsCheck = [
+    "torchao"
+  ];
+
+  nativeCheckInputs = [
+    bitsandbytes
+    expecttest
+    fire
+    parameterized
+    pytest-xdist
+    pytestCheckHook
+    tabulate
+    transformers
+    unittest-xml-reporting
+  ];
+
+  disabledTests =
+    [
+      # Requires internet access
+      "test_on_dummy_distilbert"
+
+      # FileNotFoundError: [Errno 2] No such file or directory: 'checkpoints/meta-llama/Llama-2-7b-chat-hf/model.pth'
+      "test_gptq_mt"
+    ]
+    ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
+      # RuntimeError: failed to initialize QNNPACK
+      "test_smooth_linear_cpu"
+
+      # torch._inductor.exc.InductorError: LoweringException: AssertionError: Expect L1_cache_size > 0 but got 0
+      "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"
+
+      # FileNotFoundError: [Errno 2] No such file or directory: 'test.pth'
+      "test_save_load_int4woqtensors_2_cpu"
+      "test_save_load_int8woqtensors_0_cpu"
+      "test_save_load_int8woqtensors_1_cpu"
+    ];
+
+  meta = {
+    description = "PyTorch native quantization and sparsity for training and inference";
+    homepage = "https://github.com/pytorch/ao";
+    changelog = "https://github.com/pytorch/ao/releases/tag/v${version}";
+    license = lib.licenses.bsd3;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}
@@ -17793,6 +17793,8 @@ self: super: with self; {

  torchWithoutRocm = self.torch.override { rocmSupport = false; };

+  torchao = callPackage ../development/python-modules/torchao { };
+
  torchaudio = callPackage ../development/python-modules/torchaudio { };

  torchaudio-bin = callPackage ../development/python-modules/torchaudio/bin.nix { };