pyrox.dev
1{
2 stdenv,
3 lib,
4 buildPythonPackage,
5 fetchFromGitHub,
6 pythonAtLeast,
7
8 # buildInputs
9 llvmPackages,
10
11 # build-system
12 setuptools,
13
14 # dependencies
15 huggingface-hub,
16 numpy,
17 packaging,
18 psutil,
19 pyyaml,
20 safetensors,
21 torch,
22
23 # tests
24 addBinToPathHook,
25 evaluate,
26 parameterized,
27 pytestCheckHook,
28 transformers,
29 config,
30 cudatoolkit,
31 writableTmpDirAsHomeHook,
32}:
33
34buildPythonPackage rec {
35 pname = "accelerate";
36 version = "1.10.1";
37 pyproject = true;
38
39 src = fetchFromGitHub {
40 owner = "huggingface";
41 repo = "accelerate";
42 tag = "v${version}";
43 hash = "sha256-guUlgNRSgBmi5HZNnwPA6GQ1vq5LzZEW2J9w4Wmqh6o=";
44 };
45
46 buildInputs = [ llvmPackages.openmp ];
47
48 build-system = [ setuptools ];
49
50 dependencies = [
51 huggingface-hub
52 numpy
53 packaging
54 psutil
55 pyyaml
56 safetensors
57 torch
58 ];
59
60 nativeCheckInputs = [
61 addBinToPathHook
62 evaluate
63 parameterized
64 pytestCheckHook
65 transformers
66 writableTmpDirAsHomeHook
67 ];
68
69 preCheck = lib.optionalString config.cudaSupport ''
70 export TRITON_PTXAS_PATH="${lib.getExe' cudatoolkit "ptxas"}"
71 '';
72 enabledTestPaths = [ "tests" ];
73 disabledTests = [
74 # try to download data:
75 "FeatureExamplesTests"
76 "test_infer_auto_device_map_on_t0pp"
77
78 # require socket communication
79 "test_explicit_dtypes"
80 "test_gated"
81 "test_invalid_model_name"
82 "test_invalid_model_name_transformers"
83 "test_no_metadata"
84 "test_no_split_modules"
85 "test_remote_code"
86 "test_transformers_model"
87 "test_extract_model_keep_torch_compile"
88 "test_extract_model_remove_torch_compile"
89 "test_regions_are_compiled"
90
91 # nondeterministic, tests GC behaviour by thresholding global ram usage
92 "test_free_memory_dereferences_prepared_components"
93
94 # set the environment variable, CC, which conflicts with standard environment
95 "test_patch_environment_key_exists"
96 ]
97 ++ lib.optionals ((pythonAtLeast "3.13") || (torch.rocmSupport or false)) [
98 # RuntimeError: Dynamo is not supported on Python 3.13+
99 # OR torch.compile tests broken on torch 2.5 + rocm
100 "test_can_unwrap_distributed_compiled_model_keep_torch_compile"
101 "test_can_unwrap_distributed_compiled_model_remove_torch_compile"
102 "test_convert_to_fp32"
103 "test_send_to_device_compiles"
104 ]
105 ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
106 # usual aarch64-linux RuntimeError: DataLoader worker (pid(s) <...>) exited unexpectedly
107 "CheckpointTest"
108 # TypeError: unsupported operand type(s) for /: 'NoneType' and 'int' (it seems cpuinfo doesn't work here)
109 "test_mpi_multicpu_config_cmd"
110 ]
111 ++ lib.optionals (!config.cudaSupport) [
112 # requires ptxas from cudatoolkit, which is unfree
113 "test_dynamo_extract_model"
114 ]
115 ++ lib.optionals stdenv.hostPlatform.isDarwin [
116 # RuntimeError: 'accelerate-launch /nix/store/a7vhm7b74a7bmxc35j26s9iy1zfaqjs...
117 "test_accelerate_test"
118 "test_init_trackers"
119 "test_init_trackers"
120 "test_log"
121 "test_log_with_tensor"
122
123 # After enabling MPS in pytorch, these tests started failing
124 "test_accelerated_optimizer_step_was_skipped"
125 "test_auto_wrap_policy"
126 "test_autocast_kwargs"
127 "test_automatic_loading"
128 "test_backward_prefetch"
129 "test_can_resume_training"
130 "test_can_resume_training_checkpoints_relative_path"
131 "test_can_resume_training_with_folder"
132 "test_can_unwrap_model_fp16"
133 "test_checkpoint_deletion"
134 "test_cpu_offload"
135 "test_cpu_ram_efficient_loading"
136 "test_grad_scaler_kwargs"
137 "test_invalid_registration"
138 "test_map_location"
139 "test_mixed_precision"
140 "test_mixed_precision_buffer_autocast_override"
141 "test_project_dir"
142 "test_project_dir_with_config"
143 "test_sharding_strategy"
144 "test_state_dict_type"
145 "test_with_save_limit"
146 "test_with_scheduler"
147
148 # torch._inductor.exc.InductorError: TypeError: cannot determine truth value of Relational
149 "test_regional_compilation_cold_start"
150 "test_regional_compilation_inference_speedup"
151
152 # Fails in nixpkgs-review due to a port conflict with simultaneous python builds
153 "test_config_compatibility"
154
155 # Fails with `sandbox=false` by mis-configuring the model it's using.
156 # AttributeError: 'DistributedDataParallel' object has no attribute '_ignored_modules'. Did you mean: 'named_modules'?
157 "test_ignored_modules_regex"
158 ]
159 ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [
160 # RuntimeError: torch_shm_manager: execl failed: Permission denied
161 "CheckpointTest"
162 ];
163
164 disabledTestPaths = lib.optionals (!(stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64)) [
165 # numerous instances of torch.multiprocessing.spawn.ProcessRaisedException:
166 "tests/test_cpu.py"
167 "tests/test_grad_sync.py"
168 "tests/test_metrics.py"
169 "tests/test_scheduler.py"
170 ];
171
172 pythonImportsCheck = [ "accelerate" ];
173
174 __darwinAllowLocalNetworking = true;
175
176 meta = {
177 homepage = "https://huggingface.co/docs/accelerate";
178 description = "Simple way to train and use PyTorch models with multi-GPU, TPU, mixed-precision";
179 changelog = "https://github.com/huggingface/accelerate/releases/tag/${src.tag}";
180 license = lib.licenses.asl20;
181 maintainers = with lib.maintainers; [ bcdarwin ];
182 mainProgram = "accelerate";
183 };
184}