jiuge.py 16.4 KB
Newer Older
PanZezhong's avatar
PanZezhong committed
1
from ctypes import POINTER, c_int, c_uint, c_void_p, byref
Pan Zezhong's avatar
Pan Zezhong committed
2
import os
PanZezhong's avatar
PanZezhong committed
3
4
from pathlib import Path
import safetensors
PanZezhong's avatar
PanZezhong committed
5
import sys
PanZezhong's avatar
PanZezhong committed
6
import time
Pan Zezhong's avatar
Pan Zezhong committed
7
import json
PanZezhong's avatar
PanZezhong committed
8

PanZezhong's avatar
PanZezhong committed
9
10
11
12
13
14
15
from libinfinicore_infer import (
    JiugeMeta,
    JiugeWeights,
    KVCache,
    DataType,
    DeviceType,
    create_jiuge_model,
PanZezhong's avatar
PanZezhong committed
16
    destroy_jiuge_model,
PanZezhong's avatar
PanZezhong committed
17
18
19
20
21
22
23
    create_kv_cache,
    drop_kv_cache,
    infer_batch,
)
import torch
import transformers

Pan Zezhong's avatar
Pan Zezhong committed
24
torch.set_default_device("cpu")
PanZezhong's avatar
PanZezhong committed
25
26
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

class LlamaWeightsNaming:
    def input_embd(self):
        return "model.embed_tokens.weight"

    def output_norm(self):
        return "model.norm.weight"

    def output_embd(self):
        return "lm_head.weight"

    def attn_norm(self, i):
        return f"model.layers.{i}.input_layernorm.weight"

    def attn_q(self, i):
        return f"model.layers.{i}.self_attn.q_proj.weight"

    def attn_k(self, i):
        return f"model.layers.{i}.self_attn.k_proj.weight"

    def attn_v(self, i):
        return f"model.layers.{i}.self_attn.v_proj.weight"

    def attn_o(self, i):
        return f"model.layers.{i}.self_attn.o_proj.weight"

    def attn_q_b(self, i):
        return f"model.layers.{i}.self_attn.q_proj.bias"

    def attn_k_b(self, i):
        return f"model.layers.{i}.self_attn.k_proj.bias"

    def attn_v_b(self, i):
        return f"model.layers.{i}.self_attn.v_proj.bias"

PanZezhong's avatar
PanZezhong committed
60
61
62
    def ffn_norm(self, i):
        return f"model.layers.{i}.post_attention_layernorm.weight"

PanZezhong's avatar
PanZezhong committed
63
64
65
66
67
68
69
70
71
    def gate(self, i):
        return f"model.layers.{i}.mlp.gate_proj.weight"

    def up(self, i):
        return f"model.layers.{i}.mlp.up_proj.weight"

    def down(self, i):
        return f"model.layers.{i}.mlp.down_proj.weight"

PanZezhong's avatar
PanZezhong committed
72
73
74
75
76
77
    def match(state_dict):
        return (
            "model.norm.weight" in state_dict
            and "model.layers.0.self_attn.q_proj.weight" in state_dict
        )

PanZezhong's avatar
PanZezhong committed
78
79

class JiugeMetaFromLlama(JiugeMeta):
Pan Zezhong's avatar
Pan Zezhong committed
80
81
    def __init__(self, config, dtype=torch.float16):
        if dtype == torch.float16:
PanZezhong's avatar
PanZezhong committed
82
83
84
85
86
            dt_ = DataType.INFINI_DTYPE_F16
        elif dtype == torch.float32:
            dt_ = DataType.INFINI_DTYPE_F32
        else:
            dt_ = DataType.INFINI_DTYPE_F16
