-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathlm_eval_awq_mixtral.py
149 lines (116 loc) · 6.12 KB
/
lm_eval_awq_mixtral.py
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
138
139
140
141
142
143
144
145
146
147
148
import argparse
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import json
import numpy as np
from lm_eval import evaluator
from lm_eval.models.huggingface import HFLM
from lm_eval.tasks import initialize_tasks
from transformers import AutoModelForCausalLM, AutoTokenizer
from optimum.gptq import GPTQQuantizer, load_quantized_model, GPTQQuantizer_deepseek
import torch
import random
from argparse import ArgumentParser
from transformers import AutoTokenizer, TextGenerationPipeline
import logging
from datasets import load_dataset
from awq import AutoAWQForCausalLM
DEEPSEEK_MODEL_COMPONENTS = [
'model.embed_tokens', 'model.layers.0', 'model.layers.1', 'model.layers.2', 'model.layers.3', 'model.layers.4', 'model.layers.5', 'model.layers.6', 'model.layers.7', 'model.layers.8',
'model.layers.9', 'model.layers.10', 'model.layers.11', 'model.layers.12', 'model.layers.13', 'model.layers.14', 'model.layers.15', 'model.layers.16', 'model.layers.17', 'model.layers.18',
'model.layers.19', 'model.layers.20', 'model.layers.21', 'model.layers.22', 'model.layers.23', 'model.layers.24', 'model.layers.25', 'model.layers.26', 'model.layers.27', 'model.norm', 'lm_head'
]
LM_EVAL_TASK_KWARGS_DICT = {
"winogrande": {"task": "winogrande", "num_fewshot": 0, "batch_size": 128, "metric": "acc"},
"copa": {"task": "copa", "num_fewshot": 0, "batch_size": 128, "metric": "acc"},
"openbookqa": {"task": "openbookqa", "num_fewshot": 0, "batch_size": 128, "metric": "acc_norm"},
"hellaswag": {"task": "hellaswag", "num_fewshot": 0, "batch_size": 128, "metric": "acc_norm"},
"lambada_openai": {"task": "lambada_openai", "num_fewshot": 0, "batch_size": 128, "metric": "acc"},
"rte": {"task": "rte", "num_fewshot": 0, "batch_size": 128, "metric": "acc"},
"piqa": {"task": "piqa", "num_fewshot": 0, "batch_size": 128, "metric": "acc"},
}
# "mmlu": {"task": "mmlu", "num_fewshot": 5, "batch_size": 8, "metric": "acc"},
def create_device_map(components):
num_gpus = torch.cuda.device_count()
device_map = {}
if num_gpus == 0:
print("No GPU found. Please check your system.")
return device_map
part_size = len(components) // num_gpus
for i, component in enumerate(components):
device_id = i // part_size
device_id = min(device_id, num_gpus - 1)
device_map[component] = device_id
print(f"Device map: {device_map}")
return device_map
if __name__ == "__main__":
print("Start evaluating model...")
print(f"Evaluating task: {LM_EVAL_TASK_KWARGS_DICT}")
parser = argparse.ArgumentParser(description="Calculate Perplexity for a model.")
parser.add_argument("--model_path", type=str, default='')
parser.add_argument("--bits", type=str, default=None)
parser.add_argument("--is_quantized", action="store_true", help="Is the model GPTQ quantized?", default=True)
args = parser.parse_args()
if args.is_quantized:
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# args.model_path = 'llama-2-7b-awq-w_bit.4-group_size.128'
# args.model_path = 'quantized_deepseek-moe-16b-base-awq-w_bit4-group_size64'
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
if not tokenizer.pad_token_id:
tokenizer.pad_token_id = tokenizer.eos_token_id
if 'deepseek' in args.model_path:
device_map = create_device_map(DEEPSEEK_MODEL_COMPONENTS)
else:
device_map = 'auto'
model = AutoAWQForCausalLM.from_quantized(args.model_path, quant_file='', fuse_layers=False, device_map=device_map)
# Prepare input text
input_text = "The purpose of life is"
input_ids = tokenizer.encode(input_text, return_tensors='pt')
output = model.generate(input_ids)
else:
if 'deepseek' in args.model_path:
device_map = create_device_map(DEEPSEEK_MODEL_COMPONENTS)
else:
device_map = 'auto'
model_name = "deepseek-ai/deepseek-moe-16b-base"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map=device_map)
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
for name, param in model.model.named_parameters():
print(name, param.dtype)
total_params = sum(p.numel() for p in model.model.parameters() if p.requires_grad)
print("Total number of parameters:", total_params)
total_memory = sum(p.element_size() * p.numel() for p in model.model.parameters() if p.requires_grad)
print("Total memory used by model parameters (in bytes):", total_memory)
all_metrics = {}
save_file_path = os.path.join("awq_eval_result", f"eval_result_{args.model_path}.log")
os.makedirs("awq_eval_result", exist_ok=True)
print(LM_EVAL_TASK_KWARGS_DICT)
for task_kwargs in LM_EVAL_TASK_KWARGS_DICT.values():
print(f"Evaluating task: {task_kwargs['task']}")
task_name = task_kwargs["task"]
initialize_tasks(verbosity="ERROR")
lm = HFLM(
pretrained=model,
tokenizer=tokenizer,
batch_size=task_kwargs["batch_size"],
)
results = evaluator.simple_evaluate(
model=lm,
tasks=task_name,
num_fewshot=task_kwargs["num_fewshot"],
batch_size=task_kwargs["batch_size"],
log_samples=False,
)
metric = task_kwargs["metric"]
for key, value in results["results"][task_name].items():
if key.startswith(metric + ","):
all_metrics[f"{task_name}_{metric}"] = value
print(">>>>> Results <<<<<")
print(f"Metrics: {all_metrics}")
with open(save_file_path, 'w') as file:
json.dump(all_metrics, file, indent=4)
print(">>>>> Results <<<<<")
print(f"Metrics: {all_metrics}")
# nohup python lm_eval_awq.py --model_path 'quantized_mistral-instruct-v0.2-awq-w_bit.8-group_size.64' --is_quantized > awq_eval_deepseek-moe-16b-base-awq-w_bit4-group_size64.log 2>&1 &