python3Packages.unsloth: add CUDA tests using GPU (#441728)

This commit is contained in:
Gaétan Lepage
2025-10-11 21:29:28 +00:00
committed by GitHub
@@ -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";