From 6dc0eb03c0209ae340f0fd2b26fa6ab491e50fcd Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Fri, 20 Mar 2026 16:52:29 +0000
Subject: [PATCH] python3Packages.unsloth-zoo: 2026.1.3 -> 2026.3.4

---
 .../python-modules/unsloth-zoo/default.nix    | 25 ++++++++-----
 .../unsloth-zoo/dont-require-unsloth.patch    | 36 +++++++++++--------
 2 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/pkgs/development/python-modules/unsloth-zoo/default.nix b/pkgs/development/python-modules/unsloth-zoo/default.nix
index 8549a8c0ee5c..ef2be47691bf 100644
--- a/pkgs/development/python-modules/unsloth-zoo/default.nix
+++ b/pkgs/development/python-modules/unsloth-zoo/default.nix
@@ -25,18 +25,28 @@
   tyro,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "unsloth-zoo";
-  version = "2026.1.3";
+  version = "2026.3.4";
   pyproject = true;
 
   # no tags on GitHub
   src = fetchPypi {
     pname = "unsloth_zoo";
-    inherit version;
-    hash = "sha256-njArRI9y6EKwZamJnQoWQIvmYzV6t0syUIi1d7DzX4U=";
+    inherit (finalAttrs) version;
+    hash = "sha256-24w8UV5cLG5XWrU73xfmg+Jk32zl+QSPdqXLzMtmP1E=";
   };
 
+  postPatch = ''
+    substituteInPlace pyproject.toml \
+      --replace-fail \
+        "setuptools==80.9.0" \
+        "setuptools" \
+      --replace-fail \
+        "setuptools-scm==9.2.0" \
+        "setuptools-scm"
+  '';
+
   # pyproject.toml requires an obsolete version of protobuf,
   # but it is not used.
   # Upstream issue: https://github.com/unslothai/unsloth-zoo/pull/68
@@ -44,6 +54,7 @@ buildPythonPackage rec {
     "datasets"
     "protobuf"
     "transformers"
+    "trl"
     "torch"
   ];
 
@@ -78,9 +89,7 @@ buildPythonPackage rec {
   # No tests
   doCheck = false;
 
-  pythonImportsCheck = [
-    "unsloth_zoo"
-  ];
+  pythonImportsCheck = [ "unsloth_zoo" ];
 
   meta = {
     description = "Utils for Unsloth";
@@ -88,4 +97,4 @@ buildPythonPackage rec {
     license = lib.licenses.mit;
     maintainers = with lib.maintainers; [ hoh ];
   };
-}
+})
diff --git a/pkgs/development/python-modules/unsloth-zoo/dont-require-unsloth.patch b/pkgs/development/python-modules/unsloth-zoo/dont-require-unsloth.patch
index f1f2990d318f..b5f7c9d8ddd7 100644
--- a/pkgs/development/python-modules/unsloth-zoo/dont-require-unsloth.patch
+++ b/pkgs/development/python-modules/unsloth-zoo/dont-require-unsloth.patch
@@ -1,29 +1,37 @@
 diff --git a/unsloth_zoo/__init__.py b/unsloth_zoo/__init__.py
-index 2332907..2eed5e9 100644
+index 8c4404c..dc1666f 100644
 --- a/unsloth_zoo/__init__.py
 +++ b/unsloth_zoo/__init__.py
-@@ -64,8 +64,6 @@ pass
+@@ -90,8 +90,6 @@ if (os.environ.get("UNSLOTH_COMPILE_DEBUG", "0") == "1"):
  
  
  from importlib.util import find_spec
 -if find_spec("unsloth") is None:
 -    raise ImportError("Please install Unsloth via `pip install unsloth`!")
- pass
- del find_spec
+ if find_spec("torch") is None:
+     raise ImportError(
+         "Unsloth: Pytorch is not installed. Go to https://pytorch.org/.\n"\
+@@ -275,9 +273,6 @@ del remove_expandable_segments, delete_key, IS_HIP_RUNTIME, IS_TORCH_2_9_OR_NEWE
+ del clean_expandable_segments_value
+ del _ORIGINAL_PYTORCH_CUDA_ALLOC_CONF, _ORIGINAL_PYTORCH_HIP_ALLOC_CONF, _HAS_ORIGINAL_PYTORCH_ALLOC_CONF
  
-@@ -75,12 +73,13 @@ def get_device_type():
+-if not ("UNSLOTH_IS_PRESENT" in os.environ):
+-    raise ImportError("Please install Unsloth via `pip install unsloth`!")
+-
+ try:
+     print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.")
+ except:
+diff --git a/unsloth_zoo/device_type.py b/unsloth_zoo/device_type.py
+index 11136fb..8f8dafc 100644
+--- a/unsloth_zoo/device_type.py
++++ b/unsloth_zoo/device_type.py
+@@ -209,6 +209,9 @@ def get_device_type():
          return "cuda"
      elif hasattr(torch, "xpu") and torch.xpu.is_available():
          return "xpu"
 +    else:
 +        # Allow import during tests
 +        return None
-     raise NotImplementedError("Unsloth currently only works on NVIDIA GPUs and Intel GPUs.")
- pass
- DEVICE_TYPE : str = get_device_type()
- 
--if not ("UNSLOTH_IS_PRESENT" in os.environ):
--    raise ImportError("Please install Unsloth via `pip install unsloth`!")
- pass
- 
- try:
+     # Check torch.accelerator
+     if hasattr(torch, "accelerator"):
+         if not torch.accelerator.is_available():