File size: 4,957 Bytes
b3a3b15 |
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 |
# python model_sparse_generation.py --model_index 11 --attn_ckpt_dir <PATH_TO_ATTENTION_ROUTER_CHECKPOINTS> --attn_topk 0.5
# python model_sparse_generation.py --model_index 5 --mlp_ckpt_dir <PATH_TO_MLP_ROUTER_CHECKPOINTS> --attn_ckpt_dir <PATH_TO_ATTENTION_ROUTER_CHECKPOINTS> --batch_stats_dir configs/mlp_router/opt-6.7b --attn_topk 0.5
import torch
import argparse
from HybridTensor.utils.utils import _get_device
from HybridTensor.utils.activations import MODELS
from HybridTensor.models.opt import SparseConfig, build_sparse_opt
from HybridTensor.models.llama import build_sparse_llama
from HybridTensor.benchmarks.generation.gen_util import tokenize_dataset, get_random_batch
from HybridTensor.utils.activations import build_mlp_topk_lookup
from HybridTensor.routers.mlp.mlp_router_optim import load_router_dict_from_csv
from datasets import load_dataset
from transformers.models.opt import OPTConfig
from transformers import AutoTokenizer
from flash_attn.models.opt import opt_config_to_gpt2_config
from flash_attn.utils.generation import update_graph_cache
def update_router_config(model, num_layers, mlp_topk_lookup, attn_topk):
for i in range(num_layers):
if mlp_topk_lookup is not None:
model.transformer.layers[i].mlp_topk = mlp_topk_lookup[i]
# model.transformer.layers[i].mlp_topk = 512
model.transformer.layers[i].mha_router.topk = attn_topk
# dense attention in layer 0
model.transformer.layers[0].mha_router.topk = 1.0
def arg_parser():
parser = argparse.ArgumentParser(description='Inference benchmarking')
parser.add_argument('--batch_size', type=int, default=16)
parser.add_argument('--model_index', type=int, default=5)
parser.add_argument('--print_results', type=bool, default=True)
parser.add_argument('--iterations', type=int, default=1)
parser.add_argument('--gpu', type=int, default=0)
parser.add_argument('--attn_topk', type=float, default=0.5, help='Attention topk for sparse model')
parser.add_argument('--mlp_ckpt_dir', type=str, default='<PATH_TO_MLP_ROUTER_CHECKPOINTS>')
parser.add_argument('--attn_ckpt_dir', type=str, default='<PATH_TO_ATTENTION_ROUTER_CHECKPOINTS>')
parser.add_argument('--batch_stats_dir', type=str, default='configs/mlp_router/opt-6.7b')
parser.add_argument('--delta', type=int, default=256, help='Delta value for MLP topk calculation')
parser.add_argument('--use_cuda_graph', type=bool, default=False, help='Use CUDA graph for inference')
return parser.parse_args()
if __name__ == "__main__":
args = arg_parser()
model_name = MODELS[args.model_index-1]
print(f"Model name: {model_name}")
dtype = torch.float16
device= _get_device(args.gpu)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
if "llama" in model_name.lower():
tokenizer.pad_token = tokenizer.unk_token
else:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.pad_token_id
if "llama" in model_name:
model = build_sparse_llama(args, model_name,
args.attn_ckpt_dir,
device = device, dtype=dtype)
update_router_config(model, model.config.n_layer, None, args.attn_topk) # this sets the router config for all layers using a single config
else:
mlp_topk_lookup = load_router_dict_from_csv(args.batch_stats_dir, args.batch_size)
model = build_sparse_opt(args, model_name,
args.mlp_ckpt_dir,
args.attn_ckpt_dir,
device = device, dtype=dtype)
update_router_config(model, model.config.n_layer, mlp_topk_lookup, args.attn_topk) # this sets the router config for all layers using a single config
model.eval()
print(model)
# test input
input_texts = ["The future of AI is", "In a distant galaxy,"]
# input_ids = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=False).input_ids.to(device)
encoding = tokenizer(
input_texts,
return_tensors="pt",
padding=True,
truncation=False # Or True with a specified max_length if inputs can be long
)
input_ids = encoding.input_ids.to(device)
attention_mask = encoding.attention_mask.to(device) # Crucial: get the attention mask
position_ids = None
eos_token_id = tokenizer.eos_token_id
return_dict_in_generate=True,
max_length = 50
# Generate output
with torch.no_grad():
out = model.generate(
input_ids=input_ids,
max_length=max_length,
eos_token_id=eos_token_id,
return_dict_in_generate=True,
output_scores=True,
enable_timing=False,
cg=False,
)
print(tokenizer.batch_decode(out.sequences.tolist()))
|