feat: Offload the vision tower to host RAM
The vision tower sat on GPU 0 and 1 all the time, and did work only when a request carried an image. It now streams from pinned host RAM and encodes over H2D, which frees 432 MiB on GPU 0 and 1. Stage 1 carries no tower, so GPU 2 and 3 gain nothing, and the KV cache stays at 924,993 tokens because --kv-cache-memory pins it. The freed memory is headroom, not capacity. An image request passes: 2223 image tokens in 11.77 s, with no VRAM growth. Key the prebuilt ops cache on the compiled inputs, so a Python-only patch installs from the shipped tarball instead of a 20-40 minute compile. Raise --mm-processor-cache-gb to 16. That cache is duplicated for each API process and engine core process, so 16 GiB costs 32 GiB of host RAM and no VRAM, and repeated images stop re-processing. Co-Authored-By: Hinata <hinata@idihore.id> Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
+30
-4
@@ -181,9 +181,9 @@ docker build -t vllm:latest .
|
||||
|
||||
`docker compose build` mengirim context `ops` dari `artifacts/` sendiri.
|
||||
|
||||
Build memeriksa tarball sebelum ia memasangnya: head tree, arsitektur GPU,
|
||||
CPython minor, CUDA major, versi torch, dan sha256 setiap file hasil ekstrak.
|
||||
Tarball yang tidak cocok dengan tree menghentikan build dan menyebut item yang
|
||||
Build memeriksa tarball sebelum ia memasangnya: input yang di-compile (lihat 7.5),
|
||||
arsitektur GPU, CPython minor, CUDA major, versi torch, dan sha256 setiap file hasil
|
||||
ekstrak. Tarball yang tidak cocok dengan tree menghentikan build dan menyebut item yang
|
||||
tidak cocok.
|
||||
|
||||
Build image butuh akses jaringan: ia mengunduh vLLM v0.30.0 dan wheel terkunci.
|
||||
@@ -260,6 +260,32 @@ docker run -d --name vllm --gpus all --stop-timeout 20 --shm-size=16g \
|
||||
|
||||
Tanpa Compose, mount sendiri sidecar-nya: `-v <host>/scales.json:/opt/vllm/qsa_kv_scales.json:ro -e SCALES=/opt/vllm/qsa_kv_scales.json`.
|
||||
|
||||
### 7.5. Tambah Patch Tanpa Compile Ulang
|
||||
|
||||
Tarball ops berisi kernel yang sudah di-compile. `build.sh` menguncinya pada `OPS_OBJ`,
|
||||
yaitu hash dari input yang di-compile saja: `csrc`, `CMakeLists.txt`, `cmake`, `setup.py`,
|
||||
`pyproject.toml`, `MANIFEST.in`, `requirements`, `tools`, `rust-toolchain.toml`,
|
||||
`vllm/third_party`, dan `vllm/vllm_flash_attn`. Daftarnya ada di `OPS_PATHS` pada `build.sh`.
|
||||
|
||||
Patch yang hanya menyentuh Python tidak mengubah `OPS_OBJ`. `artifacts/` tetap valid, jadi
|
||||
build image tetap di jalur ~2 menit. Patch yang menyentuh satu path di daftar itu membatalkan
|
||||
tarball, dan build akan meng-compile.
|
||||
|
||||
Cara menambah patch ke seri:
|
||||
|
||||
1. Taruh patch format mail di akhir `patches/`, dinamai `NNNN-<slug>.patch`.
|
||||
`git format-patch` menulis format itu.
|
||||
2. Jalankan `./build.sh`. Step 2 memasang seri dan mencetak tree yang ia build kalau tree itu
|
||||
tidak cocok dengan `EXPECTED_TREE`.
|
||||
3. Tulis hash tree itu ke `EXPECTED_TREE` di `build.sh`, beserta tanggal verifikasi.
|
||||
4. Jalankan `./build.sh` lagi. Ia memasang seri, mencocokkan tree, dan memakai ulang kernel.
|
||||
|
||||
Patch dari fork upstream jarang langsung cocok ke seri ini. Seri ini menggeser file yang sama.
|
||||
Uji di tree bersih sebelum menambahnya:
|
||||
`git -C vllm worktree add -q --detach /tmp/wt HEAD`, lalu `git -C /tmp/wt am --keep-non-patch
|
||||
<patch>`, lalu `git -C vllm worktree remove --force /tmp/wt`. Kalau konflik, rebase patch itu
|
||||
ke seri dan ekspor ulang dengan `git format-patch`.
|
||||
|
||||
## 8. File
|
||||
|
||||
| File | Guna |
|
||||
@@ -271,7 +297,7 @@ Tanpa Compose, mount sendiri sidecar-nya: `-v <host>/scales.json:/opt/vllm/qsa_k
|
||||
| `systemd/` | `vllm.service` dan kebijakan logrotate |
|
||||
| `Dockerfile`, `docker-compose.yml`, `.env.example` | Jalur container |
|
||||
| `scripts/` | Pemilih topologi, requirements terkunci, stub driver CUDA |
|
||||
| `patches/` | 103 patch, dalam urutan pasang |
|
||||
| `patches/` | 105 patch, dalam urutan pasang |
|
||||
| `benchmarks/` | Wrapper `bench.sh`, dan benchmark power limit (Inggris, Indonesia) |
|
||||
| `abliteration/` | File arah refusal untuk `QWEN38_ABLITERATION` |
|
||||
|
||||
|
||||
@@ -174,10 +174,10 @@ docker build -t vllm:latest .
|
||||
|
||||
`docker compose build` passes the `ops` context from `artifacts/` by itself.
|
||||
|
||||
The build checks the tarball before it installs it: the tree head, the GPU architecture,
|
||||
the CPython minor, the CUDA major, the torch version, and the sha256 of every unpacked
|
||||
file. A tarball that does not fit the tree stops the build and names the item that does
|
||||
not match.
|
||||
The build checks the tarball before it installs it: the compiled inputs (see 7.5), the GPU
|
||||
architecture, the CPython minor, the CUDA major, the torch version, and the sha256 of every
|
||||
unpacked file. A tarball that does not fit the tree stops the build and names the item that
|
||||
does not match.
|
||||
|
||||
The image build needs network access: it downloads vLLM v0.30.0 and the pinned wheels. The
|
||||
base is Debian trixie. Kernels built on Debian 13 need glibc 2.38, and bookworm (glibc
|
||||
@@ -248,6 +248,32 @@ docker run -d --name vllm --gpus all --stop-timeout 20 --shm-size=16g \
|
||||
|
||||
Without Compose, mount the sidecar yourself: `-v <host>/scales.json:/opt/vllm/qsa_kv_scales.json:ro -e SCALES=/opt/vllm/qsa_kv_scales.json`.
|
||||
|
||||
### 7.5. Add a patch without a recompile
|
||||
|
||||
The ops tarball carries compiled kernels. `build.sh` keys it on `OPS_OBJ`, the hash of the
|
||||
compiled inputs only: `csrc`, `CMakeLists.txt`, `cmake`, `setup.py`, `pyproject.toml`,
|
||||
`MANIFEST.in`, `requirements`, `tools`, `rust-toolchain.toml`, `vllm/third_party`, and
|
||||
`vllm/vllm_flash_attn`. The list is `OPS_PATHS` in `build.sh`.
|
||||
|
||||
A patch that touches Python only leaves `OPS_OBJ` alone. `artifacts/` stays valid, so the
|
||||
image build keeps the ~2 minute route. A patch that moves any path in that list invalidates
|
||||
the tarball, and the build compiles.
|
||||
|
||||
To add a patch to the series:
|
||||
|
||||
1. Put the mail-format patch at the end of `patches/`, named `NNNN-<slug>.patch`.
|
||||
`git format-patch` writes that format.
|
||||
2. Run `./build.sh`. Step 2 applies the series and prints the tree it built when it does not
|
||||
match `EXPECTED_TREE`.
|
||||
3. Write that tree hash into `EXPECTED_TREE` in `build.sh`, with the date you verified it.
|
||||
4. Run `./build.sh` again. It applies the series, matches the tree, and reuses the kernels.
|
||||
|
||||
A patch from a fork of upstream rarely applies to this series as-is. The series moves the
|
||||
same files. Test it on a clean tree before you add it:
|
||||
`git -C vllm worktree add -q --detach /tmp/wt HEAD`, then `git -C /tmp/wt am --keep-non-patch
|
||||
<patch>`, then `git -C vllm worktree remove --force /tmp/wt`. On a conflict, rebase the patch
|
||||
onto the series and re-export it with `git format-patch`.
|
||||
|
||||
## 8. Files
|
||||
|
||||
| File | Use |
|
||||
@@ -260,7 +286,7 @@ Without Compose, mount the sidecar yourself: `-v <host>/scales.json:/opt/vllm/qs
|
||||
| `Dockerfile`, `docker-compose.yml`, `.env.example` | The container path |
|
||||
| `README.md`, `README-ID.md` | This guide, and the Indonesian translation |
|
||||
| `scripts/` | Topology selector, pinned requirements, CUDA driver stub |
|
||||
| `patches/` | The 103 patches, in apply order |
|
||||
| `patches/` | The 105 patches, in apply order |
|
||||
| `benchmarks/` | `bench.sh` wrapper, and the power-limit benchmark (English, Indonesian) |
|
||||
| `abliteration/` | The refusal-direction file for `QWEN38_ABLITERATION` |
|
||||
|
||||
|
||||
@@ -16,10 +16,10 @@
|
||||
# installs the tree into .venv in place, and gives the tree its version and console script.
|
||||
#
|
||||
# ARTIFACTS=<tarball> installs compiled kernels of this same tree instead of compiling them. step 6
|
||||
# unpacks them after it checks the tree head, the arch, the CPython minor, the CUDA major and the
|
||||
# torch version in the tarball's manifest, and hashes every file after the unpack. step 7 then
|
||||
# installs the tree without compiling (VLLM_TARGET_DEVICE=empty). step 8 writes that tarball, so the
|
||||
# Dockerfile uses this route when you pass --build-context ops=$PWD/artifacts.
|
||||
# unpacks them after it checks the compiled inputs (OPS_OBJ below), the arch, the CPython minor,
|
||||
# the CUDA major and the torch version in the tarball's manifest, and hashes every file after the
|
||||
# unpack. step 7 then installs the tree without compiling (VLLM_TARGET_DEVICE=empty). step 8 writes
|
||||
# that tarball, so the Dockerfile uses this route when you pass --build-context ops=$PWD/artifacts.
|
||||
#
|
||||
# step 8 packs the kernels into artifacts/ for the image build (PACK=0 or --no-pack skips it; the
|
||||
# Dockerfile sets it, an image build must not grow a 340 MB tarball inside a layer).
|
||||
@@ -102,7 +102,7 @@ step 2 "serving tree (upstream v0.30.0 + $(ls "$HERE"/patches/*.patch | wc -l) p
|
||||
# The tree is not in this repo: this repo ships the patches, not the tree. Fetch applies the series
|
||||
# to the upstream tag and asserts the result, so a patch series that drifts fails here, not at serve.
|
||||
BASE=ced6857afa0ea7b2e3f0846a62e1394e90f15607 # tag v0.30.0
|
||||
EXPECTED_TREE=b7fe074a6c11b5e60ef15f31f432c67626607454 # BASE + patches/*.patch, verified 2026-10-11
|
||||
EXPECTED_TREE=a47150d66ebe25da5b1da08bdbb86820a4aa3fac # BASE + patches/*.patch (105), verified 2026-10-18
|
||||
URL="${VLLM_UPSTREAM_URL:-https://github.com/vllm-project/vllm.git}"
|
||||
# git am makes a commit, so git needs a committer identity: the host's ~/.gitconfig is not there (a
|
||||
# container, a build user, a fresh box). The committer is not in the tree hash, so the assertion
|
||||
@@ -137,6 +137,12 @@ TREE_HEAD="$(git -C "$SRC" rev-parse --short HEAD 2>/dev/null || echo unknown)"
|
||||
# box that applied the patches, so the host commit and the container commit of one patch set are
|
||||
# not the same SHA. The tree object is identical. Compare that, not the SHA.
|
||||
TREE_OBJ="$(git -C "$SRC" rev-parse 'HEAD^{tree}' 2>/dev/null || echo unknown)"
|
||||
# The ops cache key: the compiled inputs only. A patch that touches Python alone leaves this key
|
||||
# alone, so a tarball of prebuilt kernels stays valid for it and the image build keeps the ~2 min
|
||||
# route. A patch that moves any path below invalidates the tarball, and the build compiles.
|
||||
# A path missing from the tree yields "missing", which no manifest can match: the fail-safe way.
|
||||
OPS_PATHS="csrc CMakeLists.txt cmake setup.py pyproject.toml MANIFEST.in requirements tools rust-toolchain.toml vllm/third_party vllm/vllm_flash_attn"
|
||||
OPS_OBJ="$(for p in $OPS_PATHS; do git -C "$SRC" rev-parse -q --verify "HEAD^{tree}:$p" 2>/dev/null || echo missing; done | LC_ALL=C sha256sum | cut -d' ' -f1)"
|
||||
MARK="$SRC/.build-products"; SUMS="$SRC/.build-products.sha256"
|
||||
LIST="$HERE/scripts/build-artifacts.list"
|
||||
|
||||
@@ -238,13 +244,19 @@ if [ -n "${ARTIFACTS:-}" ]; then
|
||||
# spaces, which the shell cannot source.
|
||||
mf(){ sed -n "s/^$1=//p" "$adir/$stem.manifest" | head -1; }
|
||||
m_tree="$(mf tree)"; m_tree_obj="$(mf tree_obj)"; m_arch="$(mf arch)"; m_py="$(mf python)"; m_cuda="$(mf cuda)"
|
||||
m_torch="$(mf torch)"; m_glibc="$(mf glibc)"; m_files="$(mf files)"; m_built="$(mf built)"
|
||||
echo " manifest: tree=$m_tree tree_obj=${m_tree_obj:0:7} arch=$m_arch python=$m_py cuda=$m_cuda torch=$m_torch glibc=$m_glibc files=$m_files built=$m_built"
|
||||
m_torch="$(mf torch)"; m_glibc="$(mf glibc)"; m_files="$(mf files)"; m_built="$(mf built)"; m_ops_obj="$(mf ops_obj)"
|
||||
echo " manifest: tree=$m_tree tree_obj=${m_tree_obj:0:7} ops_obj=${m_ops_obj:0:7} arch=$m_arch python=$m_py cuda=$m_cuda torch=$m_torch glibc=$m_glibc files=$m_files built=$m_built"
|
||||
echo " this tree: tree_obj=${TREE_OBJ:0:7} ops_obj=${OPS_OBJ:0:7}"
|
||||
fail=0
|
||||
[ -n "$m_tree_obj" ] || { echo "$adir/$stem.manifest has no tree_obj= line: pack the ops again with ./build.sh"; exit 1; }
|
||||
# The tree content is what the kernels were compiled from, so the tree object is the check. The
|
||||
# The compiled inputs are what the kernels were built from, so OPS_OBJ is the check. A Python-only
|
||||
# patch moves the whole tree object, and must not push this build into a 20-40 min compile. The
|
||||
# commit SHA is printed next to it and is allowed to differ between machines (committer date).
|
||||
[ "$m_tree_obj" = "$TREE_OBJ" ] || { echo "the ops were packed from tree ${m_tree_obj:0:7}, this tree is ${TREE_OBJ:0:7} (commits $m_tree / $TREE_HEAD): pack them from this tree, or compile"; fail=1; }
|
||||
if [ -z "$m_ops_obj" ]; then
|
||||
echo " manifest has no ops_obj= line, it predates this check: the strict tree_obj test applies"
|
||||
m_ops_obj="$m_tree_obj"
|
||||
fi
|
||||
[ "$m_ops_obj" = "$OPS_OBJ" ] || { echo "the ops were compiled from inputs ${m_ops_obj:0:7}, this tree compiles from ${OPS_OBJ:0:7} (trees $m_tree_obj / $TREE_OBJ, commits $m_tree / $TREE_HEAD): pack them from this tree, or compile"; fail=1; }
|
||||
venv_py="$("$VENV/bin/python" -c 'import sys; print("%d.%d" % sys.version_info[:2])')"
|
||||
[ "$m_py" = "$venv_py" ] || { echo "the ops are for CPython $m_py, the venv is $venv_py"; fail=1; }
|
||||
[ "$m_cuda" = "$nv_rel" ] || { echo "the ops are for CUDA $m_cuda, this build uses CUDA $nv_rel"; fail=1; }
|
||||
@@ -447,6 +459,7 @@ else
|
||||
{
|
||||
echo "tree=$TREE_HEAD"
|
||||
echo "tree_obj=$TREE_OBJ"
|
||||
echo "ops_obj=$OPS_OBJ"
|
||||
echo "tree_full=$(git -C "$SRC" rev-parse HEAD)"
|
||||
echo "arch=$ARCH"
|
||||
echo "python=$pyv"
|
||||
|
||||
@@ -0,0 +1,517 @@
|
||||
From f5497ecb0ab526ed95511ca781e698226aa0b808 Mon Sep 17 00:00:00 2001
|
||||
From: MOVZX <movzx@yahoo.com>
|
||||
Date: Fri, 16 Oct 2026 14:32:14 +0700
|
||||
Subject: [PATCH 104/105] [Offloader] Stream multimodal towers from host RAM with
|
||||
bulk-copy
|
||||
|
||||
Add --offload-towers: place multimodal tower weights in pinned host
|
||||
memory and bulk-copy each unit to the GPU only during its forward,
|
||||
keeping towers out of VRAM between encodes. Unlike UVA zero-copy,
|
||||
whole-tensor H2D transfers avoid re-reading operand tiles over PCIe
|
||||
inside the GEMM inner loop. Tower streaming is independent of layer
|
||||
group scheduling, so group_size may be 0 for tower-only offload.
|
||||
|
||||
Co-authored-by: Hinata <hinata@idihore.id>
|
||||
Signed-off-by: MOVZX <movzx@yahoo.com>
|
||||
---
|
||||
.../model_executor/offloader/test_prefetch.py | 65 +++++-
|
||||
vllm/config/offload.py | 30 ++-
|
||||
vllm/engine/arg_utils.py | 5 +
|
||||
vllm/model_executor/offloader/base.py | 12 +-
|
||||
vllm/model_executor/offloader/prefetch.py | 209 +++++++++++++++++-
|
||||
5 files changed, 307 insertions(+), 14 deletions(-)
|
||||
|
||||
diff --git a/tests/model_executor/offloader/test_prefetch.py b/tests/model_executor/offloader/test_prefetch.py
|
||||
index 560ac8d..41fc6c2 100644
|
||||
--- a/tests/model_executor/offloader/test_prefetch.py
|
||||
+++ b/tests/model_executor/offloader/test_prefetch.py
|
||||
@@ -2,8 +2,13 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
+import torch
|
||||
+import torch.nn as nn
|
||||
|
||||
-from vllm.model_executor.offloader.prefetch import _get_next_prefetch_index
|
||||
+from vllm.model_executor.offloader.prefetch import (
|
||||
+ _get_next_prefetch_index,
|
||||
+ _TowerUnitStreamer,
|
||||
+)
|
||||
|
||||
pytestmark = pytest.mark.skip_global_cleanup
|
||||
|
||||
@@ -35,3 +40,61 @@ def test_next_prefetch_index_preserves_slot_ownership(
|
||||
target % prefetch_step == index % prefetch_step
|
||||
for index, target in enumerate(targets)
|
||||
)
|
||||
+
|
||||
+
|
||||
+def test_tower_unit_streamer_roundtrip_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
+ """Host-resident weights must compute identically to resident weights,
|
||||
+ and be back on host storage after each forward."""
|
||||
+ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
|
||||
+
|
||||
+ lin = nn.Linear(4, 4, bias=False)
|
||||
+ weight = lin.weight.detach().clone()
|
||||
+
|
||||
+ streamer = _TowerUnitStreamer(lin, ["weight"], torch.device("cpu"))
|
||||
+ streamer.offload_to_host_now()
|
||||
+ assert lin.weight.device.type == "cpu"
|
||||
+
|
||||
+ streamer.sync_storage()
|
||||
+ assert streamer.offloaded_bytes == weight.numel() * weight.element_size()
|
||||
+
|
||||
+ streamer.hook_forward()
|
||||
+ x = torch.randn(2, 4)
|
||||
+ out = lin(x)
|
||||
+ assert torch.allclose(out, x @ weight.T)
|
||||
+ # Weights released back to host after the forward.
|
||||
+ assert lin.weight.device.type == "cpu"
|
||||
+
|
||||
+ # Re-running post_init-style hooking must not double-wrap the forward.
|
||||
+ streamer.hook_forward()
|
||||
+ assert torch.allclose(lin(x), out)
|
||||
+
|
||||
+
|
||||
+def test_offload_config_tower_only_prefetch_is_valid() -> None:
|
||||
+ from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
|
||||
+
|
||||
+ cfg = OffloadConfig(
|
||||
+ offload_backend="prefetch",
|
||||
+ prefetch=PrefetchOffloadConfig(offload_towers=True),
|
||||
+ )
|
||||
+ assert cfg.prefetch.offload_towers
|
||||
+ assert cfg.prefetch.offload_group_size == 0
|
||||
+
|
||||
+
|
||||
+def test_offload_config_prefetch_requires_layers_or_towers() -> None:
|
||||
+ from vllm.config.offload import OffloadConfig
|
||||
+
|
||||
+ with pytest.raises(ValueError):
|
||||
+ OffloadConfig(offload_backend="prefetch")
|
||||
+
|
||||
+
|
||||
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
|
||||
+def test_create_offloader_selects_prefetch_for_towers() -> None:
|
||||
+ from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
|
||||
+ from vllm.model_executor.offloader.base import create_offloader
|
||||
+ from vllm.model_executor.offloader.prefetch import PrefetchOffloader
|
||||
+
|
||||
+ offloader = create_offloader(
|
||||
+ OffloadConfig(prefetch=PrefetchOffloadConfig(offload_towers=True))
|
||||
+ )
|
||||
+ assert isinstance(offloader, PrefetchOffloader)
|
||||
+ assert offloader.supports_tower_offload
|
||||
diff --git a/vllm/config/offload.py b/vllm/config/offload.py
|
||||
index ad65e8a..bb465e3 100644
|
||||
--- a/vllm/config/offload.py
|
||||
+++ b/vllm/config/offload.py
|
||||
@@ -73,6 +73,19 @@ class PrefetchOffloadConfig:
|
||||
parameters of each offloaded layer are offloaded.
|
||||
Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight"
|
||||
but not "mlp.experts.w13_weight_scale".
|
||||
+ Applies to decoder layers only; `offload_towers` offloads the full
|
||||
+ tower regardless of this set.
|
||||
+ """
|
||||
+
|
||||
+ offload_towers: bool = False
|
||||
+ """Offload multimodal tower weights (vision/audio encoders) to pinned
|
||||
+ host memory, and bulk-copy them to the GPU during each encoder forward.
|
||||
+ Towers hold no VRAM between image encodes. Unlike UVA zero-copy, weights
|
||||
+ are streamed with whole-tensor H2D copies instead of being re-read tile
|
||||
+ by tile inside the GEMM inner loops, so encode cost stays near the
|
||||
+ resident-tower baseline.
|
||||
+ Works together with, or instead of, layer prefetch:
|
||||
+ `offload_group_size` may be 0 when only towers are offloaded.
|
||||
"""
|
||||
|
||||
|
||||
@@ -83,7 +96,8 @@ class OffloadConfig:
|
||||
offload_backend: OffloadBackend = "auto"
|
||||
"""The backend for weight offloading. Options:
|
||||
- "auto": Selects based on which sub-config has non-default values
|
||||
- (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0).
|
||||
+ (prefetch if offload_group_size > 0 or offload_towers, uva if
|
||||
+ cpu_offload_gb > 0).
|
||||
- "uva": UVA (Unified Virtual Addressing) zero-copy offloading.
|
||||
- "prefetch": Async prefetch with group-based layer offloading.
|
||||
"""
|
||||
@@ -97,7 +111,15 @@ class OffloadConfig:
|
||||
@model_validator(mode="after")
|
||||
def validate_offload_config(self) -> "OffloadConfig":
|
||||
"""Validate offload configuration constraints."""
|
||||
- if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0:
|
||||
+ layer_offload = self.prefetch.offload_group_size > 0
|
||||
+ if self.offload_backend == "prefetch" and not (
|
||||
+ layer_offload or self.prefetch.offload_towers
|
||||
+ ):
|
||||
+ raise ValueError(
|
||||
+ "offload_backend='prefetch' requires offload_group_size > 0"
|
||||
+ " or offload_towers=True"
|
||||
+ )
|
||||
+ if layer_offload:
|
||||
if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size:
|
||||
raise ValueError(
|
||||
f"offload_num_in_group ({self.prefetch.offload_num_in_group})"
|
||||
@@ -114,7 +136,9 @@ class OffloadConfig:
|
||||
|
||||
# Warn if both backends have non-default values
|
||||
uva_active = self.uva.cpu_offload_gb > 0
|
||||
- prefetch_active = self.prefetch.offload_group_size > 0
|
||||
+ prefetch_active = (
|
||||
+ self.prefetch.offload_group_size > 0 or self.prefetch.offload_towers
|
||||
+ )
|
||||
if self.offload_backend == "uva" and prefetch_active:
|
||||
warnings.warn(
|
||||
"Prefetch offload fields are set but offload_backend='uva'. "
|
||||
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
|
||||
index b420e84..ae1e2df 100644
|
||||
--- a/vllm/engine/arg_utils.py
|
||||
+++ b/vllm/engine/arg_utils.py
|
||||
@@ -551,6 +551,7 @@ class EngineArgs:
|
||||
offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group
|
||||
offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
|
||||
offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
|
||||
+ offload_towers: bool = PrefetchOffloadConfig.offload_towers
|
||||
gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
|
||||
kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
|
||||
max_num_batched_tokens: int | None = None
|
||||
@@ -1356,6 +1357,9 @@ class EngineArgs:
|
||||
offload_group.add_argument(
|
||||
"--offload-params", **prefetch_kwargs["offload_params"]
|
||||
)
|
||||
+ offload_group.add_argument(
|
||||
+ "--offload-towers", **prefetch_kwargs["offload_towers"]
|
||||
+ )
|
||||
|
||||
# Multimodal related configs
|
||||
multimodal_kwargs = get_kwargs(MultiModalConfig)
|
||||
@@ -2623,6 +2627,7 @@ class EngineArgs:
|
||||
offload_num_in_group=self.offload_num_in_group,
|
||||
offload_prefetch_step=self.offload_prefetch_step,
|
||||
offload_params=self.offload_params,
|
||||
+ offload_towers=self.offload_towers,
|
||||
),
|
||||
)
|
||||
|
||||
diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py
|
||||
index cb9c35b..64eba91 100644
|
||||
--- a/vllm/model_executor/offloader/base.py
|
||||
+++ b/vllm/model_executor/offloader/base.py
|
||||
@@ -56,8 +56,9 @@ class BaseOffloader(ABC):
|
||||
`SupportsMultiModal._mark_tower_model`, outside the `make_layers` call.
|
||||
|
||||
Offloaders whose `wrap_modules` may only be called on the decoder layer
|
||||
- stack (e.g. `PrefetchOffloader`, which schedules prefetches over a
|
||||
- circular layer stack) must keep this `False`.
|
||||
+ stack must keep this `False`. `PrefetchOffloader` enables it per-instance
|
||||
+ when `offload_towers` is configured, via a tower path that does not use
|
||||
+ the circular layer stack or the static buffer pool.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -142,8 +143,8 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
|
||||
"""Create an offloader based on the offload configuration.
|
||||
|
||||
Uses the explicit ``offload_backend`` selector. When set to ``"auto"``,
|
||||
- selects prefetch if ``offload_group_size > 0``, UVA if
|
||||
- ``cpu_offload_gb > 0``, otherwise noop.
|
||||
+ selects prefetch if ``offload_group_size > 0`` or ``offload_towers`` is
|
||||
+ set, UVA if ``cpu_offload_gb > 0``, otherwise noop.
|
||||
"""
|
||||
from vllm.model_executor.offloader.prefetch import PrefetchOffloader
|
||||
from vllm.model_executor.offloader.uva import UVAOffloader
|
||||
@@ -153,7 +154,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
|
||||
prefetch = offload_config.prefetch
|
||||
|
||||
if backend == "auto":
|
||||
- if prefetch.offload_group_size > 0:
|
||||
+ if prefetch.offload_group_size > 0 or prefetch.offload_towers:
|
||||
backend = "prefetch"
|
||||
elif uva.cpu_offload_gb > 0:
|
||||
backend = "uva"
|
||||
@@ -166,6 +167,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
|
||||
num_in_group=prefetch.offload_num_in_group,
|
||||
prefetch_step=prefetch.offload_prefetch_step,
|
||||
offload_params=prefetch.offload_params,
|
||||
+ offload_towers=prefetch.offload_towers,
|
||||
mode="cpu",
|
||||
)
|
||||
elif backend == "uva":
|
||||
diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py
|
||||
index e682fd0..f276673 100644
|
||||
--- a/vllm/model_executor/offloader/prefetch.py
|
||||
+++ b/vllm/model_executor/offloader/prefetch.py
|
||||
@@ -150,6 +150,10 @@ class PrefetchOffloader(BaseOffloader):
|
||||
group_size: Group every N layers together.
|
||||
num_in_group: Offload this many layers per group (last N of each group).
|
||||
prefetch_step: Number of layers to prefetch ahead.
|
||||
+ offload_towers: Also offload multimodal tower weights to pinned host
|
||||
+ memory, streaming them to the GPU during each encoder forward
|
||||
+ (bulk-copy). Group scheduling still applies to decoder layers
|
||||
+ only; `group_size` may be 0 when only towers are offloaded.
|
||||
mode: Offload mode ("cpu" is currently supported).
|
||||
"""
|
||||
|
||||
@@ -159,19 +163,24 @@ class PrefetchOffloader(BaseOffloader):
|
||||
num_in_group: int,
|
||||
prefetch_step: int,
|
||||
offload_params: set[str] | None = None,
|
||||
+ offload_towers: bool = False,
|
||||
mode: str = "cpu",
|
||||
):
|
||||
self.group_size = group_size
|
||||
self.num_in_group = num_in_group
|
||||
self.prefetch_step = prefetch_step
|
||||
self.offload_params = offload_params or set()
|
||||
+ self.offload_towers = offload_towers
|
||||
self.mode = mode
|
||||
+ if offload_towers:
|
||||
+ self.supports_tower_offload = True
|
||||
|
||||
# Copy stream for async H2D transfers
|
||||
self.copy_stream = torch.cuda.Stream()
|
||||
|
||||
# Module offloaders and buffer pool (populated in wrap_modules/post_init)
|
||||
self.module_offloaders: list[_ModuleOffloader] = []
|
||||
+ self.tower_streamers: list[_TowerUnitStreamer] = []
|
||||
self.buffer_pool: StaticBufferPool | None = None
|
||||
self.total_offloaded_bytes = 0
|
||||
|
||||
@@ -181,17 +190,29 @@ class PrefetchOffloader(BaseOffloader):
|
||||
prefix: str = "",
|
||||
) -> list[nn.Module]:
|
||||
"""Wrap modules with prefetch offloading logic."""
|
||||
- assert len(self.module_offloaders) == 0, (
|
||||
- "wrap_modules should only be called once"
|
||||
- )
|
||||
+ modules = list(modules_generator)
|
||||
|
||||
if prefix:
|
||||
- prefix = f"{prefix}."
|
||||
+ # Routed by _mark_tower_model: a whole tower module, handled by
|
||||
+ # the streaming path, independent of the decoder group schedule.
|
||||
+ if self.offload_towers:
|
||||
+ for module in modules:
|
||||
+ self._wrap_tower(module, prefix)
|
||||
+ return modules
|
||||
+
|
||||
+ if self.group_size <= 0:
|
||||
+ # Prefetch backend selected for towers only; decoder layers
|
||||
+ # stay resident on the device.
|
||||
+ return modules
|
||||
+
|
||||
+ assert len(self.module_offloaders) == 0, (
|
||||
+ "wrap_modules should only be called once for the decoder stack"
|
||||
+ )
|
||||
|
||||
all_modules = []
|
||||
offload_modules = []
|
||||
|
||||
- for module_index, module in enumerate(modules_generator):
|
||||
+ for module_index, module in enumerate(modules):
|
||||
all_modules.append(module)
|
||||
|
||||
# Select layers to offload based on group pattern
|
||||
@@ -227,6 +248,39 @@ class PrefetchOffloader(BaseOffloader):
|
||||
|
||||
return all_modules
|
||||
|
||||
+ def _wrap_tower(self, tower: nn.Module, prefix: str):
|
||||
+ """Split a tower into parameter-bearing units and place each one
|
||||
+ in pinned host RAM until its forward executes (bulk-copy).
|
||||
+ """
|
||||
+ first_param = next(tower.parameters(), None)
|
||||
+ if first_param is None:
|
||||
+ # Tower was staged out (e.g. --limit-mm-per-prompt zero for all
|
||||
+ # its modalities) or carries no parameters.
|
||||
+ return
|
||||
+ device = first_param.device
|
||||
+ if device.type == "cpu":
|
||||
+ return
|
||||
+
|
||||
+ num_units = 0
|
||||
+ for submodule in tower.modules():
|
||||
+ param_names = [
|
||||
+ name for name, _ in submodule.named_parameters(recurse=False)
|
||||
+ ]
|
||||
+ if not param_names:
|
||||
+ continue
|
||||
+ streamer = _TowerUnitStreamer(submodule, param_names, device)
|
||||
+ streamer.offload_to_host_now()
|
||||
+ self.tower_streamers.append(streamer)
|
||||
+ num_units += 1
|
||||
+
|
||||
+ logger.info(
|
||||
+ "[PrefetchOffloader] Queued tower '%s' for host streaming: "
|
||||
+ "%d parameter units (weights move to the GPU only during "
|
||||
+ "encoder forwards).",
|
||||
+ prefix,
|
||||
+ num_units,
|
||||
+ )
|
||||
+
|
||||
def _hook_module_forward(self, index: int, module: nn.Module):
|
||||
"""Hook module's forward with torch.compile-compatible sync."""
|
||||
original_forward = module.forward
|
||||
@@ -339,6 +393,21 @@ class PrefetchOffloader(BaseOffloader):
|
||||
(in _CpuParamOffloader.__init__), so GPU memory is available for the
|
||||
static buffer pool.
|
||||
"""
|
||||
+ # Tower streamers: re-snapshot host storage from the final processed
|
||||
+ # weights, then install the streaming forwards.
|
||||
+ for streamer in self.tower_streamers:
|
||||
+ streamer.sync_storage()
|
||||
+ for streamer in self.tower_streamers:
|
||||
+ streamer.hook_forward()
|
||||
+ if self.tower_streamers:
|
||||
+ tower_bytes = sum(s.offloaded_bytes for s in self.tower_streamers)
|
||||
+ self.total_offloaded_bytes += tower_bytes
|
||||
+ logger.info_once(
|
||||
+ f"[PrefetchOffloader] Tower offload active: "
|
||||
+ f"{len(self.tower_streamers)} units, "
|
||||
+ f"{tower_bytes / 1e9:.4f} GB in pinned host RAM."
|
||||
+ )
|
||||
+
|
||||
# Sync CPU storage with current param.data BEFORE collecting param info.
|
||||
# This is needed because process_weights_after_loading may have:
|
||||
# 1. Transformed weights (quantization, transpose, etc.)
|
||||
@@ -390,6 +459,136 @@ class PrefetchOffloader(BaseOffloader):
|
||||
self.module_offloaders[i].start_onload_to_static()
|
||||
|
||||
|
||||
+class _TowerUnitStreamer:
|
||||
+ """Bulk-copy placement for one parameter-bearing submodule of a tower.
|
||||
+
|
||||
+ The unit's parameters live in pinned host memory between forwards.
|
||||
+ On each forward, they are bulk-copied to freshly allocated GPU tensors
|
||||
+ on the compute stream (so copies are stream-ordered ahead of the unit's
|
||||
+ kernels with no explicit sync) and released right after. At steady
|
||||
+ state only the unit currently executing holds VRAM, and freed blocks
|
||||
+ are recycled by the caching allocator for the next unit's weights.
|
||||
+
|
||||
+ This is the bulk-copy alternative to UVA zero-copy for multimodal
|
||||
+ towers: whole-tensor H2D transfers instead of operand tiles re-read
|
||||
+ over PCIe inside the GEMM inner loop.
|
||||
+
|
||||
+ Interaction notes:
|
||||
+ - If an encoder CUDA graph captures a streamed unit, correctness is
|
||||
+ preserved (the memcpy is captured and replays from host memory), but
|
||||
+ the captured allocations persist in the graph pool, so the VRAM
|
||||
+ saving is lost. Prefer eager encoder execution (`cudagraph_mm_encoder`
|
||||
+ off) together with tower offloading.
|
||||
+ - Sleep/wake weight reload is not intercepted; storage is re-synced
|
||||
+ from `param.data` at `post_init` only.
|
||||
+ """
|
||||
+
|
||||
+ def __init__(self, unit: nn.Module, param_names: list[str], device: Any):
|
||||
+ self.unit = unit
|
||||
+ self.param_names = param_names
|
||||
+ self.device = device
|
||||
+ self.offloaded_bytes = 0
|
||||
+ # name -> pinned host tensor holding the weights between forwards.
|
||||
+ # Populated/refreshed in sync_storage() (after weight loading and
|
||||
+ # process_weights_after_loading), which is the authoritative point.
|
||||
+ self._storages: dict[str, torch.Tensor] = {}
|
||||
+
|
||||
+ def _get_param(self, name: str) -> nn.Parameter:
|
||||
+ return getattr(self.unit, name)
|
||||
+
|
||||
+ @staticmethod
|
||||
+ def _empty_cpu_like(data: torch.Tensor, pin_memory: bool) -> torch.Tensor:
|
||||
+ return torch.empty_strided(
|
||||
+ size=data.size(),
|
||||
+ stride=data.stride(),
|
||||
+ dtype=data.dtype,
|
||||
+ layout=data.layout,
|
||||
+ device="cpu",
|
||||
+ pin_memory=pin_memory,
|
||||
+ )
|
||||
+
|
||||
+ def offload_to_host_now(self) -> None:
|
||||
+ """Move current parameter data to pinned host storage during
|
||||
+ construction, so weight loading writes to host memory and the
|
||||
+ device never accumulates tower weights.
|
||||
+ """
|
||||
+ pin_memory = should_pin_memory()
|
||||
+ for name in self.param_names:
|
||||
+ p = self._get_param(name)
|
||||
+ if p.device.type == "cpu":
|
||||
+ continue
|
||||
+ cpu_data = self._empty_cpu_like(p.data, pin_memory)
|
||||
+ cpu_data.copy_(p.data)
|
||||
+ p.data = cpu_data
|
||||
+
|
||||
+ def sync_storage(self) -> None:
|
||||
+ """Re-snapshot host storage from param.data.
|
||||
+
|
||||
+ Called from post_init() after weight loading and
|
||||
+ process_weights_after_loading. device_loading_context may have
|
||||
+ replaced param.data with fresh (already pinned) CPU tensors, and
|
||||
+ quantization processing may have left data on the device; handle
|
||||
+ both so _storages always references the final processed weights.
|
||||
+ """
|
||||
+ self.offloaded_bytes = 0
|
||||
+ for name in list(self.param_names):
|
||||
+ try:
|
||||
+ p = self._get_param(name)
|
||||
+ except AttributeError:
|
||||
+ # Parameter deleted by process_weights_after_loading.
|
||||
+ self.param_names.remove(name)
|
||||
+ continue
|
||||
+ data = p.data
|
||||
+ if data.device.type != "cpu":
|
||||
+ cpu_data = self._empty_cpu_like(data, pin_memory=should_pin_memory())
|
||||
+ cpu_data.copy_(data)
|
||||
+ data = cpu_data
|
||||
+ p.data = data
|
||||
+ elif should_pin_memory() and not data.is_pinned():
|
||||
+ pinned = self._empty_cpu_like(data, pin_memory=True)
|
||||
+ pinned.copy_(data)
|
||||
+ data = pinned
|
||||
+ p.data = data
|
||||
+ self._storages[name] = data
|
||||
+ self.offloaded_bytes += data.numel() * data.element_size()
|
||||
+
|
||||
+ def hook_forward(self) -> None:
|
||||
+ if getattr(self.unit, "_vllm_tower_streamed", False):
|
||||
+ return # post_init ran more than once
|
||||
+ original_forward = self.unit.forward
|
||||
+ streamer = self
|
||||
+
|
||||
+ def forward(*args, **kwargs):
|
||||
+ if torch.cuda.is_current_stream_capturing():
|
||||
+ logger.warning_once(
|
||||
+ "[_TowerUnitStreamer] Streaming a tower unit during "
|
||||
+ "CUDA graph capture: results stay correct but the "
|
||||
+ "weights will be captured into the graph pool, so no "
|
||||
+ "VRAM is saved. Disable encoder CUDA graphs "
|
||||
+ "(cudagraph_mm_encoder) to keep the saving."
|
||||
+ )
|
||||
+ for name in streamer.param_names:
|
||||
+ storage = streamer._storages.get(name)
|
||||
+ if storage is None:
|
||||
+ continue
|
||||
+ # Async H2D on the compute stream: ordered ahead of this
|
||||
+ # unit's kernels, no explicit synchronization needed.
|
||||
+ streamer._get_param(name).data = storage.to(
|
||||
+ streamer.device, non_blocking=True
|
||||
+ )
|
||||
+ try:
|
||||
+ return original_forward(*args, **kwargs)
|
||||
+ finally:
|
||||
+ for name in streamer.param_names:
|
||||
+ storage = streamer._storages.get(name)
|
||||
+ if storage is None:
|
||||
+ continue
|
||||
+ streamer._get_param(name).data = storage
|
||||
+
|
||||
+ self.unit._vllm_tower_streamed = True
|
||||
+ self.unit.forward = forward
|
||||
+
|
||||
+
|
||||
class _ModuleOffloader:
|
||||
"""Manages offloading for a single module.
|
||||
|
||||
--
|
||||
2.47.3
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
From 5441fcdcddfc11eb10ff1319a53e0613d6f46705 Mon Sep 17 00:00:00 2001
|
||||
From: MOVZX <movzx@yahoo.com>
|
||||
Date: Fri, 16 Oct 2026 16:07:53 +0700
|
||||
Subject: [PATCH 105/105] [Offloader] Stream every multimodal tower parameter to host
|
||||
RAM
|
||||
|
||||
Qwen3_VisionTransformer derived its compute device from
|
||||
patch_embed.proj.weight and fed pos_embed.weight straight into the
|
||||
position-embedding interpolation kernel, so the last two anchor weights
|
||||
had to stay in VRAM for a tower to be streamable at all.
|
||||
|
||||
Record the device the tower was built on before the offloader moves
|
||||
anything, and fall back to it when no weight is on the accelerator.
|
||||
Move the interpolation table to that device per encode, which costs one
|
||||
8.4 MiB host-to-device copy for a 456 MiB tower. Towers now leave no
|
||||
parameters in VRAM between encodes; only small buffers such as the fused
|
||||
input norm and rotary inv_freq remain.
|
||||
|
||||
Also let a tower restrict streaming through tower_offload_children, and
|
||||
warn when a tower resolves its device property to another device type,
|
||||
so a tower that reads an anchor weight reports the cause instead of a
|
||||
Triton pointer error.
|
||||
|
||||
Co-authored-by: Hinata <hinata@idihore.id>
|
||||
Signed-off-by: MOVZX <movzx@yahoo.com>
|
||||
---
|
||||
vllm/model_executor/models/qwen3_vl.py | 26 ++++++++++-
|
||||
vllm/model_executor/offloader/prefetch.py | 54 ++++++++++++++++++-----
|
||||
2 files changed, 66 insertions(+), 14 deletions(-)
|
||||
|
||||
diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py
|
||||
index 644fd13..b09e872 100644
|
||||
--- a/vllm/model_executor/models/qwen3_vl.py
|
||||
+++ b/vllm/model_executor/models/qwen3_vl.py
|
||||
@@ -666,13 +666,25 @@ class Qwen3_VisionTransformer(nn.Module):
|
||||
]
|
||||
)
|
||||
|
||||
+ # An offloader moves tower parameters to host RAM after this
|
||||
+ # constructor returns, and no weight is left on the accelerator to
|
||||
+ # read the compute device from. Record it here while they are still
|
||||
+ # placed, so `device` keeps working once every weight is streamed.
|
||||
+ first_param = next(self.parameters(), None)
|
||||
+ self._compute_device = (
|
||||
+ torch.device("cpu") if first_param is None else first_param.device
|
||||
+ )
|
||||
+
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return self.patch_embed.proj.weight.dtype
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
- return self.patch_embed.proj.weight.device
|
||||
+ weight_device = self.patch_embed.proj.weight.device
|
||||
+ if weight_device.type != "cpu":
|
||||
+ return weight_device
|
||||
+ return self._compute_device
|
||||
|
||||
def rot_pos_ids(self, h: int, w: int, spatial_merge_size: int) -> torch.Tensor:
|
||||
cache_key = (h, w, spatial_merge_size)
|
||||
@@ -739,7 +751,7 @@ class Qwen3_VisionTransformer(nn.Module):
|
||||
for t, h, w in grid_thw:
|
||||
outputs.append(
|
||||
interpolate_fn(
|
||||
- self.pos_embed.weight,
|
||||
+ self.pos_embed_table(),
|
||||
t,
|
||||
h,
|
||||
w,
|
||||
@@ -750,6 +762,16 @@ class Qwen3_VisionTransformer(nn.Module):
|
||||
)
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
+ def pos_embed_table(self) -> torch.Tensor:
|
||||
+ """The position embedding table on the compute device. The whole tower
|
||||
+ may be streamed from host RAM, but the interpolation kernel cannot read
|
||||
+ host memory.
|
||||
+ """
|
||||
+ table = self.pos_embed.weight
|
||||
+ if table.device.type == "cpu":
|
||||
+ table = table.to(self.device, non_blocking=True)
|
||||
+ return table
|
||||
+
|
||||
def prepare_encoder_metadata(
|
||||
self,
|
||||
grid_thw_list: list[list[int]],
|
||||
diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py
|
||||
index f276673..cd70138 100644
|
||||
--- a/vllm/model_executor/offloader/prefetch.py
|
||||
+++ b/vllm/model_executor/offloader/prefetch.py
|
||||
@@ -251,6 +251,12 @@ class PrefetchOffloader(BaseOffloader):
|
||||
def _wrap_tower(self, tower: nn.Module, prefix: str):
|
||||
"""Split a tower into parameter-bearing units and place each one
|
||||
in pinned host RAM until its forward executes (bulk-copy).
|
||||
+
|
||||
+
|
||||
+ A tower restricts this to its heavy compute subtrees by declaring
|
||||
+ `tower_offload_children`. Tower code reads a small anchor weight to find
|
||||
+ its compute device, and passes embedding weights straight to a kernel,
|
||||
+ so weights outside that list stay on the device.
|
||||
"""
|
||||
first_param = next(tower.parameters(), None)
|
||||
if first_param is None:
|
||||
@@ -261,26 +267,50 @@ class PrefetchOffloader(BaseOffloader):
|
||||
if device.type == "cpu":
|
||||
return
|
||||
|
||||
+ declared = getattr(tower, "tower_offload_children", None)
|
||||
+ children = dict(tower.named_children())
|
||||
+ if declared is None:
|
||||
+ roots: list[nn.Module] = [tower]
|
||||
+ kept: list[str] = []
|
||||
+ else:
|
||||
+ roots = [c for name, c in children.items() if name in declared]
|
||||
+ kept = [name for name in children if name not in declared]
|
||||
+
|
||||
num_units = 0
|
||||
- for submodule in tower.modules():
|
||||
- param_names = [
|
||||
- name for name, _ in submodule.named_parameters(recurse=False)
|
||||
- ]
|
||||
- if not param_names:
|
||||
- continue
|
||||
- streamer = _TowerUnitStreamer(submodule, param_names, device)
|
||||
- streamer.offload_to_host_now()
|
||||
- self.tower_streamers.append(streamer)
|
||||
- num_units += 1
|
||||
+ for root in roots:
|
||||
+ for submodule in root.modules():
|
||||
+ param_names = [
|
||||
+ name for name, _ in submodule.named_parameters(recurse=False)
|
||||
+ ]
|
||||
+ if not param_names:
|
||||
+ continue
|
||||
+ streamer = _TowerUnitStreamer(submodule, param_names, device)
|
||||
+ streamer.offload_to_host_now()
|
||||
+ self.tower_streamers.append(streamer)
|
||||
+ num_units += 1
|
||||
|
||||
logger.info(
|
||||
"[PrefetchOffloader] Queued tower '%s' for host streaming: "
|
||||
- "%d parameter units (weights move to the GPU only during "
|
||||
- "encoder forwards).",
|
||||
+ "%d parameter units, %d children kept on device (%s) (weights "
|
||||
+ "move to the GPU only during encoder forwards).",
|
||||
prefix,
|
||||
num_units,
|
||||
+ len(kept),
|
||||
+ ", ".join(kept) or "none",
|
||||
)
|
||||
|
||||
+ if isinstance(getattr(type(tower), "device", None), property):
|
||||
+ resolved = getattr(tower, "device", None)
|
||||
+ if isinstance(resolved, torch.device) and resolved.type != device.type:
|
||||
+ logger.warning(
|
||||
+ "[PrefetchOffloader] Tower '%s' resolves its compute "
|
||||
+ "device to %s from offloaded weights, so code that calls "
|
||||
+ ".to(tower.device) runs on the host. Add the anchor "
|
||||
+ "submodule to tower_offload_children.",
|
||||
+ prefix,
|
||||
+ resolved,
|
||||
+ )
|
||||
+
|
||||
def _hook_module_forward(self, index: int, module: nn.Module):
|
||||
"""Hook module's forward with torch.compile-compatible sync."""
|
||||
original_forward = module.forward
|
||||
--
|
||||
2.47.3
|
||||
|
||||
@@ -265,7 +265,7 @@ exec ${topo_launch_prefix:+$topo_launch_prefix} "$VENV/bin/vllm" serve "$CKPT" \
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}' \
|
||||
--max-model-len 262144 --max-num-seqs 4 --max-num-batched-tokens 2048 \
|
||||
--gpu-memory-utilization 0.96 --kv-cache-memory 4100000000 --kv-cache-dtype fp8_e4m3 \
|
||||
--limit-mm-per-prompt '{"image":42,"video":0}' --mm-processor-cache-gb 0 \
|
||||
--offload-towers --limit-mm-per-prompt '{"image":42,"video":0}' --mm-processor-cache-gb 16 \
|
||||
--mm-processor-kwargs '{"size":{"longest_edge":4194304,"shortest_edge":65536}}' \
|
||||
--compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,2,4,8,12,16]}' \
|
||||
--enable-prefix-caching --enable-prompt-tokens-details \
|
||||
|
||||
Reference in New Issue
Block a user