jiuge.py 7.45 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
    parser.add_argument(
        "--enable-paged-attn",
        action="store_true",
        help="use paged cache",
    )

92
93
94
95
96
97
98
99
    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
100
    tp=1,
101
    enable_paged_attn=False,
102
103
104
):
    model_path = os.path.expanduser(model_path)
    # ---------------------------------------------------------------------------- #
105
    #                        Create Model
106
    # ---------------------------------------------------------------------------- #
107
    model = InferEngine(
108
109
        model_path,
        device=infini_device,
Your Name's avatar
Your Name committed
110
        distributed_config=DistConfig(tp),
111
112
113
    )

    # ---------------------------------------------------------------------------- #
114
    #                        Load Weights
115
    # ---------------------------------------------------------------------------- #
116
    load_model_state_dict_by_file(model, model_path, dtype=model.config.dtype)
117
118

    # ---------------------------------------------------------------------------- #
119
    #                        create tokenizer
120
121
    # ---------------------------------------------------------------------------- #
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
Your Name's avatar
Your Name committed
122

123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
    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(),
                ]
            )

    # ---------------------------------------------------------------------------- #
142
    #                        tokenize
143
144
145
146
147
148
149
150
151
152
153
154
    # ---------------------------------------------------------------------------- #
    # 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
155

156
157
158
159
    input_ids_list = tokenizer.batch_encode_plus(input_contents)[
        "input_ids"
    ]  # List: [[1, 1128, 526, 366, 29892]]

160
    # ---------------------------------------------------------------------------- #
161
    #                       Create KVCache
162
163
    # ---------------------------------------------------------------------------- #
    if enable_paged_attn:
164
165
        batch_size = 1 if prompts is str else len(prompts)
        max_total_tokens = max_new_tokens + len(input_ids_list[0])
166
        cache_config = PagedKVCacheConfig(
167
            num_blocks=(max_total_tokens // 16 + 1) * batch_size, block_size=16
168
169
170
171
172
173
174
175
176
        )
    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
177

178
    # ---------------------------------------------------------------------------- #
179
    #                        Generate
180
    # ---------------------------------------------------------------------------- #
PanZezhong's avatar
PanZezhong committed
181
    print(input_contents[0], end="", flush=True)
182
183
184
185
    input_ids_infini = infinicore.from_list(input_ids_list)

    t1 = time.time()
    print("=================== start generate ====================")
186
    output_ids = model.generate(
187
        input_ids_infini,
188
189
190
191
        GenerationConfig(
            max_new_tokens=max_new_tokens, temperature=1, top_k=1, top_p=0.8
        ),
        _measure_and_log_time=True,
192
193
194
    )
    t2 = time.time()

195
196
197
    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))

198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
    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"
219
220
    elif args.cambricon:
        device_str = "mlu"
221
222
    else:
        print(
pengcheng888's avatar
pengcheng888 committed
223
224
            "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"
225
226
227
228
229
230
231
        )
        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
232
    tp = args.tp
233
    enable_paged_attn = args.enable_paged_attn
234
235
236
    if backend != "cpp":
        raise ValueError(f"Unsupported backend: {backend}.")

237
238
239
240
241
242
243
    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
244
        tp=tp,
245
        enable_paged_attn=enable_paged_attn,
246
    )