File size: 4,406 Bytes
b943e32 |
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 |
import json
from argparse import ArgumentParser
import datasets
import torch
import transformers
from transformers import AutoModelForCausalLM, BatchEncoding
"""
Usage examples (with the best batch sizes on A100-80GB-400W)
============================================================
python -m benchmark_hf_model --model_name_or_path="Deci/DeciLM-7B" --batch_size=352
python -m benchmark_hf_model --model_name_or_path="mistralai/Mistral-7B-v0.1" --batch_size=192 --model_kwargs_json='{"use_flash_attention_2": true}'
python -m benchmark_hf_model --model_name_or_path="meta-llama/Llama-2-7b-hf" --batch_size=48 --model_kwargs_json='{"use_flash_attention_2": true}'
"""
def parse_args():
parser = ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
required=True,
)
parser.add_argument(
"--warmup_iters",
type=int,
default=10,
)
parser.add_argument(
"--iterations",
type=int,
default=5,
)
parser.add_argument(
"--batch_size",
type=int,
default=32,
)
parser.add_argument(
"--prompt_length",
type=int,
default=512,
)
parser.add_argument(
"--max_new_tokens",
type=int,
default=512,
)
parser.add_argument(
"--precision",
type=str,
default="bf16",
help="Model precision, from: fp32, fp16 or bf16",
)
parser.add_argument(
"--model_kwargs_json",
type=str,
default=None,
)
return parser.parse_args()
def main():
args = parse_args()
transformers.logging.set_verbosity_error()
datasets.logging.set_verbosity_error()
dict_precisions = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}
if args.precision not in dict_precisions:
raise ValueError(
f"Non valid precision {args.precision}, choose from: fp16, fp32, bf16"
)
dtype = dict_precisions[args.precision]
model_kwargs = {}
if args.model_kwargs_json is not None:
model_kwargs = json.loads(args.model_kwargs_json)
print(f"loading model...")
model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, trust_remote_code=True,
torch_dtype=dtype, **model_kwargs)
try:
print(model.model.layers[0].self_attn)
except:
print("couldn't print the model's attention module")
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
model.cuda()
model.eval()
prompt = torch.ones(args.prompt_length, dtype=torch.long)
inputs = BatchEncoding({"input_ids": prompt.repeat(args.batch_size, 1)})
inputs = inputs.to(model.device)
# warmup
print(f"warming up for {args.warmup_iters} iterations...")
for _ in range(args.warmup_iters):
with torch.no_grad():
_ = model.generate(
**inputs,
max_new_tokens=1,
do_sample=False,
eos_token_id=-1234,
)
print('finished warmup')
torch.cuda.synchronize()
print(
f"prefill ({args.prompt_length} tokens{f' x {args.batch_size} batch' if args.batch_size > 1 else ''}) + generation ({args.max_new_tokens} tokens{f' x {args.batch_size} batch' if args.batch_size > 1 else ''}):")
tokens_generated = args.max_new_tokens * args.batch_size
prefill_and_generation = []
for gen_iter in range(args.iterations):
starter.record()
with torch.no_grad():
_ = model.generate(
**inputs,
max_new_tokens=args.max_new_tokens,
do_sample=False,
eos_token_id=-1234,
)
ender.record()
torch.cuda.synchronize()
t = starter.elapsed_time(ender) / 1000
prefill_and_generation.append(t)
print(f" iter {gen_iter + 1}: {t:.03f} sec total, {tokens_generated / t:.02f} generated tokens/sec")
aver = sum(prefill_and_generation) / len(prefill_and_generation)
print(f" average: {aver:.03f} sec total, {tokens_generated / aver:.02f} generated tokens/sec")
print(f"These results are obtained for model '{args.model_name_or_path}' with {args.batch_size=}.")
if __name__ == "__main__":
main()
|