tiny-llm/main.py at main · aeft/tiny-llm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
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
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
129
130
131
132
133
134
135
136
137
from mlx_lm import load
import mlx_lm
import mlx.core as mx
import argparse

import mlx_lm.sample_utils

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="qwen2-7b")
parser.add_argument("--draft-model", type=str, default=None)
parser.add_argument(
    "--prompt",
    type=str,
    default="Give me a short introduction to large language model.",
)
parser.add_argument("--solution", type=str, default="tiny_llm")
parser.add_argument("--loader", type=str, default="week1")
parser.add_argument("--device", type=str, default="gpu")
parser.add_argument("--sampler-temp", type=float, default=0)
parser.add_argument("--sampler-top-p", type=float, default=None)
parser.add_argument("--sampler-top-k", type=int, default=None)
parser.add_argument("--enable-thinking", action="store_true")
parser.add_argument("--enable-flash-attn", action="store_true")

parser.add_argument(
    "--quantized-mode",
    type=str,
    default="dequantize",
    choices=["custom", "dequantize", "mlx"],
    help="Quantized linear implementation: custom (C++), dequantize (expand), mlx (built-in)"
)

args = parser.parse_args()

use_mlx = False
if args.solution == "tiny_llm":
    print("Using your tiny_llm solution")
    from tiny_llm import (
        models,
        simple_generate,
        simple_generate_with_kv_cache,
        speculative_generate,
        sampler,
    )

    from tiny_llm.quantize import set_quantized_linear_mode
    set_quantized_linear_mode(args.quantized_mode)

elif args.solution == "tiny_llm_ref" or args.solution == "ref":
    print("Using tiny_llm_ref solution")
    from tiny_llm_ref import (
        models,
        simple_generate,
        simple_generate_with_kv_cache,
        speculative_generate,
        sampler,
    )

elif args.solution == "mlx":
    use_mlx = True
    from mlx_lm.generate import stream_generate

    print("Using the original mlx model")
else:
    raise ValueError(f"Solution {args.solution} not supported")

args.model = models.shortcut_name_to_full_name(args.model)
mlx_model, tokenizer = load(args.model)

if args.draft_model:
    args.draft_model = models.shortcut_name_to_full_name(args.draft_model)
    draft_mlx_model, draft_tokenizer = load(args.draft_model)
    if args.loader == "week1":
        raise ValueError("Draft model not supported for week1")
else:
    draft_mlx_model = None
    draft_tokenizer = None

with mx.stream(mx.gpu if args.device == "gpu" else mx.cpu):
    if use_mlx:
        tiny_llm_model = mlx_model
    else:
        if args.loader == "week1":
            print(f"Using week1 loader for {args.model}")
            tiny_llm_model = models.dispatch_model(args.model, mlx_model, week=1)
        elif args.loader == "week2":
            print(
                f"Using week2 loader with flash_attn={args.enable_flash_attn} thinking={args.enable_thinking} for {args.model}"
            )
            tiny_llm_model = models.dispatch_model(
                args.model, mlx_model, week=2, enable_flash_attn=args.enable_flash_attn
            )
            if draft_mlx_model is not None:
                print(f"Using draft model {args.draft_model}")
                draft_tiny_llm_model = models.dispatch_model(
                    args.draft_model,
                    draft_mlx_model,
                    week=2,
                    enable_flash_attn=args.enable_flash_attn,
                )
            else:
                draft_tiny_llm_model = None
        else:
            raise ValueError(f"Loader {args.loader} not supported")
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": args.prompt},
    ]
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=args.enable_thinking,
    )
    if not use_mlx:
        sampler = sampler.make_sampler(
            args.sampler_temp, top_p=args.sampler_top_p, top_k=args.sampler_top_k
        )
        if args.loader == "week1":
            print(simple_generate(tiny_llm_model, tokenizer, prompt, sampler=sampler))
        elif args.loader == "week2":
            if draft_tiny_llm_model is not None:
                speculative_generate(
                    draft_tiny_llm_model,
                    tiny_llm_model,
                    draft_tokenizer,
                    tokenizer,
                    prompt,
                )
            else:
                print(simple_generate_with_kv_cache(tiny_llm_model, tokenizer, prompt))
    else:
        sampler = mlx_lm.sample_utils.make_sampler(
            args.sampler_temp, top_p=args.sampler_top_p, top_k=args.sampler_top_k
        )
        for resp in stream_generate(tiny_llm_model, tokenizer, prompt, sampler=sampler):
            print(resp.text, end="", flush=True)