bench.py 16.2 KB
Newer Older
pengcheng888's avatar
pengcheng888 committed
1
2
3
4
import infinicore
from transformers import AutoTokenizer
from infinilm.modeling_utils import load_model_state_dict_by_file
from infinilm.distributed import DistConfig
5
from infinilm.infer_engine import GenerationConfig, InferEngine
6
from infinilm.cache import StaticKVCacheConfig, PagedKVCacheConfig
pengcheng888's avatar
pengcheng888 committed
7
8
9
10
import argparse
import sys
import time
import os
pengcheng888's avatar
pengcheng888 committed
11
12
import json
from collections import OrderedDict
13
import numpy as np
pengcheng888's avatar
pengcheng888 committed
14
from tqdm import tqdm
pengcheng888's avatar
pengcheng888 committed
15
16
17
18

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../python"))


pengcheng888's avatar
pengcheng888 committed
19
20
21
22
23
24
DATA_TYPE_BYTES = {
    "bfloat16": 2,
    "float16": 2,
    "float32": 4,
}

25
26
_PAGED_KV_BLOCK_SIZE = 256

pengcheng888's avatar
pengcheng888 committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
# BATCH_SIZES = [1, 4, 8, 16, 32, 64, 128]
# INPUT_LENS = [32, 256, 1024, 4096]
# OUTPUT_LENS = [256, 1024, 4096]


def read_json_file(file_path):
    """Load and return JSON content from file_path."""
    with open(file_path, "r") as file:
        return json.load(file)


def parse_list(value: str):
    """Parse parse_list argument: can be a single int or a list of ints.

    Examples:
        "1" -> 1
        "[1,2,4]" -> [1, 2, 4]
        "1,2,4" -> [1, 2, 4]
    """
    value = value.strip()
    # Try to parse as JSON list first
    if value.startswith("[") and value.endswith("]"):
        try:
            result = json.loads(value)
            if isinstance(result, list):
                return [int(x) for x in result]
            return int(result)
        except (json.JSONDecodeError, ValueError):
            pass

    # Try to parse as comma-separated values
    if "," in value:
        try:
            return [int(x.strip()) for x in value.split(",")]
        except ValueError:
            pass

    # Try to parse as a single integer
    try:
        return int(value)
    except ValueError:
        raise argparse.ArgumentTypeError(
            f"batch-size must be an int or list[int], got: {value}"
        )


def get_test_cases(
    model_path: str,
    batch_size_list: list[int],
    input_len_list: list[int],
    output_len_list: list[int],
):
    model_path = os.path.expanduser(model_path)

    """Generate cases ordered by ascending KV cache memory usage."""
    # Load model config to derive attention dimensions
    config = read_json_file(os.path.join(model_path, "config.json"))
    head_dim = config.get(
        "head_dim", config.get("hidden_size") // config.get("num_attention_heads")
    )
    # KV heads and layers drive cache size
    num_key_value_heads = config.get("num_key_value_heads")
    num_hidden_layers = config.get("num_hidden_layers")

    # Enumerate all batch/input/output combinations and compute KV cache size
    case_list = []
    for batch_size in batch_size_list:
        for input_len in input_len_list:
            for output_len in output_len_list:
                for data_type in ["bfloat16"]:
                    data_type_bytes = DATA_TYPE_BYTES[data_type]

                    total_seq_len = input_len + output_len
                    kvcache_memory_bytes = (
                        data_type_bytes
                        * (batch_size * total_seq_len * num_key_value_heads * head_dim)
                        * num_hidden_layers
                    )
                    kvcache_memory_gb = kvcache_memory_bytes / (1024 * 1024 * 1024)

                    case_list.append(
                        {
                            "idx": len(case_list),
                            "batch_size": batch_size,
                            "input_len": input_len,
                            "output_len": output_len,
                            "data_type": data_type,
                            "kvcache_memory": round(kvcache_memory_gb, 3),
                        }
                    )

    # Sort by KV cache size and wrap in OrderedDict with index keys
    case_dict = OrderedDict(
        (idx, case)
        for idx, case in enumerate(
            sorted(case_list, key=lambda case: case["kvcache_memory"])
        )
    )

    return case_dict


