test_toy_llama.py 15.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
4
"""
Test the piecewise compilation with a simple model, comparing the output
with and without the piecewise compilation.
5
6
7
8

This is a tractable model, the weights and computation are specially designed
if the config `tractable_init` is set to True. Otherwise, the weights are
initialized randomly with a fixed seed.
9
10
"""
from dataclasses import dataclass
11
from typing import Any, Optional
12
13
14

import torch
from torch import nn
15
from torch.library import Library
16
17
18

from vllm.compilation.counter import compilation_counter
from vllm.compilation.decorators import support_torch_compile
19
20
from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
                         set_current_vllm_config)
21
22
23
24
from vllm.utils import direct_register_custom_op

# create a library to hold the custom op
silly_lib = Library("silly", "FRAGMENT")  # noqa
25
26
27
28
29
30
31
32
33


def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                    out: torch.Tensor) -> None:
    out.copy_(q)
    out += k
    out += v


34
35
def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                         out: torch.Tensor) -> None:
36
37
38
    return


39
40
41
42
43
44
45
46
47
direct_register_custom_op(
    op_name="attention",
    op_func=silly_attention,
    mutates_args=["out"],
    fake_impl=silly_attention_fake,
    target_lib=silly_lib,
)


48
49
50
51
52
53
@dataclass
class LlamaConfig:
    hidden_size: int = 128
    mlp_size: int = 256
    vocab_size: int = 128
    num_layers: int = 2
54
55
56
57
    init_value: float = 1.0
    tractable_init: bool = False
    random_seed: int = 0

58
    def compute_hash(self) -> str:
59
        factors: list[Any] = []
60
61
62
63
64
65
        for k, v in self.__dict__.items():
            if k == "random_seed":
                continue
            factors.append((k, v))
        factors.sort()
        import hashlib
66
67
        return hashlib.md5(str(factors).encode(),
                           usedforsecurity=False).hexdigest()
68

69
70
    def __post_init__(self):
        assert self.mlp_size >= self.hidden_size
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87


class LlamaMLP(nn.Module):

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.gate_up_projection = nn.Linear(
            in_features=config.hidden_size,
            out_features=config.mlp_size * 2,
            bias=False,
        )
        self.down_projection = nn.Linear(
            in_features=config.mlp_size,
            out_features=config.hidden_size,
            bias=False,
        )

88
89
90
91
92
93
94
95
96
97
98
99
100
        if config.tractable_init:
            nn.init.eye_(self.gate_up_projection.weight.data[:config.mlp_size])
            nn.init.eye_(self.gate_up_projection.weight.data[config.mlp_size:])
            nn.init.eye_(self.down_projection.weight.data)
        else:
            nn.init.xavier_normal_(self.gate_up_projection.weight.data,
                                   generator=torch.Generator().manual_seed(
                                       config.random_seed),
                                   gain=0.001)
            nn.init.xavier_normal_(self.down_projection.weight.data,
                                   generator=torch.Generator().manual_seed(
                                       config.random_seed),
                                   gain=0.001)
101
102

    def forward(self, x):
