MOVZXandHinata
18c06ed406
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>
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# 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,
|
|
_TowerUnitStreamer,
|
|
)
|
|
|
|
pytestmark = pytest.mark.skip_global_cleanup
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("module_count", "prefetch_step", "expected_targets"),
|
|
[
|
|
(3, 1, [1, 2, 0]),
|
|
(3, 2, [2, 1, 0]),
|
|
(4, 2, [2, 3, 0, 1]),
|
|
(5, 2, [2, 3, 4, 1, 0]),
|
|
(3, 3, [0, 1, 2]),
|
|
(3, 4, [1, 2, 0]),
|
|
],
|
|
)
|
|
def test_next_prefetch_index_preserves_slot_ownership(
|
|
module_count: int,
|
|
prefetch_step: int,
|
|
expected_targets: list[int],
|
|
) -> None:
|
|
targets = [
|
|
_get_next_prefetch_index(index, prefetch_step, module_count)
|
|
for index in range(module_count)
|
|
]
|
|
|
|
assert targets == expected_targets
|
|
if prefetch_step < module_count:
|
|
assert all(
|
|
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
|