iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Offloader] Stream every multimodal tower parameter to host RAM

Browse Source 
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>
This commit is contained in:
MOVZX
2026-10-16 16:07:53 +07:00
co-authored by Hinata
parent 18c06ed406
commit 8dc54f35dd
2 changed files with 65 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/models/qwen3_vl.py
+24 -2
View File
@@ -676,13 +676,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)
@@ -751,7 +763,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,
@@ -762,6 +774,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: Sequence[Sequence[int]],
vllm/model_executor/offloader/prefetch.py
+41 -12
View File
@@ -252,6 +252,11 @@ 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:
@@ -262,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