diff --git a/pkgs/development/python-modules/finetuning-scheduler/default.nix b/pkgs/development/python-modules/finetuning-scheduler/default.nix new file mode 100644 index 000000000000..d4a2752c8ea3 --- /dev/null +++ b/pkgs/development/python-modules/finetuning-scheduler/default.nix @@ -0,0 +1,55 @@ +{ + stdenv, + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + pythonOlder, + pytestCheckHook, + torch, + pytorch-lightning, +}: + +buildPythonPackage rec { + pname = "finetuning-scheduler"; + version = "2.3.3"; + pyproject = true; + + disabled = pythonOlder "3.8"; + + src = fetchFromGitHub { + owner = "speediedan"; + repo = "finetuning-scheduler"; + rev = "refs/tags/v${version}"; + hash = "sha256-Nci+NcfcaZV3NH3MTCyr8PwPdDykJTdq8+CmtngWdiM="; + }; + + build-system = [ setuptools ]; + + dependencies = [ + pytorch-lightning + torch + ]; + + # needed while lightning is installed as package `pytorch-lightning` rather than`lightning`: + env.PACKAGE_NAME = "pytorch"; + + nativeCheckInputs = [ pytestCheckHook ]; + pytestFlagsArray = [ "tests" ]; + disabledTests = lib.optionals (stdenv.isAarch64 && stdenv.isLinux) [ + # slightly exceeds numerical tolerance on aarch64-linux: + "test_fts_frozen_bn_track_running_stats" + ]; + + pythonImportsCheck = [ "finetuning_scheduler" ]; + + meta = { + description = "PyTorch Lightning extension for foundation model experimentation with flexible fine-tuning schedules"; + homepage = "https://finetuning-scheduler.readthedocs.io"; + changelog = "https://github.com/speediedan/finetuning-scheduler/blob/${src.rev}/CHANGELOG.md"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ bcdarwin ]; + # "No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package" at import time: + broken = stdenv.isDarwin; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 68047d862928..35b2babc7e28 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -4290,6 +4290,8 @@ self: super: with self; { findpython = callPackage ../development/python-modules/findpython { }; + finetuning-scheduler = callPackage ../development/python-modules/finetuning-scheduler { }; + fingerprints = callPackage ../development/python-modules/fingerprints { }; finitude = callPackage ../development/python-modules/finitude { };