jiuge.py 7.5 KB
Newer Older
1
2
3
4
5
import infinicore
from transformers import AutoTokenizer
from tokenizers import decoders as _dec
from infinilm.modeling_utils import load_model_state_dict_by_file
from infinilm.distributed import DistConfig
6
from infinilm.infer_engine import GenerationConfig, InferEngine
7
8
9
10
import argparse
import sys
import time
import os
11
import numpy as np
12
from infinilm.cache import StaticKVCacheConfig, PagedKVCacheConfig
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

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


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",
    )
    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",
    )
45
46
47
48
49
    parser.add_argument(
        "--cambricon",
        action="store_true",
        help="Run cambricon test",
    )
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    parser.add_argument(
        "--model_path",
        type=str,
        required=True,
        help="model_path",
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=100,
        help="max_new_tokens",
    )
    parser.add_argument(
        "--backend",
        type=str,
Your Name's avatar
Your Name committed
65
        default="cpp",
66
67
68
        help="python or cpp model",
    )
    parser.add_argument(
pengcheng888's avatar
pengcheng888 committed
69
        "--batch-size",
70
71
72
73
74
75
76
77
78
79
80
81
82
        type=int,
        default=1,
        help="number of prompts in a batch",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default="How are you",
        help="input prompt",
    )
    parser.add_argument(
        "--tp",
        type=int,
Your Name's avatar
Your Name committed
83
        default=1,
84
85
        help="total rank for tensor parallel",
    )
86
87
88
89
90
91
92
93
94
95
96
97
    parser.add_argument(
        "--enable-paged-attn",
        action="store_true",
        help="use paged cache",
    )

    parser.add_argument(
        "--max-kvcache-size",
        type=int,
        default=8 * 1024 * 1024 * 1024,
        help="max size (in bytes) allocated to paged kv cache",
    )
98
99
100
101
102
103
104
105
106

    return parser.parse_args()


def test(
    prompts: str | list[str],
    model_path,
    max_new_tokens=100,
    infini_device=infinicore.device("cpu", 0),
Your Name's avatar
Your Name committed
107
    tp=1,
108
    enable_paged_attn=False,
109
110
111
112
113
):
    model_path = os.path.expanduser(model_path)
    # ---------------------------------------------------------------------------- #
    #                        创建模型,
    # ---------------------------------------------------------------------------- #
114
    model = InferEngine(
115
116
        model_path,
        device=infini_device,
Your Name's avatar
Your Name committed
117
        distributed_config=DistConfig(tp),
118
119
120
121
122
    )

    # ---------------------------------------------------------------------------- #
    #                        加载权重
    # ---------------------------------------------------------------------------- #
123
    load_model_state_dict_by_file(model, model_path, dtype=model.config.dtype)
124
125
126
127
128

    # ---------------------------------------------------------------------------- #
    #                        创建 tokenizer
    # ---------------------------------------------------------------------------- #
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
Your Name's avatar
Your Name committed
129

130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
    if "llama" == model.config.model_type:
        backend = getattr(tokenizer, "backend_tokenizer", None)
        target = getattr(backend, "_tokenizer", backend)
        norm = getattr(target, "normalizer", None)
        dec = getattr(target, "decoder", None)
        sn = repr(norm)[:800] if norm is not None else ""
        sd = repr(dec)[:800] if dec is not None else ""
        has_prepend = "Prepend" in sn
        has_strip = "Strip" in sd
        if has_prepend and has_strip:
            target.decoder = _dec.Sequence(
                [
                    _dec.Replace("▁", " "),
                    _dec.ByteFallback(),
                    _dec.Fuse(),
                ]
            )

    # ---------------------------------------------------------------------------- #
    #                        token编码
    # ---------------------------------------------------------------------------- #
    # prompt = "山东最高的山是?"
    if isinstance(prompts, str):
        prompts = [prompts]
    input_contents = [
        tokenizer.apply_chat_template(
            conversation=[{"role": "user", "content": prompt}],
            add_generation_prompt=True,
            tokenize=False,
        )
        for prompt in prompts
    ]
PanZezhong's avatar
PanZezhong committed
162

163
164
165
166
    input_ids_list = tokenizer.batch_encode_plus(input_contents)[
        "input_ids"
    ]  # List: [[1, 1128, 526, 366, 29892]]

167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
    # ---------------------------------------------------------------------------- #
    #                        创建KVCache
    # ---------------------------------------------------------------------------- #
    if enable_paged_attn:
        cache_config = PagedKVCacheConfig(
            max_kv_memory_bytes=args.max_kvcache_size, block_size=16
        )
    else:
        batch_size = 1 if prompts is str else len(prompts)
        initial_capacity = max_new_tokens + len(input_ids_list[0])
        cache_config = StaticKVCacheConfig(
            max_batch_size=batch_size, max_cache_len=initial_capacity
        )

    model.reset_cache(cache_config)
PanZezhong's avatar
PanZezhong committed
182

183
184
185
    # ---------------------------------------------------------------------------- #
    #                        自回归生成
    # ---------------------------------------------------------------------------- #
PanZezhong's avatar
PanZezhong committed
186
    print(input_contents[0], end="", flush=True)
187
188
189
190
    input_ids_infini = infinicore.from_list(input_ids_list)

    t1 = time.time()
    print("=================== start generate ====================")
191
    output_ids = model.generate(
192
        input_ids_infini,
193
194
195
196
        GenerationConfig(
            max_new_tokens=max_new_tokens, temperature=1, top_k=1, top_p=0.8
        ),
        _measure_and_log_time=True,
197
198
199
    )
    t2 = time.time()

200
201
202
    numpy_output_ids = np.array([output_id.to_numpy()[0] for output_id in output_ids])
    print(tokenizer.decode(numpy_output_ids, skip_special_tokens=True))

203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
    print(
        f"total_time: {round((t2 - t1) * 1000, 2)} ms",
    )


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"
    elif args.metax:
        device_str = "cuda"
    elif args.moore:
        device_str = "musa"
    elif args.iluvatar:
        device_str = "cuda"
224
225
    elif args.cambricon:
        device_str = "mlu"
226
227
    else:
        print(
pengcheng888's avatar
pengcheng888 committed
228
229
            "Usage:  python examples/jiuge.py [--cpu | --nvidia | --metax | --moore | --iluvatar] --model_path=<path/to/model_dir>\n"
            "such as, python examples/jiuge.py --nvidia --model_path=~/TinyLlama-1.1B-Chat-v1.0"
230
231
232
233
234
235
236
        )
        sys.exit(1)
    prompts = [args.prompt for _ in range(args.batch_size)]

    model_path = args.model_path
    max_new_tokens = args.max_new_tokens
    backend = args.backend
Your Name's avatar
Your Name committed
237
    tp = args.tp
238
    enable_paged_attn = args.enable_paged_attn
239
240
241
    if backend != "cpp":
        raise ValueError(f"Unsupported backend: {backend}.")

242
243
244
245
246
247
248
    infini_device = infinicore.device(device_str, 0)

    test(
        prompts,
        model_path,
        max_new_tokens,
        infini_device=infini_device,
Your Name's avatar
Your Name committed
249
        tp=tp,
250
        enable_paged_attn=enable_paged_attn,
251
    )