Commit 88ff1cfb authored by Hugo Herter's avatar Hugo Herter
Browse files

python3Packages.unsloth: Add test using CUDA 



Adds a test from the Unsloth repository that runs QLora training
and merging, ensuring basic functionality of the library.

Co-authored-by: default avatarSomeone <else@someonex.net>
parent daebeba7

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

+52 −0
Original line number Diff line number Diff line
@@ -27,8 +27,30 @@ 
  hf-transfer,

  diffusers,

  torchvision,



  # tests

  fetchFromGitHub,

  cudaPackages,

  python,

  gcc,

}:



let

  # Test files are absent from the PyPI package, so we fetch them separately.

  testSrc = fetchFromGitHub {

    owner = "unslothai";

    repo = "unsloth";

    rev = "cb78f0e83dc2d61fb1571b6e904eb2f064510d63";

    hash = "sha256-0oR3m8jnjSdfjH+NslW6SsVj+0cQ4VUhKXZ38U/VBy0=";

    # Keep only the tests directory, use the PyPI package for everything else.

    postFetch = ''

      mv $out/tests $TMPDIR/tests

      rm -rf $out/*

      mv $TMPDIR/tests $out/tests

    '';

  };

in



buildPythonPackage rec {

  pname = "unsloth";

  version = "2025.9.4";
@@ -85,6 +107,36 @@ buildPythonPackage rec { 
  # NotImplementedError: Unsloth: No NVIDIA GPU found? Unsloth currently only supports GPUs!

  dontUsePythonImportsCheck = true;



  passthru.tests = {

    qlora-train-and-merge =

      # FIXME: Replace python3.pkgs with python3Packages once possible, as to unbeak splicing.

      # Cf. https://github.com/NixOS/nixpkgs/pull/394838#issuecomment-3319287038

      cudaPackages.writeGpuTestPython.override { python3Packages = python.pkgs; }

        {

          libraries = ps: [

            ps.unsloth

            ps.unsloth-zoo

          ];

          # In-derivation test would require fetching the model from hugging-face which will not be trivial.

          gpuCheckArgs.meta.broken = true;

        }

        # Triton JIT requires a C compiler at runtime and imports files from the test directory.

        ''

          import os, sys, runpy



          os.environ["CC"]  = "${lib.getExe' gcc "cc"}";

          os.environ["CXX"] = "${lib.getExe' gcc "cxx"}";



          sys.path.insert(0, "${testSrc}")



          # Execute the test file from the Nix store at runtime (no eval-time IFD).

          runpy.run_path(

              "${testSrc}/tests/qlora/test_unsloth_qlora_train_and_merge.py",

              run_name="__main__"

          )

        '';

  };



  meta = {

    description = "Finetune Llama 3.3, DeepSeek-R1 & Reasoning LLMs 2x faster with 70% less memory";

    homepage = "https://github.com/unslothai/unsloth";