pengcheng888's avatar
pengcheng888 committed
129
130
131
132
133
134
135
136
137
138
139
140
141
def get_args():
    parser = argparse.ArgumentParser(description="run Llama args")

    parser.add_argument(
        "--cpu",
        action="store_true",
        help="Run cpu test",
    )
    parser.add_argument(
        "--nvidia",
        action="store_true",
        help="Run nvidia test",
    )
142
143
144
145
146
    parser.add_argument(
        "--qy",
        action="store_true",
        help="Run qy test",
    )
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
    parser.add_argument(
        "--metax",
        action="store_true",
        help="Run metax test",
    )
    parser.add_argument(
        "--moore",
        action="store_true",
        help="Run moore test",
    )
    parser.add_argument(
        "--iluvatar",
        action="store_true",
        help="Run iluvatar test",
    )
162
163
164
165
166
    parser.add_argument(
        "--cambricon",
        action="store_true",
        help="Run cambricon test",
    )
wooway777's avatar
wooway777 committed
167
168
169
170
171
    parser.add_argument(
        "--ali",
        action="store_true",
        help="Run alippu test",
    )
172
173
174
175
176
    parser.add_argument(
        "--hygon",
        action="store_true",
        help="Run hygon test",
    )
pengcheng888's avatar
pengcheng888 committed
177
178
179
180
181
182
183
184
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="model path",
    )
    parser.add_argument(
        "--batch-size",
pengcheng888's avatar
pengcheng888 committed
185
        type=parse_list,
pengcheng888's avatar
pengcheng888 committed
186
        default=1,
pengcheng888's avatar
pengcheng888 committed
187
        help="number of prompts in a batch (can be an int or a list of ints, e.g., '1' or '[1,2,4]' or '1,2,4')",
pengcheng888's avatar
pengcheng888 committed
188
189
190
191
192
193
194
195
196
197
    )
    parser.add_argument(
        "--tensor-parallel-size",
        "--tp",
        type=int,
        default=1,
        help="total rank for tensor parallel",
    )
    parser.add_argument(
        "--input-len",
pengcheng888's avatar
pengcheng888 committed
198
199
        type=parse_list,
        default=10,
pengcheng888's avatar
pengcheng888 committed
200
201
202
203
204
        help="output tokens",
    )

    parser.add_argument(
        "--output-len",
pengcheng888's avatar
pengcheng888 committed
205
206
        type=parse_list,
        default=20,
pengcheng888's avatar
pengcheng888 committed
207
208
        help="output tokens",
    )
209
210
211
212
213
    parser.add_argument(
        "--skip-load",
        action="store_true",
        help="skip loading model weights",
    )
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
    parser.add_argument(
        "--top-k",
        type=int,
        default=1,
        help="top k sampling",
    )

    parser.add_argument(
        "--top-p",
        type=float,
        default=1.0,
        help="top p sampling",
    )

    parser.add_argument(
        "--temperature",
        type=float,
        default=1.0,
        help="sampling temperature",
    )
234
235
236
237
238
    parser.add_argument(
        "--enable-paged-attn",
        action="store_true",
        help="use paged cache",
    )
239
240
241
242
243
244
    parser.add_argument(
        "--paged_kv_block_size",
        type=int,
        default=256,
        help="num tokens each kv block can hold",
    )
245
246
247
248
249
    parser.add_argument(
        "--enable-graph",
        action="store_true",
        help="enable graph compiling",
    )
250
251
252
    parser.add_argument(
        "--warmup",
        action="store_true",
wooway777's avatar
wooway777 committed
253
        help="Perform a warmup run before benchmarking/inference.",
254
    )
255
256
257
    parser.add_argument(
        "--attn",
        type=str,
258
        default="default",
259
260
261
        choices=["default", "flash-attn"],
        help="attention backend to use: 'default' or 'flash-attn'",
    )