PanZezhong's avatar
PanZezhong committed
87
        super().__init__(
PanZezhong's avatar
PanZezhong committed
88
            dt_logits=dt_,
Pan Zezhong's avatar
Pan Zezhong committed
89
90
91
            nlayer=config["num_hidden_layers"],
            d=config["hidden_size"],
            nh=config["num_attention_heads"],
PanZezhong's avatar
PanZezhong committed
92
            nkvh=(
Pan Zezhong's avatar
Pan Zezhong committed
93
94
95
                config["num_key_value_heads"]
                if "num_key_value_heads" in config
                else config["num_attention_heads"]
PanZezhong's avatar
PanZezhong committed
96
            ),
Pan Zezhong's avatar
Pan Zezhong committed
97
98
99
100
101
102
            dh=config["hidden_size"] // config["num_attention_heads"],
            di=config["intermediate_size"],
            dctx=config["max_position_embeddings"],
            dvoc=config["vocab_size"],
            epsilon=config["rms_norm_eps"],
            theta=(config["rope_theta"] if "rope_theta" in config else 100000.0),
PanZezhong's avatar
PanZezhong committed
103
104
            end_token=2,
        )
PanZezhong's avatar
PanZezhong committed
105
        self.torch_dtype_logits = dtype
PanZezhong's avatar
PanZezhong committed
106
107
108


class JiugeWeightsImpl(JiugeWeights):
Pan Zezhong's avatar
Pan Zezhong committed
109
110
111
112
113
114
115
116
117
    def __init__(
        self,
        meta,
        naming,
        state_dict,
        torch_dt_mat=torch.float16,
        torch_dt_norm=torch.float32,
        ndev=1,
    ):
PanZezhong's avatar
PanZezhong committed
118
119
120
121
122
123
124
125
126
127
        nlayer = meta.nlayer
        nh = meta.nh
        nkvh = meta.nkvh
        dh = meta.dh
        d = meta.d
        di = meta.di
        assert nh % nkvh == 0
        assert nh % ndev == 0
        assert nkvh % ndev == 0
        assert di % ndev == 0
PanZezhong's avatar
PanZezhong committed
128
129
130
131
132
133
134
135
136
137
138
139
140
        torch_dt_logits = meta.torch_dtype_logits
        if torch_dt_mat == torch.float16:
            self.dt_mat = DataType.INFINI_DTYPE_F16
        elif torch_dt_mat == torch.float32:
            self.dt_mat = DataType.INFINI_DTYPE_F32
        else:
            raise ValueError("Unsupported proj weight data type")
        if torch_dt_norm == torch.float16:
            self.dt_norm = DataType.INFINI_DTYPE_F16
        elif torch_dt_norm == torch.float32:
            self.dt_norm = DataType.INFINI_DTYPE_F32
        else:
            raise ValueError("Unsupported norm weight data type")
PanZezhong's avatar
PanZezhong committed
141

Pan Zezhong's avatar
Pan Zezhong committed
142
143
144
145
146
147
148
149
150
151
152
        input_embd_naming = (
            naming.input_embd()
            if naming.input_embd() in state_dict
            else naming.output_embd()
        )
        output_embd_naming = (
            naming.output_embd()
            if naming.output_embd() in state_dict
            else naming.input_embd()
        )

PanZezhong's avatar
PanZezhong committed
153
        self.nlayer = nlayer
Pan Zezhong's avatar
Pan Zezhong committed
154
        self.input_embd_tensor = state_dict[input_embd_naming].to(torch_dt_logits)
PanZezhong's avatar
PanZezhong committed
155
        self.input_embd = self.input_embd_tensor.data_ptr()
PanZezhong's avatar
PanZezhong committed
156
        self.output_norm_tensor = state_dict[naming.output_norm()].to(torch_dt_norm)
PanZezhong's avatar
PanZezhong committed
157
        self.output_norm = self.output_norm_tensor.data_ptr()
Pan Zezhong's avatar
Pan Zezhong committed
158
        self.output_embd_tensor = state_dict[output_embd_naming].to(torch_dt_mat)
PanZezhong's avatar
PanZezhong committed
159
160
161
        self.output_embd = self.output_embd_tensor.data_ptr()

        self.attn_norm_tensors = [
PanZezhong's avatar
PanZezhong committed
162
            state_dict[naming.attn_norm(i)].to(torch_dt_norm) for i in range(nlayer)
PanZezhong's avatar
PanZezhong committed
163
164
165
166
167
        ]
        self.attn_norm_ptrs = [
            self.attn_norm_tensors[i].data_ptr() for i in range(nlayer)
        ]
        self.attn_norm = (c_void_p * nlayer)(*self.attn_norm_ptrs)
