|
"""Run a chatbot with FlexGen and OPT models.""" |
|
import argparse |
|
|
|
from transformers import AutoTokenizer |
|
from flexgen.flex_opt import (Policy, OptLM, TorchDevice, TorchDisk, TorchMixedDevice, |
|
CompressionConfig, Env, Task, get_opt_config) |
|
|
|
|
|
def main(args): |
|
|
|
gpu = TorchDevice("cuda:0") |
|
cpu = TorchDevice("cpu") |
|
disk = TorchDisk(args.offload_dir) |
|
env = Env(gpu=gpu, cpu=cpu, disk=disk, mixed=TorchMixedDevice([gpu, cpu, disk])) |
|
|
|
|
|
policy = Policy(1, 1, |
|
args.percent[0], args.percent[1], |
|
args.percent[2], args.percent[3], |
|
args.percent[4], args.percent[5], |
|
overlap=True, sep_layer=True, pin_weight=True, |
|
cpu_cache_compute=False, attn_sparsity=1.0, |
|
compress_weight=args.compress_weight, |
|
comp_weight_config=CompressionConfig( |
|
num_bits=4, group_size=64, |
|
group_dim=0, symmetric=False), |
|
compress_cache=args.compress_cache, |
|
comp_cache_config=CompressionConfig( |
|
num_bits=4, group_size=64, |
|
group_dim=2, symmetric=False)) |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-30b", padding_side="left") |
|
tokenizer.add_bos_token = False |
|
stop = tokenizer("\n").input_ids[0] |
|
|
|
print("Initialize...") |
|
opt_config = get_opt_config(args.model) |
|
model = OptLM(opt_config, env, args.path, policy) |
|
model.init_all_weights() |
|
|
|
context = ( |
|
"A chat between a curious human and a knowledgeable artificial intelligence assistant.\n" |
|
"Human: Hello! What can you do?\n" |
|
"Assistant: As an AI assistant, I can answer questions and chat with you.\n" |
|
"Human: What is the name of the tallest mountain in the world?\n" |
|
"Assistant: Everest.\n" |
|
) |
|
|
|
|
|
print(context, end="") |
|
while True: |
|
inp = input("Human: ") |
|
if not inp: |
|
print("exit...") |
|
break |
|
|
|
context += "Human: " + inp + "\n" |
|
inputs = tokenizer([context]) |
|
output_ids = model.generate( |
|
inputs.input_ids, |
|
do_sample=True, |
|
temperature=0.7, |
|
max_new_tokens=96, |
|
stop=stop) |
|
outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0] |
|
try: |
|
index = outputs.index("\n", len(context)) |
|
except ValueError: |
|
outputs += "\n" |
|
index = outputs.index("\n", len(context)) |
|
|
|
outputs = outputs[:index + 1] |
|
print(outputs[len(context):], end="") |
|
context = outputs |
|
|
|
|
|
|
|
|
|
model.delete_all_weights() |
|
disk.close_copy_threads() |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--model", type=str, default="facebook/opt-6.7b", |
|
help="The model name.") |
|
parser.add_argument("--path", type=str, default="~/opt_weights", |
|
help="The path to the model weights. If there are no cached weights, " |
|
"FlexGen will automatically download them from HuggingFace.") |
|
parser.add_argument("--offload-dir", type=str, default="~/flexgen_offload_dir", |
|
help="The directory to offload tensors. ") |
|
parser.add_argument("--percent", nargs="+", type=int, |
|
default=[100, 0, 100, 0, 100, 0], |
|
help="Six numbers. They are " |
|
"the percentage of weight on GPU, " |
|
"the percentage of weight on CPU, " |
|
"the percentage of attention cache on GPU, " |
|
"the percentage of attention cache on CPU, " |
|
"the percentage of activations on GPU, " |
|
"the percentage of activations on CPU") |
|
parser.add_argument("--compress-weight", action="store_true", |
|
help="Whether to compress weight.") |
|
parser.add_argument("--compress-cache", action="store_true", |
|
help="Whether to compress cache.") |
|
args = parser.parse_args() |
|
|
|
assert len(args.percent) == 6 |
|
|
|
main(args) |
|
|