pengcheng888's avatar
pengcheng888 committed
262
263
264
    return parser.parse_args()


wooway777's avatar
wooway777 committed
265
266
with open("examples/bench_prompt.md", "r") as f:
    prompt = f.read()
pengcheng888's avatar
pengcheng888 committed
267
268
269
270
271
272
273
274


def repeat_prompt(input_ids: list[int], target_length: int):
    num = len(input_ids)
    repeat_times = (target_length + num - 1) // num
    return (input_ids * repeat_times)[:target_length]


pengcheng888's avatar
pengcheng888 committed
275
276
277
278
class TestModel:
    model: infinicore.nn.Module
    tokenizer: AutoTokenizer
    input_ids_list: list[int]
pengcheng888's avatar
pengcheng888 committed
279

pengcheng888's avatar
pengcheng888 committed
280
281
    def __init__(
        self,
pengcheng888's avatar
pengcheng888 committed
282
        model_path,
pengcheng888's avatar
pengcheng888 committed
283
284
        infini_device=infinicore.device("cpu", 0),
        tp=1,
285
        skip_load=False,
286
287
        cache_config=None,
        enable_graph=False,
288
        attn_backend="default",
pengcheng888's avatar
pengcheng888 committed
289
290
291
292
293
    ) -> None:
        model_path = os.path.expanduser(model_path)
        # ---------------------------------------------------------------------------- #
        #                        创建模型,
        # ---------------------------------------------------------------------------- #
294
        model = InferEngine(
pengcheng888's avatar
pengcheng888 committed
295
296
297
            model_path,
            device=infini_device,
            distributed_config=DistConfig(tp),
298
299
            cache_config=cache_config,
            enable_graph_compiling=enable_graph,
300
            attention_backend=attn_backend,
pengcheng888's avatar
pengcheng888 committed
301
        )
pengcheng888's avatar
pengcheng888 committed
302

pengcheng888's avatar
pengcheng888 committed
303
304
305
        # ---------------------------------------------------------------------------- #
        #                        加载权重
        # ---------------------------------------------------------------------------- #
306
307
        if not skip_load:
            load_model_state_dict_by_file(model, model_path, dtype=model.config.dtype)
pengcheng888's avatar
pengcheng888 committed
308
309
310
311
312

        # ---------------------------------------------------------------------------- #
        #                        创建 tokenizer
        # ---------------------------------------------------------------------------- #
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
wooway777's avatar
wooway777 committed
313

314
315
316
317
318
        if tokenizer.pad_token is None:
            if tokenizer.eos_token is not None:
                tokenizer.pad_token = tokenizer.eos_token
                tokenizer.pad_token_id = tokenizer.eos_token_id
            else:
wooway777's avatar
wooway777 committed
319
                tokenizer.add_special_tokens({"pad_token": "[PAD]"})
pengcheng888's avatar
pengcheng888 committed
320
321
322
323
324
325
326
327
328
329
330
331
332

        # ---------------------------------------------------------------------------- #
        #                        token编码
        # ---------------------------------------------------------------------------- #
        input_content = [
            tokenizer.apply_chat_template(
                conversation=[{"role": "user", "content": prompt}],
                add_generation_prompt=True,
                tokenize=False,
            )
        ]

        # print(input_content, end="", flush=True)
333
334
335
336
337
        # Support Transformers >= 5.0 for batch_encode_plus deprecation
        encoding = tokenizer(
            input_content,
            padding=True,
            truncation=True,
wooway777's avatar
wooway777 committed
338
339
            max_length=8192,
        )
340
341

        input_ids_list = encoding["input_ids"]