103
104
        # for tractable_init and positive input, this is
        # essentially an elementwise-square
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        x = self.gate_up_projection(x)
        x = x[:, :x.size(1) // 2] * torch.nn.functional.relu(
            x[:, x.size(1) // 2:])
        x = self.down_projection(x)
        return x


class LlamaAttention(nn.Module):

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.qkv_projection = nn.Linear(
            in_features=config.hidden_size,
            out_features=config.hidden_size * 3,
119
            bias=False,
120
121
122
123
124
        )

        self.output_projection = nn.Linear(
            in_features=config.hidden_size,
            out_features=config.hidden_size,
125
            bias=False,
126
127
        )

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
        if config.tractable_init:
            nn.init.eye_(self.qkv_projection.weight.data[:config.hidden_size])
            nn.init.eye_(self.qkv_projection.weight.data[config.hidden_size:2 *
                                                         config.hidden_size])
            nn.init.eye_(self.qkv_projection.weight.data[2 *
                                                         config.hidden_size:])
            nn.init.eye_(self.output_projection.weight.data)
        else:
            nn.init.xavier_normal_(self.qkv_projection.weight.data,
                                   generator=torch.Generator().manual_seed(
                                       config.random_seed),
                                   gain=0.001)
            nn.init.xavier_normal_(self.output_projection.weight.data,
                                   generator=torch.Generator().manual_seed(
                                       config.random_seed),
                                   gain=0.001)
144
145
146
147
148
149

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
150
151
        # for tractable_init, this is:
        # output = (hidden_states * 3 + positions * 2)
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
        qkv = self.qkv_projection(hidden_states)
        hidden_size = qkv.size(-1) // 3
        q, k, v = qkv.split([hidden_size, hidden_size, hidden_size], dim=-1)

        q = q + positions.unsqueeze(1)
        k = k + positions.unsqueeze(1)

        attn_output = torch.empty_like(q)
        torch.ops.silly.attention(q, k, v, attn_output)

        output = self.output_projection(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

    def __init__(self, config: LlamaConfig) -> None:
        super().__init__()
        self.self_attention = LlamaAttention(config)
        self.mlp = LlamaMLP(config)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
178
    ) -> tuple[torch.Tensor, torch.Tensor]:
179
180
181
182
183
184
185
186
187
        """
        For tractable computation:
        - if residual is None, the outputs are:
            - residual = (hidden_states + 1) * 3 + positions * 2 + hidden_states = hidden_states * 4 + positions * 2 + 3
            - hidden_states = (residual + 1) ** 2
        - if residual is not None, the outputs are:
            - residual = (hidden_states + residual + 1) * 3 + positions * 2 + hidden_states + residual = (hidden_states + residual) * 4 + positions * 2 + 3
            - hidden_states = (residual + 1) ** 2
        """ # noqa
188
189
        if residual is None:
            residual = hidden_states
190
            hidden_states = hidden_states + 1
191
192
193
        else:
            hidden_states = hidden_states + residual
            residual = hidden_states
194
            hidden_states = hidden_states + 1
195
196
197
198
199
200

        hidden_states = self.self_attention(positions=positions,
                                            hidden_states=hidden_states)

        hidden_states = hidden_states + residual
        residual = hidden_states
201
        hidden_states = hidden_states + 1
202
203
204
205
206
        hidden_states = self.mlp(hidden_states)

        return hidden_states, residual


207
@support_torch_compile
208
209
class LlamaModel(nn.Module):

210
211
212
213
214
215
    def __init__(self,
                 *,
                 vllm_config: VllmConfig,
                 config: LlamaConfig,
                 prefix: str = '',
                 **kwargs) -> None:
216
217
218
219
220
221
222
223
        super().__init__()
        self.embedding_tokens = nn.Embedding(
            num_embeddings=config.vocab_size,
            embedding_dim=config.hidden_size,
        )
        self.layers = nn.ModuleList(
            [LlamaDecoderLayer(config) for _ in range(config.num_layers)])

224
225
        # this is the initial value of the hidden states
        self.embedding_tokens.weight.data.fill_(config.init_value)
226
227
228
229
230
231
232
233
234
235
236
237
238

    def forward(
        self,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = self.embedding_tokens(input_ids)
        residual = None
        for layer in self.layers:
            hidden_states, residual = layer(positions, hidden_states, residual)
        return hidden_states


239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
def tractable_computation(input_ids: torch.Tensor,
                          positions: torch.Tensor,
                          config: LlamaConfig,
                          init_value: float = 1.0) -> torch.Tensor:
    hidden_states = torch.ones(input_ids.size(0),
                               config.hidden_size,
                               device=input_ids.device,
                               dtype=input_ids.dtype) * init_value

    # first layer
    residual = hidden_states * 4 + positions.unsqueeze(1) * 2 + 3
    hidden_states = (residual + 1)**2

    # following layers
    for _ in range(config.num_layers - 1):
        hidden_states = hidden_states + residual
        residual = hidden_states * 4 + positions.unsqueeze(1) * 2 + 3
        hidden_states = (residual + 1)**2

    return hidden_states


261
262
263
264
265
266
@torch.inference_mode
def run_model(llama_config,
              use_compile: bool,
              split_attn: bool = False) -> torch.Tensor:

    if use_compile:
267
268
269
        compilation_config = CompilationConfig(
            level=CompilationLevel.PIECEWISE,
            use_cudagraph=True,
270
            cudagraph_capture_sizes=[1, 2],
271
        )
272
        if split_attn:
273
            compilation_config.splitting_ops = ["silly.attention"]
274
    else:
275
276
        compilation_config = CompilationConfig(
            level=CompilationLevel.NO_COMPILATION, )
277

278
279
    vllm_config = VllmConfig(compilation_config=compilation_config,
                             additional_config=llama_config)
280
281
282
283
    with set_current_vllm_config(vllm_config):
        model = LlamaModel(config=llama_config,
                           vllm_config=vllm_config,
                           prefix="").eval().cuda()
284
285
286
287
288

    B = 16  # max batch size
    input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
    positions = torch.arange(B).cuda()

289
290
291
    model(input_ids, positions)
    model(input_ids[:2], positions[:2])
    model(input_ids[:1], positions[:1])
292
293
294
295

    input_ids[:2].zero_()
    output = model(input_ids[:2], positions[:2])

296
297
298
299
300
301
302
303
304
    output = output.cpu()

    if llama_config.tractable_init:
        expected_output = tractable_computation(input_ids[:2], positions[:2],
                                                llama_config).cpu()

        assert torch.allclose(output, expected_output)
    else:
        return output.cpu()
305
306
307
308
309
310
311
312


def test_toy_llama():
    # compare output with and without piecewise compilation

    llama_config = LlamaConfig(hidden_size=128,
                               mlp_size=256,
                               vocab_size=128,
313
314
315
316
317
318
319
                               num_layers=12)

    tractable_config = LlamaConfig(hidden_size=128,
                                   mlp_size=256,
                                   vocab_size=128,
                                   num_layers=2,
                                   tractable_init=True)
320
321
322
323
324
325

    outputs = []
    with compilation_counter.expect(
            num_graphs_seen=0,
            num_piecewise_graphs_seen=0,
            num_piecewise_capturable_graphs_seen=0,
326
            num_backend_compilations=0,
327
328
329
            num_cudagraph_caputured=0,
    ):
        outputs.append(run_model(llama_config, use_compile=False))
330
331
    run_model(tractable_config, use_compile=False)

332
333
334
335
    with compilation_counter.expect(
            num_graphs_seen=1,  # one graph for the model
            num_piecewise_graphs_seen=1,
            num_piecewise_capturable_graphs_seen=1,
336
            num_backend_compilations=1,  # num_piecewise_capturable_graphs_seen
337
338
339
340
            num_cudagraph_caputured=
            2,  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
    ):
        outputs.append(run_model(llama_config, use_compile=True))
341
    run_model(tractable_config, use_compile=True)
342
343
344
345
346
347
348

    with compilation_counter.expect(
            num_graphs_seen=1,  # one graph for the model
            num_piecewise_graphs_seen=2 * llama_config.num_layers +
            1,  # 2 * num_layers + 1
            num_piecewise_capturable_graphs_seen=1 +
            llama_config.num_layers,  # 1 + num_layers
349
            num_backend_compilations=1 +
350
351
352
353
354
355
356
            llama_config.num_layers,  # num_piecewise_capturable_graphs_seen
            num_cudagraph_caputured=2 *
        (1 + llama_config.num_layers
         ),  # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
    ):
        outputs.append(
            run_model(llama_config, use_compile=True, split_attn=True))
357
    run_model(tractable_config, use_compile=True, split_attn=True)
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389

    for i in range(1, len(outputs)):
        assert torch.allclose(outputs[0], outputs[i])


@torch.inference_mode
def benchmark():
    from triton.testing import do_bench

    # similar to llama 3.1-8B
    llama_config = LlamaConfig(hidden_size=4096,
                               mlp_size=14336,
                               vocab_size=128 * 1024,
                               num_layers=32)

    # a tiny model to measure the overhead
    # of piecewise cudagraph
    llama_config = LlamaConfig(hidden_size=40,
                               mlp_size=80,
                               vocab_size=128,
                               num_layers=2)

    cudagraph_sizes = [1, 2, 4] + [i * 8 for i in range(1, 33)]

    eager_time = {}
    full_cudagraph_time = {}
    piecewise_cudagraph_time = {}

    pool = torch.cuda.graph_pool_handle()

    for piecewise in [False, True]:
        if piecewise:
390
391
392
            compilation_config = CompilationConfig(
                level=CompilationLevel.PIECEWISE,
                use_cudagraph=True,
393
                splitting_ops=["silly.attention"],
394
                cudagraph_capture_sizes=cudagraph_sizes,
395
            )
396
        else:
397
            compilation_config = CompilationConfig(
398
399
400
                level=CompilationLevel.PIECEWISE,
                cudagraph_capture_sizes=cudagraph_sizes,
            )
401

402
        vllm_config = VllmConfig(compilation_config=compilation_config)
403
404
405
406
        with set_current_vllm_config(vllm_config):
            model = LlamaModel(config=llama_config,
                               vllm_config=vllm_config,
                               prefix="").eval().cuda().to(torch.bfloat16)
407
408
409
410
411
412
413

        B = 256  # max batch size
        input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
        positions = torch.arange(B).cuda().to(torch.bfloat16)

        graphs = {}

414
415
416
417
418
        model(input_ids, positions)
        for b in cudagraph_sizes[::-1]:
            if not piecewise:
                graph = torch.cuda.CUDAGraph()
                with torch.cuda.graph(graph, pool=pool):
419
                    output = model(input_ids[:b], positions[:b])
420
421
422
423
                graphs[b] = (graph, output)
            else:
                output = model(input_ids[:b], positions[:b])
                graphs[b] = (model, output)
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
        for b in cudagraph_sizes:
            if piecewise:
                # noqa is for `Function definition does not bind loop variable`
                # it will be problematic if we save the created lambda function
                # and use it later, because it will look up the name `b` in the
                # enclosing scope, and the value of `b` will always be 256.
                # it is fine here, because we only use the lambda function once.
                runtime = do_bench(lambda: graphs[b][0]  # noqa
                                   (input_ids[:b], positions[:b]))  # noqa
                piecewise_cudagraph_time[b] = runtime
            else:
                runtime = do_bench(lambda: graphs[b][0].replay())  # noqa
                eager_runtime = do_bench(
                    lambda: model(input_ids[:b], positions[:b]))  # noqa
                full_cudagraph_time[b] = runtime
                eager_time[b] = eager_runtime

    # print in tabular format
    print("batch size\teager mode\tfull cudagraph\tpiecewise cudagraph")
    for b in cudagraph_sizes:
444
445
        print(f"{b}\t{eager_time[b]:.3f}\t{full_cudagraph_time[b]:.3f}"
              f"\t{piecewise_cudagraph_time[b]:.3f}")
446
447
448
449


if __name__ == "__main__":
    benchmark()