Update PyTorch to 2.4.1

This change the default PyTorch 2.4.1.
Intel Gaudi (HPU) remains on Intel fork of PyTorch 2.3.1.

Fixes instructlab/instructlab#2437

Signed-off-by: Fabien Dupont <[email protected]>

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"]

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.41.2
 accelerate>=0.34.2
 datasets>=2.15.0

requirements/cpu.txt

@@ -0,0 +1 @@
+# Extra dependencies for CPU-only

requirements-cuda.txt → requirements/cuda.txt

@@ -1,2 +1,3 @@
+# Extra dependencies for NVIDIA CUDA
 flash-attn>=2.4.0
 bitsandbytes>=0.43.1

requirements/hpu.txt

@@ -0,0 +1,12 @@
+# # 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

requirements/mps.txt

@@ -0,0 +1 @@
+# Extra dependencies for Apple MPS (Metal Performance Shaders)

requirements/rocm.txt

@@ -0,0 +1,2 @@
+# Extra dependencies for AMD ROCm
+flash-attn>=2.6.2,<2.7.0