python3Packages.unsloth: init at 2025.4.1 (#388735)

maintainers/maintainer-list.nix

@@ -9903,6 +9903,12 @@
     githubId = 57007241;
     name = "hogcycle";
   };
+  hoh = {
+    email = "git@hugoherter.com";
+    github = "hoh";
+    githubId = 404665;
+    name = "Hugo Herter";
+  };
   holgerpeters = {
     name = "Holger Peters";
     email = "holger.peters@posteo.de";

pkgs/development/python-modules/cut-cross-entropy/default.nix

@@ -0,0 +1,83 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+  setuptools-scm,
+
+  # dependencies
+  torch,
+  triton,
+
+  # optional-dependencies
+  accelerate,
+  datasets,
+  fire,
+  huggingface-hub,
+  pandas,
+  pytestCheckHook,
+  pythonAtLeast,
+  tqdm,
+  transformers,
+}:
+
+buildPythonPackage {
+  pname = "cut-cross-entropy";
+  version = "25.3.1";
+  pyproject = true;
+
+  # The `ml-cross-entropy` Pypi comes from a third-party.
+  # Apple recommends installing from the repo's main branch directly
+  src = fetchFromGitHub {
+    owner = "apple";
+    repo = "ml-cross-entropy";
+    rev = "24fbe4b5dab9a6c250a014573613c1890190536c"; # no tags
+    hash = "sha256-BVPon+T7chkpozX/IZU3KZMw1zRzlYVvF/22JWKjT2Y=";
+  };
+
+  # Python 3.13 support requires PyTorch 2.6, which is not merged into master yet
+  # https://github.com/NixOS/nixpkgs/pull/377785
+  disabled = pythonAtLeast "3.13";
+
+  build-system = [
+    setuptools
+    setuptools-scm
+  ];
+
+  dependencies = [
+    torch
+    triton
+  ];
+
+  optional-dependencies = {
+    transformers = [ transformers ];
+    all = [
+      accelerate
+      datasets
+      fire
+      huggingface-hub
+      pandas
+      tqdm
+      transformers
+    ];
+    # `deepspeed` is not yet packaged in nixpkgs
+    # ++ lib.optionals (!stdenv.isDarwin) [
+    #   deepspeed
+    # ];
+  };
+
+  nativeCheckInputs = [ pytestCheckHook ];
+
+  pythonImportsCheck = [
+    "cut_cross_entropy"
+  ];
+
+  meta = {
+    description = "Memory-efficient cross-entropy loss implementation using Cut Cross-Entropy (CCE)";
+    homepage = "https://github.com/apple/ml-cross-entropy";
+    license = lib.licenses.aml;
+    maintainers = with lib.maintainers; [ hoh ];
+  };
+}

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

@@ -0,0 +1,53 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+  setuptools-scm,
+
+  # dependencies
+  accelerate,
+  datasets,
+  rich,
+  transformers,
+}:
+
+buildPythonPackage rec {
+  pname = "trl";
+  version = "0.15.2";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "huggingface";
+    repo = "trl";
+    tag = "v${version}";
+    hash = "sha256-HsSmFXFqDOWVLa6VXdPZVS9C3bjYcsliR0TwNpPiQx4=";
+  };
+
+  build-system = [
+    setuptools
+    setuptools-scm
+  ];
+
+  dependencies = [
+    accelerate
+    datasets
+    rich
+    transformers
+  ];
+
+  # Many tests require internet access.
+  doCheck = false;
+
+  pythonImportsCheck = [ "trl" ];
+
+  meta = {
+    description = "Train transformer language models with reinforcement learning";
+    homepage = "https://github.com/huggingface/trl";
+    changelog = "https://github.com/huggingface/trl/releases/tag/v${version}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ hoh ];
+  };
+}

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

@@ -0,0 +1,69 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  hatchling,
+
+  # dependencies
+  docstring-parser,
+  rich,
+  shtab,
+  typeguard,
+  typing-extensions,
+
+  # tests
+  attrs,
+  flax,
+  jax,
+  ml-collections,
+  omegaconf,
+  pydantic,
+  pytestCheckHook,
+  torch,
+}:
+
+buildPythonPackage rec {
+  pname = "tyro";
+  version = "0.9.19";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "brentyi";
+    repo = "tyro";
+    tag = "v${version}";
+    hash = "sha256-A1Vplc84Xy8TufqmklPUzIdgiPpFcIjqV0eUgdKmYRM=";
+  };
+
+  build-system = [ hatchling ];
+
+  dependencies = [
+    docstring-parser
+    rich
+    shtab
+    typeguard
+    typing-extensions
+  ];
+
+  nativeCheckInputs = [
+    attrs
+    flax
+    jax
+    ml-collections
+    omegaconf
+    pydantic
+    pytestCheckHook
+    torch
+  ];
+
+  pythonImportsCheck = [ "tyro" ];
+
+  meta = {
+    description = "CLI interfaces & config objects, from types";
+    homepage = "https://github.com/brentyi/tyro";
+    changelog = "https://github.com/brentyi/tyro/releases/tag/v${version}";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ hoh ];
+  };
+}

pkgs/development/python-modules/unsloth-zoo/default.nix

