instructlab · fabiendupont · Oct 14, 2024 · Maxusmusti · Oct 29, 2024 · fabiendupont
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,8 +44,11 @@ package-dir = { "" = "src" }
 
 [tool.setuptools.dynamic]
 dependencies = { file = ["requirements.txt"] }
-optional-dependencies.cuda = { file = ["requirements-cuda.txt"] }
-optional-dependencies.rocm = { file = ["requirements-rocm.txt"] }
+optional-dependencies.cpu = { file = ["requirements/cpu.txt"] }
+optional-dependencies.cuda = { file = ["requirements/cuda.txt"] }
+optional-dependencies.hpu = { file = ["requirements/hpu.txt"] }
+optional-dependencies.mps = { file = ["requirements/mps.txt"] }
+optional-dependencies.rocm = { file = ["requirements/rocm.txt"] }
 
 [tool.setuptools.packages.find]
 where = ["src"]

diff --git a/requirements-rocm.txt b/requirements-rocm.txt
diff --git a/requirements.txt b/requirements.txt
@@ -4,7 +4,7 @@ pyyaml
 py-cpuinfo
 # we set this to be above 0a0 so that it doesn't
 # replace custom pytorch images with the 2.3.0
-torch>=2.3.0a0
+torch>=2.3.0,<2.5.0
 transformers>=4.45.2
 accelerate>=0.34.2
 datasets>=2.15.0

diff --git a/requirements/cpu.txt b/requirements/cpu.txt
@@ -0,0 +1 @@
+# Extra dependencies for CPU-only
diff --git a/requirements-cuda.txt → requirements/cuda.txt b/requirements-cuda.txt → requirements/cuda.txt
@@ -1,2 +1,3 @@
+# Extra dependencies for NVIDIA CUDA
 flash-attn>=2.4.0
 bitsandbytes>=0.43.1
diff --git a/requirements/hpu.txt b/requirements/hpu.txt
@@ -0,0 +1,11 @@
+# # Extra dependencies for Intel Gaudi / Habana Labs HPU devices
+# Habana Labs 1.17.1 has PyTorch 2.3.1a0+gitxxx pre-release
+torch>=2.3.1a0,<2.4.0
+# Habana Labs frameworks
+habana-torch-plugin>=1.17.1
+habana_gpu_migration>=1.17.1
+# additional Habana Labs packages (installed, but not used)
+#habana-media-loader
+#habana-pyhlml
+#habana_quantization_toolkit
+#habana-torch-dataloader
diff --git a/requirements/mps.txt b/requirements/mps.txt
@@ -0,0 +1 @@
+# Extra dependencies for Apple MPS (Metal Performance Shaders)
diff --git a/requirements/rocm.txt b/requirements/rocm.txt
@@ -0,0 +1,2 @@
+# Extra dependencies for AMD ROCm
+flash-attn>=2.6.2,<2.7.0
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		# Extra dependencies for Apple MPS (Metal Performance Shaders)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		# Extra dependencies for AMD ROCm
		flash-attn>=2.6.2,<2.7.0
Copy link Contributor tiran Oct 30, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. What's the reason to limit the upper bound of `flash-attn`?