From 34fb65fc375279143249ab0667721d533dbd026c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 24 May 2026 12:18:27 +0000
Subject: [PATCH] Dockerfile: install vLLM nightly on amd64 for GRPO
 fast_inference

Unsloth's GRPO notebooks (Qwen3_4B-GRPO.ipynb, Qwen3_8B_FP8_GRPO.ipynb,
Llama_FP8_GRPO.ipynb, etc.) set `fast_inference=True` which requires
vLLM to be importable in the same venv. Install vllm pre-release wheels
from https://wheels.vllm.ai/nightly alongside the cu128 pytorch index,
holding torch==2.10.0 fixed so uv refuses any vLLM build that would
yank torch out from under unsloth.

amd64 only -- vLLM does not publish aarch64 wheels yet
(vllm-project/vllm#31128 is open). On arm64 the GRPO notebooks that
need fast_inference will fail to import vllm; non-GRPO and
fast_inference=False paths are unaffected.

Gated by ARG INSTALL_VLLM=auto so the install can be disabled for
contributors who want a smaller image or are blocked by vllm/torch
resolve conflicts during iteration.
---
 docker/Dockerfile | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 2b557345ea..ca05295723 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -164,6 +164,52 @@ RUN set -eux \
         "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@${UNSLOTH_ZOO_REF}" \
         "unsloth[${UNSLOTH_EXTRA}] @ git+https://github.com/unslothai/unsloth@${UNSLOTH_REF}"
 
+# vLLM nightly (amd64 only). Required by Unsloth's GRPO path when the
+# notebook sets fast_inference=True. We install it as a SECOND uv pass
+# rather than appending to the unified one because:
+#   * vLLM's nightly wheel typically pins a specific cu128 torch build;
+#     adding it to the unified resolve forces uv to consider whether to
+#     swap our pinned torch 2.10.0 for vLLM's choice. Splitting the
+#     install lets the unified pass settle on torch 2.10.0 first, then
+#     vLLM bolts on top.
+#   * --no-deps keeps vLLM from yanking torch / xformers / transformers
+#     out from under unsloth. Empirically vLLM's runtime deps overlap
+#     ~100% with what unsloth already installed, so we can drop them.
+#   * On arm64 vLLM does not publish wheels (vllm-project/vllm#31128 is
+#     open; source-build takes ~2-3h under QEMU and ~1h native). Skipped.
+#
+# https://docs.vllm.ai/en/latest/getting_started/installation/gpu/
+# https://wheels.vllm.ai/nightly
+ARG INSTALL_VLLM=auto
+RUN set -eux \
+ && WANT_VLLM=0 \
+ && case "${INSTALL_VLLM}" in \
+        auto) [ "${TARGETARCH:-amd64}" = "amd64" ] && WANT_VLLM=1 ;; \
+        1|true|yes) WANT_VLLM=1 ;; \
+        0|false|no) WANT_VLLM=0 ;; \
+        *) echo "ERROR: invalid INSTALL_VLLM=${INSTALL_VLLM}" >&2; exit 1 ;; \
+    esac \
+ && if [ "${WANT_VLLM}" = "1" ]; then \
+        echo ">> installing vLLM nightly (TARGETARCH=${TARGETARCH:-amd64})"; \
+        # Let uv resolve vLLM's transitive deps. We pin torch==2.10.0 so
+        # uv MUST hold our torch fixed; if vLLM nightly wants a different
+        # torch the build will fail loudly and we revisit. `unsafe-best-
+        # match` lets uv pull from whichever of the three indexes has a
+        # better wheel for each package.
+        ${VENV}/bin/uv pip install \
+            --python ${VENV}/bin/python \
+            --pre \
+            --index-strategy unsafe-best-match \
+            --extra-index-url https://wheels.vllm.ai/nightly \
+            --extra-index-url https://download.pytorch.org/whl/cu128 \
+            "torch==2.10.0" \
+            vllm; \
+        echo ">> vLLM installed:"; \
+        ${VENV}/bin/python -c "import vllm; print('vllm', vllm.__version__)"; \
+    else \
+        echo ">> vLLM skipped (INSTALL_VLLM=${INSTALL_VLLM}, TARGETARCH=${TARGETARCH:-amd64})"; \
+    fi
+
 # 5) Emit a lockfile so the next rebuild can be byte-identical even if PyPI
 #    has moved on. Bake it into the image at /opt/unsloth-venv/requirements.lock.txt
 #    so `docker run ... cat /opt/unsloth-venv/requirements.lock.txt > pins.txt`