pengcheng888's avatar
pengcheng888 committed
342
343
344
345
346
347
348
349
350
351

        self.model = model
        self.tokenizer = tokenizer
        self.input_ids_list = input_ids_list

    def run(
        self,
        batch_size: int,
        input_len: int,
        output_len: int,
352
353
354
        top_k=1,
        top_p=1.0,
        temperature=1.0,
pengcheng888's avatar
pengcheng888 committed
355
356
357
358
359
360
361
362
363
364
365
    ):
        input_ids = repeat_prompt(self.input_ids_list[0], target_length=input_len)
        input_ids_list = [input_ids] * batch_size

        # ---------------------------------------------------------------------------- #
        #                        自回归生成
        # ---------------------------------------------------------------------------- #
        input_ids_infini = infinicore.from_list(input_ids_list)

        t1 = time.time()
        print("=================== start generate ====================")
366
        output_ids = self.model.generate(
pengcheng888's avatar
pengcheng888 committed
367
            input_ids_infini,
368
369
370
371
372
373
            GenerationConfig(
                max_new_tokens=output_len,
                eos_token_id=[],
                top_k=top_k,
                top_p=top_p,
                temperature=temperature,
wooway777's avatar
wooway777 committed
374
                stop_on_eos=False,
375
            ),
PanZezhong's avatar
PanZezhong committed
376
            _measure_and_log_time=True,
pengcheng888's avatar
pengcheng888 committed
377
        )
pengcheng888's avatar
pengcheng888 committed
378
        t2 = time.time()
pengcheng888's avatar
pengcheng888 committed
379

380
381
382
383
384
        numpy_output_ids = np.array(
            [output_id.to_numpy()[0] for output_id in output_ids]
        )
        print(self.tokenizer.decode(numpy_output_ids, skip_special_tokens=True))

pengcheng888's avatar
pengcheng888 committed
385
386
387
        print(
            f"total_time: {round((t2 - t1) * 1000, 2)} ms",
        )
pengcheng888's avatar
pengcheng888 committed
388
389
390
391
392
393
394
395
396
397
398
399


if __name__ == "__main__":
    args = get_args()
    print(args)

    # Parse command line arguments
    device_str = "cpu"
    if args.cpu:
        device_str = "cpu"
    elif args.nvidia:
        device_str = "cuda"
400
401
    elif args.qy:
        device_str = "cuda"
402
403
404
405
406
407
    elif args.metax:
        device_str = "cuda"
    elif args.moore:
        device_str = "musa"
    elif args.iluvatar:
        device_str = "cuda"
408
409
    elif args.cambricon:
        device_str = "mlu"
wooway777's avatar
wooway777 committed
410
411
    elif args.ali:
        device_str = "cuda"
412
413
    elif args.hygon:
        device_str = "cuda"
pengcheng888's avatar
pengcheng888 committed
414
415
    else:
        print(
pengcheng888's avatar
pengcheng888 committed
416
            "python examples/bench.py --nvidia --model=~/TinyLlama-1.1B-Chat-v1.0/ --batch-size=2 --tp=1 --input-len=50 --output-len=50"
pengcheng888's avatar
pengcheng888 committed
417
418
        )
        sys.exit(1)
419
    _PAGED_KV_BLOCK_SIZE = args.paged_kv_block_size
pengcheng888's avatar
pengcheng888 committed
420
421
422
    # -------------------------------------------------------- #
    #             解析参数
    # -------------------------------------------------------- #
pengcheng888's avatar
pengcheng888 committed
423
424
425
426
    model_path = args.model

    infini_device = infinicore.device(device_str, 0)

pengcheng888's avatar
pengcheng888 committed
427
428
    tp = args.tensor_parallel_size

429
430
    skip_load = args.skip_load

pengcheng888's avatar
pengcheng888 committed
431
432
433
    batch_size = args.batch_size
    input_len = args.input_len
    output_len = args.output_len
434
435
    enable_paged_attn = args.enable_paged_attn
    enable_graph = args.enable_graph
pengcheng888's avatar
pengcheng888 committed
436
437
438
439
440
441
442
443
444
445
446
447
448
449

    if isinstance(batch_size, int):
        batch_size = [batch_size]

    if isinstance(input_len, int):
        input_len = [input_len]

    if isinstance(output_len, int):
        output_len = [output_len]

    cases_dict = get_test_cases(model_path, batch_size, input_len, output_len)
    # -------------------------------------------------------- #
    #             测试
    # -------------------------------------------------------- #