PanZezhong's avatar
PanZezhong committed
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189

        def qkv_slices(_i):
            _Q = (
                state_dict[naming.attn_q(_i)]
                .reshape([nh, 2, dh // 2, d])
                .transpose(1, 2)
            )
            _K = (
                state_dict[naming.attn_k(_i)]
                .reshape([nkvh, 2, dh // 2, d])
                .transpose(1, 2)
            )
            _V = state_dict[naming.attn_v(_i)].reshape([nkvh, dh // 2, 2, d])
            _result = []
            _nh = nh // ndev
            _nkvh = nkvh // ndev
            for _idev in range(ndev):
                _result.append(_Q[_idev * _nh : (_idev + 1) * _nh, :, :, :])
                _result.append(_K[_idev * _nkvh : (_idev + 1) * _nkvh, :, :, :])
                _result.append(_V[_idev * _nkvh : (_idev + 1) * _nkvh, :, :])
            return _result

Pan Zezhong's avatar
Pan Zezhong committed
190
191
192
        self.qkv_tensor = [
            torch.concat(qkv_slices(i)).to(torch_dt_mat) for i in range(nlayer)
        ]
PanZezhong's avatar
PanZezhong committed
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        self.qkv_tensor_ptrs = [self.qkv_tensor[i].data_ptr() for i in range(nlayer)]
        self.attn_qkv = (c_void_p * nlayer)(*self.qkv_tensor_ptrs)

        def qkv_b_slices(_i):
            _QB = (
                state_dict[naming.attn_q_b(_i)]
                .reshape([nh, 2, dh // 2])
                .transpose(1, 2)
            )
            _KB = (
                state_dict[naming.attn_k_b(_i)]
                .reshape([nkvh, 2, dh // 2])
                .transpose(1, 2)
            )
            _VB = state_dict[naming.attn_v_b(_i)].reshape([nkvh, dh // 2, 2])
            _result = []
            _nh = nh // ndev
            _nkvh = nkvh // ndev
            for _idev in range(ndev):
Pan Zezhong's avatar
Pan Zezhong committed
212
213
214
                _result.append(_QB[_idev * _nh : (_idev + 1) * _nh, :, :].flatten())
                _result.append(_KB[_idev * _nkvh : (_idev + 1) * _nkvh, :, :].flatten())
                _result.append(_VB[_idev * _nkvh : (_idev + 1) * _nkvh, :, :].flatten())
PanZezhong's avatar
PanZezhong committed
215
216
217
            return _result

        if naming.attn_q_b(0) in state_dict:
Pan Zezhong's avatar
Pan Zezhong committed
218
219
220
            self.qkv_b_tensors = [
                torch.concat(qkv_b_slices(i)).to(torch_dt_logits) for i in range(nlayer)
            ]
PanZezhong's avatar
PanZezhong committed
221
222
223
224
225
226
227
            self.qkv_b_tensor_ptrs = [
                self.qkv_b_tensors[i].data_ptr() for i in range(nlayer)
            ]
            self.attn_qkv_b = (c_void_p * nlayer)(*self.qkv_b_tensor_ptrs)
        else:
            self.attn_qkv_b = None

PanZezhong's avatar
PanZezhong committed
228
        self.attn_o_tensor = [
Pan Zezhong's avatar
Pan Zezhong committed
229
230
            state_dict[naming.attn_o(i)]
            .to(torch_dt_mat)
PanZezhong's avatar
PanZezhong committed
231
232
233
234
235
            .reshape([d, ndev, nh // ndev * dh])
            .transpose(0, 1)
            .contiguous()
            for i in range(nlayer)
        ]
PanZezhong's avatar
PanZezhong committed
236
237
238
        self.attn_o_ptrs = [self.attn_o_tensor[i].data_ptr() for i in range(nlayer)]
        self.attn_o = (c_void_p * nlayer)(*self.attn_o_ptrs)

Pan Zezhong's avatar
Pan Zezhong committed
239
240
241
        self.ffn_norm_tensors = [
            state_dict[naming.ffn_norm(i)].to(torch_dt_norm) for i in range(nlayer)
        ]
PanZezhong's avatar
PanZezhong committed
242
243
244
245
        self.ffn_norm_ptrs = [
            self.ffn_norm_tensors[i].data_ptr() for i in range(nlayer)
        ]
        self.ffn_norm = (c_void_p * nlayer)(*self.ffn_norm_ptrs)
PanZezhong's avatar
PanZezhong committed
246
247
248
249
250
251
252
253
254
255
256

        def gate_up_slices(_i):
            _result = []
            _di = di // ndev
            for _idev in range(ndev):
                _start = _idev * _di
                _end = (_idev + 1) * _di
                _result.append(state_dict[naming.gate(_i)][_start:_end, :])
                _result.append(state_dict[naming.up(_i)][_start:_end, :])
            return _result

Pan Zezhong's avatar
Pan Zezhong committed
257
258
259
        self.gate_up_tensors = [
            torch.concat(gate_up_slices(i)).to(torch_dt_mat) for i in range(nlayer)
        ]
PanZezhong's avatar
PanZezhong committed
260
261
        self.gate_up_ptrs = [self.gate_up_tensors[i].data_ptr() for i in range(nlayer)]
        self.ffn_gate_up = (c_void_p * nlayer)(*self.gate_up_ptrs)
PanZezhong's avatar
PanZezhong committed
262
263

        self.ffn_down_tensor = [
Pan Zezhong's avatar
Pan Zezhong committed
264
265
            state_dict[naming.down(i)]
            .to(torch_dt_mat)
PanZezhong's avatar
PanZezhong committed
266
267
268
269
270
            .reshape([d, ndev, di // ndev])
            .transpose(0, 1)
            .contiguous()
            for i in range(nlayer)
        ]
PanZezhong's avatar
PanZezhong committed
271
272
        self.ffn_down_ptrs = [self.ffn_down_tensor[i].data_ptr() for i in range(nlayer)]
        self.ffn_down = (c_void_p * nlayer)(*self.ffn_down_ptrs)
PanZezhong's avatar
PanZezhong committed
273
274
275
276


class JiugeForCauslLM:
    def __init__(self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1):
PanZezhong's avatar
PanZezhong committed
277
        def load_all_safetensors_from_dir(dir_path_: str):
PanZezhong's avatar
PanZezhong committed
278
279
280
281
282
            tensors_ = {}
            dir_path_ = Path(dir_path_)
            for file in sorted(dir_path_.glob("*.safetensors")):
                data_ = safetensors.safe_open(file, "pt")
                for name_ in data_.keys():
PanZezhong's avatar
PanZezhong committed
283
                    tensors_[name_] = data_.get_tensor(name_)
PanZezhong's avatar
PanZezhong committed
284
            return tensors_
PanZezhong's avatar
PanZezhong committed
285
286
287
        
        print("Loading model weights to host...")
        load_start_time = time.time()
PanZezhong's avatar
PanZezhong committed
288

Pan Zezhong's avatar
Pan Zezhong committed
289
290
        with open(os.path.join(model_dir_path, "config.json"), "r") as f:
            config = json.load(f)
PanZezhong's avatar
PanZezhong committed
291
292
293
            self.config = config
        eos_token_id = self.config["eos_token_id"]
        self.eos_token_id = [eos_token_id] if type(eos_token_id) == int else eos_token_id
Pan Zezhong's avatar
Pan Zezhong committed
294
295
296
        if "llama" == config["model_type"]:
            model = transformers.LlamaForCausalLM.from_pretrained(model_dir_path).cpu().half()
            self.meta = JiugeMetaFromLlama(config)
PanZezhong's avatar
PanZezhong committed
297
298
299
300
            self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir_path)
            self.weights = JiugeWeightsImpl(
                self.meta, LlamaWeightsNaming(), model.state_dict(), ndev=ndev
            )
Pan Zezhong's avatar
Pan Zezhong committed
301
        elif "fm9g" == config["model_type"]:
PanZezhong's avatar
PanZezhong committed
302
303
304
305
306
307
            if any(file.suffix == ".safetensors" for file in Path(model_dir_path).iterdir()):
                state_dict = load_all_safetensors_from_dir(model_dir_path)
            else:
                state_dict = torch.load(
                os.path.join(model_dir_path, "pytorch_model.bin"), weights_only=True, map_location="cpu"
            )
PanZezhong's avatar
PanZezhong committed
308
            if LlamaWeightsNaming.match(state_dict):
PanZezhong's avatar
PanZezhong committed
309
                self.meta = JiugeMetaFromLlama(config)
PanZezhong's avatar
PanZezhong committed
310
311
312
313
314
315
                self.weights = JiugeWeightsImpl(
                    self.meta, LlamaWeightsNaming(), state_dict, ndev=ndev
                )
                self.tokenizer = transformers.AutoTokenizer.from_pretrained(
                    model_dir_path, trust_remote_code=True
                )
PanZezhong's avatar
PanZezhong committed
316
317
            else:
                raise ValueError("Unsupported weight naming")
Pan Zezhong's avatar
Pan Zezhong committed
318
319
320
321
322
323
324
325
326
327
328
329
        elif "fm9g7b" == config["model_type"]:
            state_dict = torch.load(
                os.path.join(model_dir_path, "pytorch_model.bin"), weights_only=True, map_location="cpu"
            )
            if LlamaWeightsNaming.match(state_dict):
                self.meta = JiugeMetaFromLlama(config)
                self.weights = JiugeWeightsImpl(
                    self.meta, LlamaWeightsNaming(), state_dict, ndev=ndev
                )
                self.tokenizer = transformers.AutoTokenizer.from_pretrained(
                    model_dir_path, trust_remote_code=True
                )
PanZezhong's avatar
PanZezhong committed
330
331
332
333
334
335
336
337
338
339
340
341
            else:
                raise ValueError("Unsupported weight naming")
        elif "qwen2" == config["model_type"]:
            state_dict = load_all_safetensors_from_dir(model_dir_path)
            if LlamaWeightsNaming.match(state_dict):
                self.meta = JiugeMetaFromLlama(config)
                self.weights = JiugeWeightsImpl(
                    self.meta, LlamaWeightsNaming(), state_dict, ndev=ndev
                )
                self.tokenizer = transformers.AutoTokenizer.from_pretrained(
                    model_dir_path
                )
PanZezhong's avatar
PanZezhong committed
342
343
        else:
            raise ValueError("Unsupported model architecture")
PanZezhong's avatar
PanZezhong committed
344
345
346
347
348
349

        load_end_time = time.time()
        print(f"Time used: {load_end_time - load_start_time:.3f}s")
        
        print(f"Creating model on {ndev} devices...")
        load_start_time = time.time()
PanZezhong's avatar
PanZezhong committed
350
        dev_ids = (c_int * ndev)(*[i for i in range(ndev)])
PanZezhong's avatar
PanZezhong committed
351
352
353
354
355
356
357
        self.model_instance = create_jiuge_model(
            byref(self.meta),
            byref(self.weights),
            device,
            ndev,
            dev_ids,
        )
PanZezhong's avatar
PanZezhong committed
358
359
360
        load_end_time = time.time()
        print(f"Time used: {load_end_time - load_start_time:.3f}s")
        
PanZezhong's avatar
PanZezhong committed
361
362
363
364
365

    def infer(self, input_list, topp=1.0, topk=1, temperature=1.0):
        pass

    def generate(self, input_content, max_steps, topp=1.0, topk=1, temperature=1.0):
PanZezhong's avatar
PanZezhong committed
366
        kv_cache = create_kv_cache(self.model_instance)
Pan Zezhong's avatar
Pan Zezhong committed
367
368
369
370
371
        input_content = self.tokenizer.apply_chat_template(
            conversation=[{"role": "user", "content": input_content}],
            add_generation_prompt=True,
            tokenize=False,
        )
PanZezhong's avatar
PanZezhong committed
372
373
374
375
376
377
378
379
380
381
382
383
        print(input_content, end="", flush=True)
        tokens = self.tokenizer.encode(input_content)
        ntok = len(tokens)
        nreq = 1
        output_content = ""
        tokens = (c_uint * ntok)(*tokens)
        req_lens = (c_uint * nreq)(*[ntok])
        req_pos = (c_uint * nreq)(*[0])
        kv_caches = (POINTER(KVCache) * nreq)(*[kv_cache])
        ans = (c_uint * nreq)()

        steps = 0
Pan Zezhong's avatar
Pan Zezhong committed
384
385
386
387
        total_time = 0

        for step_i in range(max_steps):
            start_time = time.time()
PanZezhong's avatar
PanZezhong committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
            infer_batch(
                self.model_instance,
                tokens,
                ntok,
                req_lens,
                nreq,
                req_pos,
                kv_caches,
                ans,
                temperature,
                topk,
                topp,
            )
            steps += 1
            output_tokens = list(ans)
PanZezhong's avatar
PanZezhong committed
403
            end_time = time.time()
PanZezhong's avatar
PanZezhong committed
404
405
406
407
408
409
410
            output_str = (
                self.tokenizer._tokenizer.id_to_token(output_tokens[0])
                .replace("▁", " ")
                .replace("<0x0A>", "\n")
            )
            output_content += output_str
            print(output_str, end="", flush=True)
PanZezhong's avatar
PanZezhong committed
411
412
            if output_tokens[0] in self.eos_token_id:
                break
PanZezhong's avatar
PanZezhong committed
413
414
415
416
            req_pos[0] = req_pos[0] + ntok
            ntok = 1
            tokens = (c_uint * ntok)(*output_tokens)
            req_lens = (c_uint * nreq)(*[ntok])
PanZezhong's avatar
PanZezhong committed
417
            
Pan Zezhong's avatar
Pan Zezhong committed
418
419
            if step_i > 0:
                total_time += end_time - start_time
PanZezhong's avatar
PanZezhong committed
420
421

        print("\n")
Pan Zezhong's avatar
Pan Zezhong committed
422
        avg_time = total_time * 1000 / (steps - 1)
PanZezhong's avatar
PanZezhong committed
423
424
425
426
        print(f"Time per step: {avg_time:.3f}ms")
        for kv_cache in kv_caches:
            drop_kv_cache(self.model_instance, kv_cache)
        return output_content, avg_time
PanZezhong's avatar
PanZezhong committed
427
428
429
430
    
    def destroy_model_instance(self):
        destroy_jiuge_model(self.model_instance)
        print("Model destroyed")
PanZezhong's avatar
PanZezhong committed
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460


def test():
    if len(sys.argv) < 3:
        print(
            "Usage: python test_llama.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
        )
        sys.exit(1)
    model_path = sys.argv[2]
    device_type = DeviceType.DEVICE_TYPE_CPU
    if sys.argv[1] == "--cpu":
        device_type = DeviceType.DEVICE_TYPE_CPU
    elif sys.argv[1] == "--nvidia":
        device_type = DeviceType.DEVICE_TYPE_NVIDIA
    elif sys.argv[1] == "--cambricon":
        device_type = DeviceType.DEVICE_TYPE_CAMBRICON
    elif sys.argv[1] == "--ascend":
        device_type = DeviceType.DEVICE_TYPE_ASCEND
    elif sys.argv[1] == "--metax":
        device_type = DeviceType.DEVICE_TYPE_METAX
    elif sys.argv[1] == "--moore":
        device_type = DeviceType.DEVICE_TYPE_MOORE
    else:
        print(
            "Usage: python test_llama.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] <path/to/model_dir> [n_device]"
        )
        sys.exit(1)

    ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1
    model = JiugeForCauslLM(model_path, device_type, ndev)
Pan Zezhong's avatar
Pan Zezhong committed
461
    model.generate("山东最高的山是?", 500)
PanZezhong's avatar
PanZezhong committed
462
    model.destroy_model_instance()
PanZezhong's avatar
PanZezhong committed
463
464
465
466


if __name__ == "__main__":
    test()