test_swap.py 2.97 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
import torch

from vllm.engine.arg_utils import EngineArgs
7
from vllm.sequence import ExecuteModelRequest
8
from vllm.utils import get_distributed_init_method, get_ip, get_open_port
9
from vllm.worker.worker import Worker
10
11
12
13


def test_swap() -> None:
    # Configure the engine.
14
    engine_args = EngineArgs(model="distilbert/distilgpt2",
15
16
                             dtype="half",
                             load_format="dummy")
17
    engine_config = engine_args.create_engine_config()
18
19
    engine_config.cache_config.num_gpu_blocks = 1000
    engine_config.cache_config.num_cpu_blocks = 1000
20
21
22
23
24

    # Create the worker.
    distributed_init_method = get_distributed_init_method(
        get_ip(), get_open_port())
    worker = Worker(
25
        vllm_config=engine_config,
26
27
28
29
30
31
32
        local_rank=0,
        rank=0,
        distributed_init_method=distributed_init_method,
        is_driver_worker=True,
    )

    # Initialize the worker.
33
    worker.init_device()
34
    worker.load_model()
35
36
37
    worker.initialize_cache(
        num_gpu_blocks=engine_config.cache_config.num_gpu_blocks,
        num_cpu_blocks=engine_config.cache_config.num_cpu_blocks)
38
39

    # Randomly initialize the cache.
40
41
    gpu_cache = worker.cache_engine[0].gpu_cache
    cpu_cache = worker.cache_engine[0].cpu_cache
42
43
44
45
46
47
48
49
50
51
52
53
54
    num_layers = len(gpu_cache)
    for i in range(num_layers):
        gpu_key_cache, gpu_value_cache = gpu_cache[i]
        gpu_key_cache.random_()
        gpu_value_cache.random_()
        cpu_key_cache, cpu_value_cache = cpu_cache[i]
        cpu_key_cache.random_()
        cpu_value_cache.random_()

    allclose = lambda a, b: torch.allclose(
        a.cuda(), b.cuda(), rtol=0.0, atol=0.0)

    # Test swap out.
55
    blocks_to_swap_out = [(3, 72), (56, 35), (84, 34)]
56
57
    execute_model_req = ExecuteModelRequest(
        seq_group_metadata_list=[],
58
        blocks_to_swap_in=[],
59
        blocks_to_swap_out=blocks_to_swap_out,
60
        blocks_to_copy=[],
61
62
63
    )
    worker.execute_model(execute_model_req=execute_model_req)

64
65
66
    for i in range(num_layers):
        gpu_key_cache, gpu_value_cache = gpu_cache[i]
        cpu_key_cache, cpu_value_cache = cpu_cache[i]
67
        for src, dst in blocks_to_swap_out:
68
69
70
71
            assert allclose(gpu_key_cache[src], cpu_key_cache[dst])
            assert allclose(gpu_value_cache[src], cpu_value_cache[dst])

    # Test swap in.
72
73
74
75
76
77
78
79
    execute_model_req.blocks_to_swap_out = []
    execute_model_req.blocks_to_swap_in = [
        (19, 45),
        (67, 23),
        (12, 78),
        (40, 99),
        (1, 71),
    ]
80
81
    worker.execute_model(execute_model_req=execute_model_req)

82
83
84
    for i in range(num_layers):
        gpu_key_cache, gpu_value_cache = gpu_cache[i]
        cpu_key_cache, cpu_value_cache = cpu_cache[i]
85
        for src, dst in execute_model_req.blocks_to_swap_in:
86
87
            assert allclose(gpu_key_cache[dst], cpu_key_cache[src])
            assert allclose(gpu_value_cache[dst], cpu_value_cache[src])