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

4
import gc
5
import os
6
import pathlib
7
8
9
import subprocess

import pytest
10
import torch
11

zhuwenwen's avatar
zhuwenwen committed
12
from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
13
from vllm.engine.arg_utils import EngineArgs
14
# yapf conflicts with isort for this docstring
15
16
17
18
19
# yapf: disable
from vllm.model_executor.model_loader.tensorizer import (TensorizerConfig,
                                                         TensorSerializer,
                                                         is_vllm_tensorized,
                                                         open_stream,
20
                                                         tensorize_vllm_model)
21
22
from vllm.lora.request import LoRARequest

23
# yapf: enable
24
from vllm.utils import PlaceholderModule
25

zhuwenwen's avatar
zhuwenwen committed
26
from ..utils import VLLM_PATH, models_path_prefix
27

28
29
30
31
32
33
try:
    from tensorizer import EncryptionParams
except ImportError:
    tensorizer = PlaceholderModule("tensorizer")  # type: ignore[assignment]
    EncryptionParams = tensorizer.placeholder_attr("EncryptionParams")

34
EXAMPLES_PATH = VLLM_PATH / "examples"
35

36
37
38
39
40
41
42
43
44
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)

45
model_ref = os.path.join(models_path_prefix, "facebook/opt-125m")
46
47
tensorize_model_for_testing_script = os.path.join(
    os.path.dirname(__file__), "tensorize_vllm_model_for_testing.py")
48

49

50
51
52
53
54
55
56
def is_curl_installed():
    try:
        subprocess.check_call(['curl', '--version'])
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False

57

58
59
60
61
62
def write_keyfile(keyfile_path: str):
    encryption_params = EncryptionParams.random()
    pathlib.Path(keyfile_path).parent.mkdir(parents=True, exist_ok=True)
    with open(keyfile_path, 'wb') as f:
        f.write(encryption_params.key)
63
64
65
66


@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
def test_can_deserialize_s3(vllm_runner):
67
    model_ref = os.path.join(models_path_prefix, "EleutherAI/pythia-1.4b")
zhuwenwen's avatar
zhuwenwen committed
68
    tensorized_path = f"{model_ref}/fp16/model.tensors"
69

70
    with vllm_runner(model_ref,
71
72
73
74
75
76
                     load_format="tensorizer",
                     model_loader_extra_config=TensorizerConfig(
                         tensorizer_uri=tensorized_path,
                         num_readers=1,
                         s3_endpoint="object.ord1.coreweave.com",
                     )) as loaded_hf_model:
77
78
        deserialized_outputs = loaded_hf_model.generate(
            prompts, sampling_params)
79
        # noqa: E501
80

81
        assert deserialized_outputs
82
83
84
85
86


@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
def test_deserialized_encrypted_vllm_model_has_same_outputs(
        vllm_runner, tmp_path):
87
    args = EngineArgs(model=model_ref)
88
89
90
    with vllm_runner(model_ref) as vllm_model:
        model_path = tmp_path / (model_ref + ".tensors")
        key_path = tmp_path / (model_ref + ".key")
91
92
        write_keyfile(key_path)

93
        outputs = vllm_model.generate(prompts, sampling_params)
94

95
96
    config_for_serializing = TensorizerConfig(tensorizer_uri=str(model_path),
                                              encryption_keyfile=str(key_path))
97

98
    tensorize_vllm_model(args, config_for_serializing)
99

100
101
    config_for_deserializing = TensorizerConfig(
        tensorizer_uri=str(model_path), encryption_keyfile=str(key_path))
102

103
104
105
106
    with vllm_runner(model_ref,
                     load_format="tensorizer",
                     model_loader_extra_config=config_for_deserializing
                     ) as loaded_vllm_model:  # noqa: E501
107

108
109
        deserialized_outputs = loaded_vllm_model.generate(
            prompts, sampling_params)
110
        # noqa: E501
111

112
        assert outputs == deserialized_outputs
113
114
115
116


def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
                                                tmp_path):
117
118
119
120
121
122
123
124
    with hf_runner(model_ref) as hf_model:
        model_path = tmp_path / (model_ref + ".tensors")
        max_tokens = 50
        outputs = hf_model.generate_greedy(prompts, max_tokens=max_tokens)
        with open_stream(model_path, "wb+") as stream:
            serializer = TensorSerializer(stream)
            serializer.write_module(hf_model.model)

125
    with vllm_runner(model_ref,
126
127
128
129
130
                     load_format="tensorizer",
                     model_loader_extra_config=TensorizerConfig(
                         tensorizer_uri=model_path,
                         num_readers=1,
                     )) as loaded_hf_model:
131
132
        deserialized_outputs = loaded_hf_model.generate_greedy(
            prompts, max_tokens=max_tokens)
133

134
        assert outputs == deserialized_outputs
135
136


137
def test_load_without_tensorizer_load_format(vllm_runner, capfd):
138
    model = None
