File size: 1,560 Bytes
c1a41d7 |
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 |
import argparse
import os
import glog
import torch
from torch.profiler import profile, record_function, ProfilerActivity
from transformers import AutoTokenizer
from lib.utils.unsafe_import import model_from_hf_path
import time
torch.set_grad_enabled(False)
parser = argparse.ArgumentParser()
parser.add_argument('--hf_path', default='meta-llama/Llama-2-70b-hf', type=str)
parser.add_argument('--batch_size', default=1, type=int)
parser.add_argument('--seqlen', default=1, type=int)
parser.add_argument('--samples', default=100, type=int)
parser.add_argument('--no_use_cuda_graph', action='store_true')
parser.add_argument('--no_use_flash_attn', action='store_true')
def main(args):
model, model_str = model_from_hf_path(args.hf_path,
use_cuda_graph=not args.no_use_cuda_graph,
use_flash_attn=not args.no_use_flash_attn)
tokenizer = AutoTokenizer.from_pretrained(model_str)
prompt = 'It is a truth universally acknowledged that'
inputs = tokenizer(prompt, return_tensors='pt')
token = inputs['input_ids'][0:1, 0:1].cuda().repeat(args.batch_size, args.seqlen)
model(token, use_cache=False)
torch.cuda.synchronize()
start = time.time()
for _ in range(args.samples):
model(token, use_cache=False)
torch.cuda.synchronize()
end = time.time()
print('TIME', (end - start) / args.samples)
if __name__ == '__main__':
torch.set_grad_enabled(False)
torch.manual_seed(0)
args = parser.parse_args()
main(args)
|