@@ -0,0 +1,86 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchPypi,
+
+  # build-system
+  setuptools,
+  setuptools-scm,
+
+  # dependencies
+  accelerate,
+  cut-cross-entropy,
+  datasets,
+  hf-transfer,
+  huggingface-hub,
+  packaging,
+  peft,
+  psutil,
+  sentencepiece,
+  torch,
+  tqdm,
+  transformers,
+  trl,
+  tyro,
+}:
+
+buildPythonPackage rec {
+  pname = "unsloth-zoo";
+  version = "2025.4.1";
+  pyproject = true;
+
+  # no tags on GitHub
+  src = fetchPypi {
+    pname = "unsloth_zoo";
+    inherit version;
+    hash = "sha256-mRs/NMCNJWT52S7mtbQI332IQR6+/IaL29XmtMOz3fE=";
+  };
+
+  # pyproject.toml requires an obsolete version of protobuf,
+  # but it is not used.
+  # Upstream issue: https://github.com/unslothai/unsloth-zoo/pull/68
+  pythonRelaxDeps = [
+    "protobuf"
+  ];
+
+  patches = [
+    # Avoid circular dependency in Nix, since `unsloth` depends on `unsloth-zoo`.
+    ./dont-require-unsloth.patch
+  ];
+
+  build-system = [
+    setuptools
+    setuptools-scm
+  ];
+
+  dependencies = [
+    accelerate
+    cut-cross-entropy
+    datasets
+    hf-transfer
+    huggingface-hub
+    packaging
+    peft
+    psutil
+    sentencepiece
+    torch
+    tqdm
+    transformers
+    trl
+    tyro
+  ];
+
+  # No tests
+  doCheck = false;
+
+  pythonImportsCheck = [
+    "unsloth_zoo"
+  ];
+
+  meta = {
+    description = "Utils for Unsloth";
+    homepage = "https://github.com/unslothai/unsloth_zoo";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ hoh ];
+  };
+}

pkgs/development/python-modules/unsloth-zoo/dont-require-unsloth.patch

@@ -0,0 +1,18 @@
+diff --git a/unsloth_zoo/__init__.py b/unsloth_zoo/__init__.py
+--- a/unsloth_zoo/__init__.py
++++ b/unsloth_zoo/__init__.py
+@@ -17,14 +17,10 @@
+ __version__ = "2025.3.17"
+ 
+ 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
+ 
+ import os
+-if not ("UNSLOTH_IS_PRESENT" in os.environ):
+-    raise ImportError("Please install Unsloth via `pip install unsloth`!")
+ pass
+ 
+ try:

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

@@ -0,0 +1,92 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchPypi,
+
+  # build-system
+  setuptools,
+  setuptools-scm,
+
+  # dependencies
+  bitsandbytes,
+  numpy,
+  packaging,
+  torch,
+  unsloth-zoo,
+  xformers,
+  tyro,
+  transformers,
+  datasets,
+  sentencepiece,
+  tqdm,
+  accelerate,
+  trl,
+  peft,
+  protobuf,
+  huggingface-hub,
+  hf-transfer,
+  diffusers,
+  torchvision,
+}:
+
+buildPythonPackage rec {
+  pname = "unsloth";
+  version = "2025.4.1";
+  pyproject = true;
+
+  # Tags on the GitHub repo don't match
+  src = fetchPypi {
+    pname = "unsloth";
+    inherit version;
+    hash = "sha256-9LtDGfdWH7R3U/xi+aK3V4zA+vs83S6Cp0F2NQKvSdY=";
+  };
+
+  build-system = [
+    setuptools
+    setuptools-scm
+  ];
+
+  dependencies = [
+    bitsandbytes
+    numpy
+    packaging
+    torch
+    unsloth-zoo
+    xformers
+    tyro
+    transformers
+    datasets
+    sentencepiece
+    tqdm
+    accelerate
+    trl
+    peft
+    protobuf
+    huggingface-hub
+    hf-transfer
+    diffusers
+    torchvision
+  ];
+
+  # pyproject.toml requires an obsolete version of protobuf,
+  # but it is not used.
+  # Upstream issue: https://github.com/unslothai/unsloth-zoo/pull/68
+  pythonRelaxDeps = [
+    "protobuf"
+  ];
+
+  # The source repository contains no test
+  doCheck = false;
+
+  # Importing requires a GPU, else the following error is raised:
+  # NotImplementedError: Unsloth: No NVIDIA GPU found? Unsloth currently only supports GPUs!
+  dontUsePythonImportsCheck = true;
+
+  meta = {
+    description = "Finetune Llama 3.3, DeepSeek-R1 & Reasoning LLMs 2x faster with 70% less memory";
+    homepage = "https://github.com/unslothai/unsloth";
+    changelog = "https://github.com/unslothai/unsloth/releases/tag/${version}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ hoh ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -3042,6 +3042,8 @@ self: super: with self; {
 
   customtkinter = callPackage ../development/python-modules/customtkinter { };
 
+  cut-cross-entropy = callPackage ../development/python-modules/cut-cross-entropy { };
+
   cvelib = callPackage ../development/python-modules/cvelib { };
 
   cvss = callPackage ../development/python-modules/cvss { };
@@ -17665,6 +17667,8 @@ self: super: with self; {
 
   tritonclient = callPackage ../development/python-modules/tritonclient { };
 
+  trl = callPackage ../development/python-modules/trl { };
+
   trlib = toPythonModule (
     pkgs.trlib.override {
       pythonSupport = true;
@@ -18273,6 +18277,8 @@ self: super: with self; {
 
   typst = callPackage ../development/python-modules/typst { };
 
+  tyro = callPackage ../development/python-modules/tyro { };
+
   tzdata = callPackage ../development/python-modules/tzdata { };
 
   tzlocal = callPackage ../development/python-modules/tzlocal { };
@@ -18427,6 +18433,10 @@ self: super: with self; {
 
   unrpa = callPackage ../development/python-modules/unrpa { };
 
+  unsloth = callPackage ../development/python-modules/unsloth { };
+
+  unsloth-zoo = callPackage ../development/python-modules/unsloth-zoo { };
+
   unstructured = callPackage ../development/python-modules/unstructured { };
 
   unstructured-api-tools = callPackage ../development/python-modules/unstructured-api-tools { };