python3Packages.torchtune: 0.6.2-unstable-2026-02-19 -> 0.6.1-unstable-2026-04-23 (#537564)

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

@@ -28,29 +28,82 @@
 
   # tests
   comet-ml,
+  expecttest,
   mlflow,
   pytest-integration,
   pytest-mock,
   pytestCheckHook,
+  tensorboard,
   writableTmpDirAsHomeHook,
+
+  # passthru
+  nix-update-script,
 }:
 
 buildPythonPackage (finalAttrs: {
   pname = "torchtune";
-  version = "0.6.2-unstable-2026-02-19";
+  version = "0.6.1-unstable-2026-04-23";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "meta-pytorch";
     repo = "torchtune";
-    rev = "6f2aa7254458145f99d7004cbd6ebc8e53a06404";
-    hash = "sha256-ryR5iO3IwkoLdMLSFGhHCLl0P8yD+GQdZFEE6M/EYh0=";
+    rev = "bd2a0fc7c31430972728494fa01aaeeb0ebf1ba1";
+    hash = "sha256-6jE8+ZCm46qyoSOCkBjxsXNvtVEOUP6v3NEmMh+ocl8=";
   };
 
+  postPatch = ''
+    substituteInPlace \
+      tests/torchtune/modules/low_precision/test_nf4_dispatch_registration.py \
+      torchtune/modules/low_precision/_register_nf4_dispatch_ops.py \
+      torchtune/modules/low_precision/nf4_linear.py \
+      torchtune/modules/peft/dora.py \
+      torchtune/modules/peft/lora.py \
+      --replace-fail \
+        "from torchao.quantization import to_nf4" \
+        "from torchao.dtypes import to_nf4" \
+
+    substituteInPlace \
+      tests/torchtune/models/llama2/test_lora_llama2.py \
+      tests/torchtune/models/phi3/test_lora_phi3.py \
+      tests/torchtune/modules/low_precision/test_nf4_linear.py \
+      tests/torchtune/training/test_distributed.py \
+      torchtune/modules/common_utils.py \
+      torchtune/training/_activation_offloading.py \
+      --replace-fail \
+        "from torchao.quantization import NF4Tensor" \
+        "from torchao.dtypes.nf4tensor import NF4Tensor"
+
+    substituteInPlace \
+      tests/torchtune/modules/peft/test_dora.py \
+      tests/torchtune/modules/peft/test_lora.py \
+      torchtune/training/_distributed.py \
+      --replace-fail \
+        "from torchao.quantization import NF4Tensor, to_nf4" \
+        "from torchao.dtypes.nf4tensor import NF4Tensor, to_nf4"
+
+    substituteInPlace \
+      torchtune/modules/low_precision/nf4_linear.py \
+      torchtune/modules/peft/dora.py \
+      torchtune/modules/peft/lora.py \
+      --replace-fail \
+        "from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4" \
+        "from torchao.dtypes.nf4tensor import linear_nf4"
+
+    substituteInPlace torchtune/modules/low_precision/_register_nf4_dispatch_ops.py \
+      --replace-fail \
+        "from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import implements as nf4_tensor_impl" \
+        "from torchao.dtypes.nf4tensor import implements as nf4_tensor_impl"
+  '';
+
   build-system = [
     setuptools
   ];
 
+  pythonRelaxDeps = [
+    "pyarrow"
+  ];
   dependencies = [
     blobfile
     datasets
@@ -77,10 +130,12 @@ buildPythonPackage (finalAttrs: {
 
   nativeCheckInputs = [
     comet-ml
+    expecttest
     mlflow
     pytest-integration
     pytest-mock
     pytestCheckHook
+    tensorboard
     writableTmpDirAsHomeHook
   ];
 
@@ -157,6 +212,10 @@ buildPythonPackage (finalAttrs: {
     "tests/torchtune/utils"
   ];
 
+  passthru.updateScript = nix-update-script {
+    extraArgs = [ "--version=branch" ];
+  };
+
   meta = {
     description = "PyTorch native post-training library";
     homepage = "https://github.com/meta-pytorch/torchtune";