From 9f7d9119b1ec5c7cb2ea1ac29e7455cdd16e9fa5 Mon Sep 17 00:00:00 2001 From: jayeshthk Date: Thu, 6 Mar 2025 23:07:49 +0530 Subject: [PATCH] add timing measurements for autoregressive and speculative sampling methods --- main.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/main.py b/main.py index 3efd6ae..b67b97f 100644 --- a/main.py +++ b/main.py @@ -4,7 +4,7 @@ import contexttimer from colorama import Fore, Style from transformers import AutoTokenizer, AutoModelForCausalLM - +import time from sampling import autoregressive_sampling, speculative_sampling, speculative_sampling_v2 from globals import Decoder @@ -95,32 +95,40 @@ def generate(input_text, approx_model_name, target_model_name, num_tokens=20, ga top_p = 0.9 torch.manual_seed(123) + st=time.time() output = autoregressive_sampling(input_ids, large_model, num_tokens, top_k = top_k, top_p=top_p) generated_text = tokenizer.decode(output[0], skip_special_tokens=True) - color_print(f"large (target) model autoregressive_sampling: {generated_text}") + tot=time.time()-st + color_print(f"large (target) model autoregressive_sampling {tot:.4f}s: {generated_text}") if use_benchmark: benchmark(autoregressive_sampling, "AS_large", use_profiling, input_ids, large_model, num_tokens, top_k = top_k, top_p=top_p) torch.manual_seed(123) + st=time.time() output = autoregressive_sampling(input_ids, small_model, num_tokens, top_k = top_k, top_p=top_p) generated_text = tokenizer.decode(output[0], skip_special_tokens=True) - color_print(f"small (approx) model autoregressive_sampling: {generated_text}") + tot=time.time()-st + color_print(f"small (approx) model autoregressive_sampling {tot:.4f}s: {generated_text}") if use_benchmark: benchmark(autoregressive_sampling, "AS_small", use_profiling, input_ids, small_model, num_tokens, top_k = top_k, top_p=top_p) torch.manual_seed(123) + st=time.time() output = speculative_sampling_v2(input_ids, small_model, large_model, num_tokens, top_k = top_k, top_p=top_p, random_seed = random_seed) generated_text = tokenizer.decode(output[0], skip_special_tokens=True) - color_print(f"deepmind's speculative_sampling: {generated_text}") + tot=time.time()-st + color_print(f"deepmind's speculative_sampling {tot:.4f}s: {generated_text}") torch.manual_seed(123) + st=time.time() output = speculative_sampling(input_ids, small_model, large_model, num_tokens, gamma = gamma, top_k = top_k, top_p=top_p, random_seed = random_seed, verbose = verbose) generated_text = tokenizer.decode(output[0], skip_special_tokens=True) - color_print(f"google's speculative_sampling: {generated_text}") + tot=time.time()-st + color_print(f"google's speculative_sampling {tot:.4f}s: {generated_text}") if use_benchmark: benchmark(speculative_sampling, "SP", use_profiling,