450
    if enable_paged_attn:
451
        paged_kv_block_size = _PAGED_KV_BLOCK_SIZE
452
453
        max_num_blocks = max(
            [
454
455
456
457
458
                (
                    (c_["input_len"] + c_["output_len"] + (paged_kv_block_size - 1))
                    // paged_kv_block_size
                )
                * c_["batch_size"]
459
460
461
462
463
464
                for _, c_ in cases_dict.items()
            ]
        )
        cache_config = PagedKVCacheConfig(max_num_blocks, paged_kv_block_size)
    else:
        cache_config = None
pengcheng888's avatar
pengcheng888 committed
465
466

    test = TestModel(
pengcheng888's avatar
pengcheng888 committed
467
        model_path,
pengcheng888's avatar
pengcheng888 committed
468
        infini_device=infini_device,
pengcheng888's avatar
pengcheng888 committed
469
        tp=tp,
470
        skip_load=skip_load,
471
472
        cache_config=cache_config,
        enable_graph=enable_graph,
473
        attn_backend=args.attn,
pengcheng888's avatar
pengcheng888 committed
474
    )
pengcheng888's avatar
pengcheng888 committed
475

476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
    # ---------------------------------------------------------------------------- #
    #                                Warmup
    # ---------------------------------------------------------------------------- #
    if args.warmup:
        warmup_steps = 1

        # warmup cache capacity
        warmup_cache_len = 128
        warmup_batch = len(test.input_ids_list)

        test.model.reset_cache(
            StaticKVCacheConfig(
                max_batch_size=warmup_batch,
                max_cache_len=warmup_cache_len,
            )
        )

wooway777's avatar
wooway777 committed
493
        avg_prompt_len = min(64, max(len(ids) for ids in test.input_ids_list))
494
495
496
497
498
499
500
501
502
503
504
505
506
507

        warmup_ids = [
            ids[:avg_prompt_len] if len(ids) >= avg_prompt_len else ids
            for ids in test.input_ids_list
        ]

        input_ids_infini = infinicore.from_list(warmup_ids)

        print("=================== warmup start ===================")

        for _ in range(warmup_steps):
            _ = test.model.generate(
                input_ids_infini,
                GenerationConfig(
wooway777's avatar
wooway777 committed
508
                    max_new_tokens=5,  # decode kernel warmup
509
510
511
                    temperature=args.temperature,
                    top_k=args.top_k,
                    top_p=args.top_p,
wooway777's avatar
wooway777 committed
512
                    stop_on_eos=False,
513
514
515
516
517
518
519
520
521
522
523
524
525
526
                ),
                _measure_and_log_time=False,
            )

        print("=================== warmup done ====================")

        # reset cache back to benchmark config
        if cache_config is not None:
            test.model.reset_cache(cache_config)

    # ---------------------------------------------------------------------------- #
    #                                Warmup done
    # ---------------------------------------------------------------------------- #

pengcheng888's avatar
pengcheng888 committed
527
528
529
530
531
532
533
    for idx, case in tqdm(cases_dict.items(), desc="Processing cases"):
        tqdm.write(f"\033[92mProcessing : {case}\033[0m")

        batch_size = case["batch_size"]
        input_len = case["input_len"]
        output_len = case["output_len"]

534
535
536
537
538
539
540
        if not enable_paged_attn:
            # reset cache if static kvcache is used
            initial_capacity = input_len + output_len
            test.model.reset_cache(
                StaticKVCacheConfig(
                    max_batch_size=batch_size, max_cache_len=initial_capacity
                )
PanZezhong's avatar
PanZezhong committed
541
            )
pengcheng888's avatar
pengcheng888 committed
542
543
544
545
546
547

        # run test one case
        test.run(
            batch_size=batch_size,
            input_len=input_len,
            output_len=output_len,
548
549
550
            top_k=args.top_k,
            top_p=args.top_p,
            temperature=args.temperature,
pengcheng888's avatar
pengcheng888 committed
551
        )