chore: update to torch 2.4 (#2259)

* chore: update to torch 2.4

* remove un-necessary patch

* fix

Dockerfile

@@ -41,7 +41,7 @@ RUN cargo build --profile release-opt
 FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS pytorch-install
 
 # NOTE: When updating PyTorch version, beware to remove `pip install nvidia-nccl-cu12==2.22.3` below in the Dockerfile. Context: https://github.com/huggingface/text-generation-inference/pull/2099
-ARG PYTORCH_VERSION=2.3.0
+ARG PYTORCH_VERSION=2.4.0
 
 ARG PYTHON_VERSION=3.10
 # Keep in sync with `server/pyproject.toml
@@ -167,8 +167,6 @@ FROM kernel-builder AS fbgemm-builder
 WORKDIR /usr/src
 
 COPY server/Makefile-fbgemm Makefile
-COPY server/fbgemm_remove_unused.patch fbgemm_remove_unused.patch
-COPY server/fix_torch90a.sh fix_torch90a.sh
 
 RUN make build-fbgemm
 
@@ -254,10 +252,7 @@ COPY server/Makefile server/Makefile
 RUN cd server && \
     make gen-server && \
     pip install -r requirements_cuda.txt && \
-    pip install ".[bnb, accelerate, quantize, peft, outlines]" --no-cache-dir && \
-    pip install nvidia-nccl-cu12==2.22.3
-
-ENV LD_PRELOAD=/opt/conda/lib/python3.10/site-packages/nvidia/nccl/lib/libnccl.so.2
+    pip install ".[bnb, accelerate, quantize, peft, outlines]" --no-cache-dir
 
 # Deps before the binaries
 # The binaries change on every build given we burn the SHA into them

server/Makefile

@@ -30,7 +30,6 @@ install: install-cuda
 
 install-cuda: install-server install-flash-attention-v2-cuda install-vllm-cuda install-flash-attention install-fbgemm
 	pip install -e ".[bnb]"
-	pip install nvidia-nccl-cu12==2.22.3
 
 install-rocm: install-server install-flash-attention-v2-rocm  install-vllm-rocm
 

server/Makefile-fbgemm

@@ -1,10 +1,8 @@
-fbgemm_commit := 9cf0429b726931cfab72b8264730bea682f32fca
+fbgemm_commit := ddac8dd9fc0bee70a3f456df68b8aac38576c856
 
 build-fbgemm:
-	chmod +x fix_torch90a.sh && ./fix_torch90a.sh && \
 	git clone https://github.com/pytorch/FBGEMM.git fbgemm && \
-	cp fbgemm_remove_unused.patch fbgemm && \
-	cd fbgemm && git fetch && git checkout $(fbgemm_commit) && git apply fbgemm_remove_unused.patch && \
+	cd fbgemm && git fetch && git checkout $(fbgemm_commit)  && \
 	git submodule update --init --recursive && \
 	cd fbgemm_gpu && \
 	pip install -r requirements.txt && \

server/pyproject.toml

@@ -31,7 +31,7 @@ einops = "^0.6.1"
 texttable = { version = "^1.6.7", optional = true }
 datasets = { version = "^2.14.0", optional = true }
 peft = { version = "^0.10", optional = true }
-torch = { version = "^2.3.0", optional = true }
+torch = { version = "^2.4.0", optional = true }
 scipy = "^1.11.1"
 pillow = "^10.0.0"
 outlines= { version = "^0.0.34", optional = true }