tjh.dev
1{
2 stdenv,
3 lib,
4 buildPythonPackage,
5 fetchFromGitHub,
6 pythonAtLeast,
7 pythonOlder,
8 llvmPackages,
9 pytest7CheckHook,
10 setuptools,
11 numpy,
12 packaging,
13 psutil,
14 pyyaml,
15 safetensors,
16 torch,
17 config,
18 cudatoolkit,
19 evaluate,
20 parameterized,
21 transformers,
22}:
23
24buildPythonPackage rec {
25 pname = "accelerate";
26 version = "0.32.0";
27 pyproject = true;
28
29 disabled = pythonOlder "3.8";
30
31 src = fetchFromGitHub {
32 owner = "huggingface";
33 repo = "accelerate";
34 rev = "refs/tags/v${version}";
35 hash = "sha256-/Is5aKTYHxvgUJSkF7HxMbEA6dgn/y5F1B3D6qSCSaE=";
36 };
37
38 buildInputs = [ llvmPackages.openmp ];
39
40 build-system = [ setuptools ];
41
42 dependencies = [
43 numpy
44 packaging
45 psutil
46 pyyaml
47 safetensors
48 torch
49 ];
50
51 nativeCheckInputs = [
52 evaluate
53 parameterized
54 pytest7CheckHook
55 transformers
56 ];
57 preCheck =
58 ''
59 export HOME=$(mktemp -d)
60 export PATH=$out/bin:$PATH
61 ''
62 + lib.optionalString config.cudaSupport ''
63 export TRITON_PTXAS_PATH="${cudatoolkit}/bin/ptxas"
64 '';
65 pytestFlagsArray = [ "tests" ];
66 disabledTests =
67 [
68 # try to download data:
69 "FeatureExamplesTests"
70 "test_infer_auto_device_map_on_t0pp"
71
72 # require socket communication
73 "test_explicit_dtypes"
74 "test_gated"
75 "test_invalid_model_name"
76 "test_invalid_model_name_transformers"
77 "test_no_metadata"
78 "test_no_split_modules"
79 "test_remote_code"
80 "test_transformers_model"
81
82 # nondeterministic, tests GC behaviour by thresholding global ram usage
83 "test_free_memory_dereferences_prepared_components"
84
85 # set the environment variable, CC, which conflicts with standard environment
86 "test_patch_environment_key_exists"
87 ]
88 ++ lib.optionals (pythonAtLeast "3.12") [
89 # RuntimeError: Dynamo is not supported on Python 3.12+
90 "test_convert_to_fp32"
91 "test_dynamo_extract_model"
92 "test_send_to_device_compiles"
93 ]
94 ++ lib.optionals (stdenv.isLinux && stdenv.isAarch64) [
95 # usual aarch64-linux RuntimeError: DataLoader worker (pid(s) <...>) exited unexpectedly
96 "CheckpointTest"
97 # TypeError: unsupported operand type(s) for /: 'NoneType' and 'int' (it seems cpuinfo doesn't work here)
98 "test_mpi_multicpu_config_cmd"
99 ]
100 ++ lib.optionals (!config.cudaSupport) [
101 # requires ptxas from cudatoolkit, which is unfree
102 "test_dynamo_extract_model"
103 ]
104 ++ lib.optionals (stdenv.isDarwin && stdenv.isx86_64) [
105 # RuntimeError: torch_shm_manager: execl failed: Permission denied
106 "CheckpointTest"
107 ];
108
109 disabledTestPaths = lib.optionals (!(stdenv.isLinux && stdenv.isx86_64)) [
110 # numerous instances of torch.multiprocessing.spawn.ProcessRaisedException:
111 "tests/test_cpu.py"
112 "tests/test_grad_sync.py"
113 "tests/test_metrics.py"
114 "tests/test_scheduler.py"
115 ];
116
117 pythonImportsCheck = [ "accelerate" ];
118
119 __darwinAllowLocalNetworking = true;
120
121 meta = with lib; {
122 homepage = "https://huggingface.co/docs/accelerate";
123 description = "Simple way to train and use PyTorch models with multi-GPU, TPU, mixed-precision";
124 changelog = "https://github.com/huggingface/accelerate/releases/tag/v${version}";
125 license = licenses.asl20;
126 maintainers = with maintainers; [ bcdarwin ];
127 mainProgram = "accelerate";
128 };
129}