From 12898b5bef3e6376d3ceda4964c03b07fe4dface Mon Sep 17 00:00:00 2001 From: cm2435 Date: Mon, 12 Feb 2024 22:31:05 +0000 Subject: [PATCH] adding tools to be able to profile model fwds to see what to turn into kernels (cherry picked from commit 6db5b126b6f882ea1ba7bb5433ef26c0c1c1cb1d) --- tests/profiles/profile_phi2.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 tests/profiles/profile_phi2.py diff --git a/tests/profiles/profile_phi2.py b/tests/profiles/profile_phi2.py new file mode 100644 index 0000000000..39708062b7 --- /dev/null +++ b/tests/profiles/profile_phi2.py @@ -0,0 +1,27 @@ +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from unsloth.kernels.utils import profile_generate_method + +torch.set_default_device("cuda") + +model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True) +tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True) + +inputs = tokenizer('''def print_prime(n): + """ + Print all primes between 1 and n + """''', return_tensors="pt", return_attention_mask=False) + + +generate_args = { + **inputs, # Assuming model_inputs is a dictionary with appropriate keys + "max_new_tokens": 100, + "do_sample": True +} + +# Ensure your model and tokenizer are properly loaded and set up as before. + +# Now, call the profile_generate_method function +prof = profile_generate_method(model, generate_args) +