A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[ci/build] update torch nightly version for GH200 (#15135)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Mar 23, 2025 at 14:04 UTC
09b6a95551097905b4930b12d32ff102ae63104b
2 files changed
+6
-3
.buildkite/run-gh200-test.sh
+2
-1
@@ -14,6 +14,7 @@ DOCKER_BUILDKIT=1 docker build . \
14
-t gh200-test \
15
--build-arg max_jobs=66 \
16
--build-arg nvcc_threads=2 \
17
+ --build-arg RUN_WHEEL_CHECK=false \
18
--build-arg torch_cuda_arch_list="9.0+PTX" \
19
--build-arg vllm_fa_cmake_gpu_arches="90-real"
20
@@ -23,6 +24,6 @@ trap remove_docker_container EXIT
24
remove_docker_container
25
26
# Run the image and test offline inference
26
-docker run -e HF_TOKEN -v /root/.cache/huggingface:/root/.cache/huggingface --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c '
27
+docker run -e HF_TOKEN -e VLLM_WORKER_MULTIPROC_METHOD=spawn -v /root/.cache/huggingface:/root/.cache/huggingface --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c '
28
python3 examples/offline_inference/basic/generate.py --model meta-llama/Llama-3.2-1B
29
'
Dockerfile
+4
-2
@@ -52,7 +52,8 @@ WORKDIR /workspace
52
# after this step
53
RUN --mount=type=cache,target=/root/.cache/uv \
54
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
55
- uv pip install --index-url https://download.pytorch.org/whl/nightly/cu126 "torch==2.7.0.dev20250121+cu126" "torchvision==0.22.0.dev20250121"; \
55
+ uv pip install --index-url https://download.pytorch.org/whl/nightly/cu128 "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319"; \
56
+ uv pip install --index-url https://download.pytorch.org/whl/nightly/cu128 --pre pytorch_triton==3.3.0+gitab727c40; \
57
fi
58
59
COPY requirements/common.txt requirements/common.txt
@@ -200,7 +201,8 @@ RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
201
# after this step
202
RUN --mount=type=cache,target=/root/.cache/uv \
203
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
203
- uv pip install --index-url https://download.pytorch.org/whl/nightly/cu124 "torch==2.6.0.dev20241210+cu124" "torchvision==0.22.0.dev20241215"; \
204
+ uv pip install --index-url https://download.pytorch.org/whl/nightly/cu128 "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319"; \
205
+ uv pip install --index-url https://download.pytorch.org/whl/nightly/cu128 --pre pytorch_triton==3.3.0+gitab727c40; \
206
fi
207
208
# Install vllm wheel first, so that torch etc will be installed.