tjh.dev
1{
2 stdenv,
3 lib,
4 buildPythonPackage,
5 fetchFromGitHub,
6 setuptools,
7 pythonOlder,
8 pytestCheckHook,
9 torch,
10 pytorch-lightning,
11}:
12
13buildPythonPackage rec {
14 pname = "finetuning-scheduler";
15 version = "2.4.0";
16 pyproject = true;
17
18 disabled = pythonOlder "3.9";
19
20 src = fetchFromGitHub {
21 owner = "speediedan";
22 repo = "finetuning-scheduler";
23 rev = "refs/tags/v${version}";
24 hash = "sha256-uSFGZseSJv519LpaddO6yP6AsIMZutEA0Y7Yr+mEWTQ=";
25 };
26
27 build-system = [ setuptools ];
28
29 dependencies = [
30 pytorch-lightning
31 torch
32 ];
33
34 # needed while lightning is installed as package `pytorch-lightning` rather than`lightning`:
35 env.PACKAGE_NAME = "pytorch";
36
37 nativeCheckInputs = [ pytestCheckHook ];
38 pytestFlagsArray = [ "tests" ];
39 disabledTests =
40 # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
41 # LoweringException: ImportError: cannot import name 'triton_key' from 'triton.compiler.compiler'
42 lib.optionals (pythonOlder "3.12") [
43 "test_fts_dynamo_enforce_p0"
44 "test_fts_dynamo_resume"
45 "test_fts_dynamo_intrafit"
46 ]
47 ++ lib.optionals (stdenv.hostPlatform.isAarch64 && stdenv.hostPlatform.isLinux) [
48 # slightly exceeds numerical tolerance on aarch64-linux:
49 "test_fts_frozen_bn_track_running_stats"
50 ];
51
52 pythonImportsCheck = [ "finetuning_scheduler" ];
53
54 meta = {
55 description = "PyTorch Lightning extension for foundation model experimentation with flexible fine-tuning schedules";
56 homepage = "https://finetuning-scheduler.readthedocs.io";
57 changelog = "https://github.com/speediedan/finetuning-scheduler/blob/${src.rev}/CHANGELOG.md";
58 license = lib.licenses.asl20;
59 maintainers = with lib.maintainers; [ bcdarwin ];
60 # "No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package" at import time:
61 broken = stdenv.hostPlatform.isDarwin;
62 };
63}