SwiftSpec/scripts/prepare_models.py at main · ByteDance-Seed/SwiftSpec · 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
# Copyright 2025 ByteDance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tinychat.utils.exp_configs as exp_configs
import os
import subprocess
import copy
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--tp", type=int, default=8)
parser.add_argument("--target_tp", type=int, default=4)
parser.add_argument("--draft_tp", type=int, default=4)

# model arguemtns could be llama3, llama3.3, qwen2, r1llama3, r1qwen2, deepseek-coder
parser.add_argument("--model", type=str, default="llama3.3")

args = parser.parse_args()

# prepare AWQ
# exp_configs.init()

if args.model == "llama3.3":
    exp_config = exp_configs.llama_3_3_exp_config
elif args.model == "llama3":
    exp_config = exp_configs.llama_exp_config
elif args.model == "qwen":
    exp_config = exp_configs.qwen_exp_config
elif args.model == "r1llama":
    exp_config = exp_configs.r1_llama_exp_config
elif args.model == "r1qwen":
    exp_config = exp_configs.r1_qwen_exp_config
elif args.model == "deepseek":
    exp_config = exp_configs.deepseek_coder_exp_config
else:
    raise ValueError(f"Model {args.model} not supported")

tp = args.tp
target_tp = args.target_tp
draft_tp = args.draft_tp

target_model_config = exp_config.target_model_config
draft_model_config = exp_config.draft_model_config
target_model_path = target_model_config.model_path
draft_model_path = draft_model_config.model_path


target_env = copy.deepcopy(os.environ)
target_env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(target_tp))
draft_env = copy.deepcopy(os.environ)
draft_env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(draft_tp))


# # convert draft model to AWQ
# convert_model_command = ["python", "-m", "awq.entry", "--model_path", draft_model_path, "--w_bit", "4", "--q_group_size", "128", "--run_awq", "--dump_awq", f"awq_cache/{draft_model_config.model_id}-awq.pt", "--dtype", "bfloat16"]
# print(f"Launching draft command on GPU {','.join(str(i) for i in range(tp))}: {' '.join(convert_model_command)}")
# subprocess.run(convert_model_command, env=draft_env)

# # convert target model to AWQ
# convert_model_command = ["python", "-m", "awq.entry", "--model_path", target_model_path, "--w_bit", "4", "--q_group_size", "128", "--run_awq", "--dump_awq", f"awq_cache/{target_model_config.model_id}-awq.pt", "--dtype", "bfloat16"]
# print(f"Launching target command on GPU {','.join(str(i) for i in range(tp))}: {' '.join(convert_model_command)}")
# subprocess.run(convert_model_command, env=target_env)


# # convert draft model to TP checkpoint
convert_command = [
    "torchrun",
    f"--nproc-per-node={draft_tp}",
    "../tinychat/convert_model.py",
    "--model_type",
    "llama",
    "--model_path",
    draft_model_path,
    "--q_group_size",
    "128",
    "--load_awq",
    f"awq_cache/{draft_model_config.model_id}-awq.pt",
    "--store-ckpt-file-name",
    f"{draft_model_config.ckpt_prefix}/{draft_model_config.model_id}",
    "--load-int8-weight",
    f"None",
    "--precision",
    "W4A16",
    "--dtype",
    "bf16",
]
print(
    f"Launching convert command on GPU {','.join(str(i) for i in range(draft_tp))}: {' '.join(convert_command)}"
)
subprocess.run(convert_command, env=draft_env)

# convert target model to TP checkpoint
# convert_command = ["torchrun", f"--nproc-per-node={target_tp}", "../tinychat/convert_model.py", "--model_type", "llama", "--model_path", target_model_path, "--q_group_size", "128", "--load_awq", f"awq_cache/{target_model_config.model_id}-awq.pt", "--store-ckpt-file-name", f"{target_model_config.ckpt_prefix}/{target_model_config.model_id}", "--load-int8-weight", f"None", "--precision", "W4A16", "--dtype", "bf16"]
# print(f"Launching convert command on GPU {','.join(str(i) for i in range(target_tp))}: {' '.join(convert_command)}")
# subprocess.run(convert_command, env=target_env)