diffusers/tools

1#!/usr/bin/env python32from transformers import AutoTokenizer, AutoModelForCausalLM3import time4import torch5 6DEVICE = "cuda:1"7 8tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")9model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True)10model.to(DEVICE)11 12 13# forward14print("Forward benchmarks")15print(50 * "=")16 17for batch_size in (1, 4, 16):18    for input_seq in (4, 16, 256):19        input_ids = torch.ones((batch_size, input_seq), dtype=torch.long, device=DEVICE)20        attention_mask = torch.ones_like(input_ids)21        attention_mask[0, 3] = 022 23        times = []24        for _ in range(3):25            start_time = time.time()26            with torch.no_grad():27                logits = model(input_ids=input_ids, attention_mask=attention_mask).logits28            times.append(time.time() - start_time)29 30        result = min(times)31 32        print(f"Forward bsz={batch_size}, input_seq={input_seq}: {result}")33 34 35# generate36print("Generate benchmarks")37print(50 * "=")38 39for batch_size in (1, 16):40    for input_seq in (4, 256):41        input_ids = torch.ones((batch_size, input_seq), dtype=torch.long, device=DEVICE)42        attention_mask = torch.ones_like(input_ids)43        attention_mask[0, 3] = 044 45        times = []46        for _ in range(3):47            start_time = time.time()48            out = model.generate(input_ids=input_ids, max_new_tokens=256, do_sample=False)49            times.append(time.time() - start_time)50 51        result = min(times)52 53        print(f"Generate bsz={batch_size}, input_seq={input_seq}: {result}")54