From 5441fcdcddfc11eb10ff1319a53e0613d6f46705 Mon Sep 17 00:00:00 2001 From: MOVZX 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 Signed-off-by: MOVZX --- 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