139
    try:
140
        model = vllm_runner(
141
142
            model_ref,
            model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
143
144
145
146
147
148
149
150
151
152
153
154
155
    except RuntimeError:
        out, err = capfd.readouterr()
        combined_output = out + err
        assert ("ValueError: Model loader extra config "
                "is not supported for load "
                "format LoadFormat.AUTO") in combined_output
    finally:
        del model
        gc.collect()
        torch.cuda.empty_cache()


def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd):
156
    model = None
157
    try:
158
        model = vllm_runner(
159
160
161
            model_ref,
            load_format="safetensors",
            model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
162
163
164
165
166
167
168
169
170
171
    except RuntimeError:
        out, err = capfd.readouterr()

        combined_output = out + err
        assert ("ValueError: Model loader extra config is not supported "
                "for load format LoadFormat.SAFETENSORS") in combined_output
    finally:
        del model
        gc.collect()
        torch.cuda.empty_cache()
172
173


174
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
175
176
def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
    try:
177
        model_ref = os.path.join(models_path_prefix, "EleutherAI/pythia-1.4b")
zhuwenwen's avatar
zhuwenwen committed
178
        # tensorized_path = f"s3://tensorized/{model_ref}/fp16/model.tensors"
zhuwenwen's avatar
zhuwenwen committed
179
        tensorized_path = f"{model_ref}/fp16/model.tensors"
180
181
182
183

        vllm_runner(
            model_ref,
            load_format="tensorizer",
184
185
186
187
188
            model_loader_extra_config=TensorizerConfig(
                tensorizer_uri=tensorized_path,
                num_readers=1,
                s3_endpoint="object.ord1.coreweave.com",
            ),
189
            tensor_parallel_size=2,
190
            disable_custom_all_reduce=True,
191
        )
192
193
194
195
196
197
198
    except RuntimeError:
        out, err = capfd.readouterr()
        combined_output = out + err
        assert ("ValueError: For a sharded model, tensorizer_uri "
                "should include a string format template like '%04d' "
                "to be formatted with the rank "
                "of the shard") in combined_output
199

200

201
202
203
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
        vllm_runner, tmp_path):
204
    model_ref = os.path.join(models_path_prefix, "EleutherAI/pythia-1.4b")
205
    # record outputs from un-sharded un-tensorized model
206
207
208
209
210
211
    with vllm_runner(
            model_ref,
            disable_custom_all_reduce=True,
            enforce_eager=True,
    ) as base_model:
        outputs = base_model.generate(prompts, sampling_params)
212
213
214
215
216
217
218

    # load model with two shards and serialize with encryption
    model_path = str(tmp_path / (model_ref + "-%02d.tensors"))
    key_path = tmp_path / (model_ref + ".key")

    tensorizer_config = TensorizerConfig(
        tensorizer_uri=model_path,
219
        encryption_keyfile=str(key_path),
220
221
222
223
    )

    tensorize_vllm_model(
        engine_args=EngineArgs(
224
225
226
227
228
            model=model_ref,
            tensor_parallel_size=2,
            disable_custom_all_reduce=True,
            enforce_eager=True,
        ),
229
230
231
232
233
        tensorizer_config=tensorizer_config,
    )
    assert os.path.isfile(model_path % 0), "Serialization subprocess failed"
    assert os.path.isfile(model_path % 1), "Serialization subprocess failed"

234
235
236
237
238
239
240
    with vllm_runner(
            model_ref,
            tensor_parallel_size=2,
            load_format="tensorizer",
            disable_custom_all_reduce=True,
            enforce_eager=True,
            model_loader_extra_config=tensorizer_config) as loaded_vllm_model:
241
242
        deserialized_outputs = loaded_vllm_model.generate(
            prompts, sampling_params)
243
244
245

    assert outputs == deserialized_outputs

246

247
@pytest.mark.flaky(reruns=3)
248
def test_vllm_tensorized_model_has_same_outputs(vllm_runner, tmp_path):
249
250
    gc.collect()
    torch.cuda.empty_cache()
251
    model_ref = os.path.join(models_path_prefix, "facebook/opt-125m")
252
253
    model_path = tmp_path / (model_ref + ".tensors")
    config = TensorizerConfig(tensorizer_uri=str(model_path))
254
    args = EngineArgs(model=model_ref, device="cuda")
255

256
257
    with vllm_runner(model_ref) as vllm_model:
        outputs = vllm_model.generate(prompts, sampling_params)
258

259
260
    tensorize_vllm_model(args, config)
    assert is_vllm_tensorized(config)
261

262
    with vllm_runner(model_ref,
263
264
                     load_format="tensorizer",
                     model_loader_extra_config=config) as loaded_vllm_model:
265
266
        deserialized_outputs = loaded_vllm_model.generate(
            prompts, sampling_params)
267
        # noqa: E501
268

269
        assert outputs == deserialized_outputs