diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 527e2b0f3e..c9e3c3151c 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -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]], diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index abde223d80..722eae6144 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -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