From f4ec65da5b4e1d73c2f1ef2483d56820c17642b0 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Tue, 21 Jul 2026 06:58:20 -0700 Subject: [PATCH 1/4] adding Dockerfile for nemotron3 example Signed-off-by: Grzegorz Karch --- .../configs/families/nemotron3/Dockerfile | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 examples/puzzletron/configs/families/nemotron3/Dockerfile diff --git a/examples/puzzletron/configs/families/nemotron3/Dockerfile b/examples/puzzletron/configs/families/nemotron3/Dockerfile new file mode 100644 index 00000000000..cb5eafbdffe --- /dev/null +++ b/examples/puzzletron/configs/families/nemotron3/Dockerfile @@ -0,0 +1,41 @@ +ARG CUDA_VERSION=12.9.2 +ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04 + +FROM ${BUILD_BASE_IMAGE} AS base + +ARG DEBIAN_FRONTEND=noninteractive + +RUN apt-get update +RUN apt-get install -y build-essential cmake git ninja-build python3 python3-dev python3-pip python3-venv + +ARG MODEL_OPT_ROOT=/workspace/modelopt +ARG VLLM_ROOT=/workspace/vllm +ARG AUTOMODEL_ROOT=/workspace/Automodel + +RUN git clone --branch feature/add_anymodel_to_vllm --single-branch https://github.com/Separius/vllm.git "${VLLM_ROOT}" +RUN git clone --branch puzzletron --single-branch https://github.com/Separius/Automodel.git "${AUTOMODEL_ROOT}" + +RUN python3 -m venv /venv +ENV VIRTUAL_ENV=/venv +ENV PATH=/venv/bin/:$PATH + +RUN python3 -m pip install --upgrade pip "setuptools>=80,<81" "setuptools-scm>=8" setuptools-rust wheel "packaging>=24.2" "cmake>=3.26.1" ninja jinja2 + +RUN python -m pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu129 + +RUN VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 python -m pip install --no-build-isolation -e "${VLLM_ROOT}" + +RUN python -m pip install -e "${AUTOMODEL_ROOT}" +RUN python -m pip install aiperf + +COPY . "${MODEL_OPT_ROOT}"/ +RUN python -m pip install -e "${MODEL_OPT_ROOT}[hf]" + +# Mixture of experts +RUN python -m pip install --no-build-isolation "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" + +# Mamba +RUN python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation + +# Linear attention +RUN python -m pip install "flash-linear-attention[cuda]" From 4929d0852f380d9efa3f94e8fcc849bc8796463d Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Wed, 22 Jul 2026 12:06:19 +0200 Subject: [PATCH 2/4] fixed dockerfile Signed-off-by: Grzegorz Karch --- examples/puzzletron/configs/families/nemotron3/Dockerfile | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/puzzletron/configs/families/nemotron3/Dockerfile b/examples/puzzletron/configs/families/nemotron3/Dockerfile index cb5eafbdffe..f1dedfff3bc 100644 --- a/examples/puzzletron/configs/families/nemotron3/Dockerfile +++ b/examples/puzzletron/configs/families/nemotron3/Dockerfile @@ -31,7 +31,12 @@ RUN python -m pip install aiperf COPY . "${MODEL_OPT_ROOT}"/ RUN python -m pip install -e "${MODEL_OPT_ROOT}[hf]" -# Mixture of experts +# Tell the build system to proceed without a live GPU +ARG FORCE_CUDA=1 +# Set TORCH_CUDA_ARCH_LIST so the compiler doesn't need to query hardware. +ARG TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0" + +# Mixture of experts — TORCH_CUDA_ARCH_LIST avoids GPU detection at build time RUN python -m pip install --no-build-isolation "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" # Mamba From 7527a4cb77283832ada0e688155a7f0cf6b4f318 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Thu, 23 Jul 2026 01:24:40 -0700 Subject: [PATCH 3/4] added dockerfile, updated readme Signed-off-by: Grzegorz Karch --- examples/puzzletron/README.md | 50 ++++++++++++++++--- .../configs/families/nemotron3/Dockerfile | 42 ++++++---------- 2 files changed, 58 insertions(+), 34 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 4eeb32688e4..a8d86381f8b 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -43,7 +43,45 @@ extension against that same installation; mixing PyTorch or CUDA builds can cause import failures or incorrect GPU execution. AIPerf uses the official PyPI package; no custom AIPerf fork is required. -### 1. Start from the CUDA development image +### Quick start + +Two automated options are provided. Both are driven by the same script so +there is no duplication between them: + +**Option A — Docker image** (recommended for reproducibility): + +```bash +# Build from the repo root so the full source tree is available as build context +docker build \ + -f examples/puzzletron/configs/families/nemotron3/Dockerfile \ + -t puzzletron-nemotron3 \ + . +``` + +**Option B — Bare-metal / cluster** (no Docker required): + +```bash +export MODEL_OPT_ROOT=/path/to/modelopt # defaults to repo root when run from there +export VLLM_ROOT=/workspace/vllm +export AUTOMODEL_ROOT=/workspace/Automodel +export VIRTUAL_ENV=/venv + +# SKIP_APT=1 if build-essential/cmake/git/ninja are already installed +bash examples/puzzletron/configs/families/nemotron3/setup_env.sh + +source "${VIRTUAL_ENV}/bin/activate" +``` + +The script is idempotent: existing git clones and the venv are reused. +`FORCE_CUDA=1` and `TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0"` are set by +default so CUDA extensions build without a live GPU. + +### Manual setup (reference) + +The steps below document what the script does. Follow them if you need to +customise individual stages or debug a failed install. + +#### 1. Start from the CUDA development image Use this image for local containers and cluster jobs: @@ -71,7 +109,7 @@ DEBIAN_FRONTEND=noninteractive apt-get install -y \ python3 python3-dev python3-pip python3-venv ``` -### 2. Clone the tracked forks +#### 2. Clone the tracked forks Keep ModelOpt and the two Puzzletron forks as siblings: @@ -93,7 +131,7 @@ git clone --branch puzzletron --single-branch \ └── Automodel/ ``` -### 3. Create the environment and install runtime packages +#### 3. Create the environment and install runtime packages The patched vLLM branch uses PyTorch 2.11.0 with CUDA 12.9. Install that combination before anything that compiles CUDA code: @@ -116,7 +154,6 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \ python -m pip install -e "${AUTOMODEL_ROOT}" python -m pip install aiperf python -m pip install -e "${MODEL_OPT_ROOT}[hf]" -python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements.txt" ``` Do not add `--no-deps`: the packages need their declared Python dependencies. @@ -127,7 +164,8 @@ Install only the model-specific kernels required by the target architecture: ```bash # Mixture of experts -python -m pip install --no-build-isolation \ +FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0" \ + python -m pip install --no-build-isolation \ "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" # Mamba @@ -137,7 +175,7 @@ python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation python -m pip install "flash-linear-attention[cuda]" ``` -### 4. Verify the exact environment +#### 4. Verify the exact environment Run these checks inside the same container and venv used for Puzzletron jobs: diff --git a/examples/puzzletron/configs/families/nemotron3/Dockerfile b/examples/puzzletron/configs/families/nemotron3/Dockerfile index f1dedfff3bc..9aa792e8bcd 100644 --- a/examples/puzzletron/configs/families/nemotron3/Dockerfile +++ b/examples/puzzletron/configs/families/nemotron3/Dockerfile @@ -5,42 +5,28 @@ FROM ${BUILD_BASE_IMAGE} AS base ARG DEBIAN_FRONTEND=noninteractive -RUN apt-get update -RUN apt-get install -y build-essential cmake git ninja-build python3 python3-dev python3-pip python3-venv - ARG MODEL_OPT_ROOT=/workspace/modelopt ARG VLLM_ROOT=/workspace/vllm ARG AUTOMODEL_ROOT=/workspace/Automodel -RUN git clone --branch feature/add_anymodel_to_vllm --single-branch https://github.com/Separius/vllm.git "${VLLM_ROOT}" -RUN git clone --branch puzzletron --single-branch https://github.com/Separius/Automodel.git "${AUTOMODEL_ROOT}" - -RUN python3 -m venv /venv -ENV VIRTUAL_ENV=/venv -ENV PATH=/venv/bin/:$PATH - -RUN python3 -m pip install --upgrade pip "setuptools>=80,<81" "setuptools-scm>=8" setuptools-rust wheel "packaging>=24.2" "cmake>=3.26.1" ninja jinja2 - -RUN python -m pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu129 - -RUN VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 python -m pip install --no-build-isolation -e "${VLLM_ROOT}" - -RUN python -m pip install -e "${AUTOMODEL_ROOT}" -RUN python -m pip install aiperf - -COPY . "${MODEL_OPT_ROOT}"/ -RUN python -m pip install -e "${MODEL_OPT_ROOT}[hf]" +# Copy only the setup script first so that the expensive dep-install layer +# stays cached as long as the script hasn't changed, regardless of source edits. +COPY examples/puzzletron/configs/families/nemotron3/setup_env.sh /tmp/setup_env.sh # Tell the build system to proceed without a live GPU ARG FORCE_CUDA=1 -# Set TORCH_CUDA_ARCH_LIST so the compiler doesn't need to query hardware. +# Avoid GPU detection at compile time ARG TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0" -# Mixture of experts — TORCH_CUDA_ARCH_LIST avoids GPU detection at build time -RUN python -m pip install --no-build-isolation "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" +RUN MODEL_OPT_ROOT="${MODEL_OPT_ROOT}" \ + VLLM_ROOT="${VLLM_ROOT}" \ + AUTOMODEL_ROOT="${AUTOMODEL_ROOT}" \ + bash /tmp/setup_env.sh --deps -# Mamba -RUN python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation +ENV VIRTUAL_ENV=/venv +ENV PATH=/venv/bin/:$PATH + +# Copy the full repo after deps so source changes don't bust the cache above. +COPY . "${MODEL_OPT_ROOT}"/ -# Linear attention -RUN python -m pip install "flash-linear-attention[cuda]" +RUN MODEL_OPT_ROOT="${MODEL_OPT_ROOT}" bash /tmp/setup_env.sh --modelopt From 99f051b20221708061ef378837aa237c989573cd Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Thu, 23 Jul 2026 01:27:59 -0700 Subject: [PATCH 4/4] setup script Signed-off-by: Grzegorz Karch --- .../configs/families/nemotron3/setup_env.sh | 119 ++++++++++++++++++ 1 file changed, 119 insertions(+) create mode 100755 examples/puzzletron/configs/families/nemotron3/setup_env.sh diff --git a/examples/puzzletron/configs/families/nemotron3/setup_env.sh b/examples/puzzletron/configs/families/nemotron3/setup_env.sh new file mode 100755 index 00000000000..2cd54282e0a --- /dev/null +++ b/examples/puzzletron/configs/families/nemotron3/setup_env.sh @@ -0,0 +1,119 @@ +#!/usr/bin/env bash +# Shared setup script for the Nemotron3 Puzzletron environment. +# +# Used by the Dockerfile (--deps before COPY, --modelopt after COPY) and +# can be run standalone to set up a bare-metal environment (no arguments). +# +# Usage: +# ./setup_env.sh [--deps | --modelopt] +# --deps System packages, git clones, venv, pip deps — safe to Docker-cache +# --modelopt Install the local ModelOpt source only (run after --deps) +# (no args) Full setup — suitable for bare-metal +# +# Configurable via environment variables (all have defaults): +# MODEL_OPT_ROOT Path to the ModelOpt repo root +# Default: 5 directories above this script +# VLLM_ROOT Where to clone vLLM (default: /workspace/vllm) +# AUTOMODEL_ROOT Where to clone Automodel (default: /workspace/Automodel) +# VIRTUAL_ENV Python venv path (default: /venv) +# SKIP_APT Set to 1 to skip apt-get (default: 0) +# FORCE_CUDA Build CUDA exts without a live GPU (default: 1) +# TORCH_CUDA_ARCH_LIST Architectures to compile for +# (default: "8.0;8.6;9.0;10.0") + +set -euo pipefail + +MODE="${1:-all}" + +# ── Paths ────────────────────────────────────────────────────────────────────── +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# examples/puzzletron/configs/families/nemotron3 → 5 levels up to repo root +MODEL_OPT_ROOT="${MODEL_OPT_ROOT:-$(cd "${SCRIPT_DIR}/../../../../.." && pwd)}" +VLLM_ROOT="${VLLM_ROOT:-/workspace/vllm}" +AUTOMODEL_ROOT="${AUTOMODEL_ROOT:-/workspace/Automodel}" +VIRTUAL_ENV="${VIRTUAL_ENV:-/venv}" + +# ── Build flags ──────────────────────────────────────────────────────────────── +SKIP_APT="${SKIP_APT:-0}" +export FORCE_CUDA="${FORCE_CUDA:-1}" +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-8.0;8.6;9.0;10.0}" + +activate_venv() { + export VIRTUAL_ENV + export PATH="${VIRTUAL_ENV}/bin:${PATH}" +} + +# ══════════════════════════════════════════════════════════════════════════════ +# DEPS phase — nothing here depends on the local ModelOpt source. +# In Docker this runs before `COPY .` so it stays cached across source changes. +# ══════════════════════════════════════════════════════════════════════════════ +install_deps() { + if [[ "${SKIP_APT}" != "1" ]]; then + apt-get update + apt-get install -y build-essential cmake git ninja-build \ + python3 python3-dev python3-pip python3-venv + fi + + # Clone external repos (idempotent) + [[ -d "${VLLM_ROOT}" ]] || \ + git clone --branch feature/add_anymodel_to_vllm --single-branch \ + https://github.com/Separius/vllm.git "${VLLM_ROOT}" + [[ -d "${AUTOMODEL_ROOT}" ]] || \ + git clone --branch puzzletron --single-branch \ + https://github.com/Separius/Automodel.git "${AUTOMODEL_ROOT}" + + [[ -d "${VIRTUAL_ENV}" ]] || python3 -m venv "${VIRTUAL_ENV}" + activate_venv + + python -m pip install --upgrade pip \ + "setuptools>=80,<81" "setuptools-scm>=8" setuptools-rust wheel \ + "packaging>=24.2" "cmake>=3.26.1" ninja jinja2 + + python -m pip install \ + torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 \ + --index-url https://download.pytorch.org/whl/cu129 + + VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \ + python -m pip install --no-build-isolation -e "${VLLM_ROOT}" + + python -m pip install -e "${AUTOMODEL_ROOT}" + python -m pip install aiperf + + # CUDA extension packages placed here (before COPY .) so they stay cached + # across ModelOpt source changes in Docker builds. + python -m pip install --no-build-isolation \ + "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" + python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation + python -m pip install "flash-linear-attention[cuda]" +} + +# ══════════════════════════════════════════════════════════════════════════════ +# MODELOPT phase — installs the local ModelOpt source. +# In Docker this runs after `COPY .` and re-runs on every source change. +# ══════════════════════════════════════════════════════════════════════════════ +install_modelopt() { + activate_venv + python -m pip install -e "${MODEL_OPT_ROOT}[hf]" +} + +# ── Dispatch ─────────────────────────────────────────────────────────────────── +case "${MODE}" in + --deps) + install_deps + ;; + --modelopt) + install_modelopt + ;; + all) + install_deps + install_modelopt + echo "" + echo "Setup complete. Activate your environment with:" + echo " source ${VIRTUAL_ENV}/bin/activate" + ;; + *) + echo "Unknown argument: ${MODE}" >&2 + echo "Usage: $0 [--deps | --modelopt]" >&2 + exit 1 + ;; +esac