Update

Author: jianc99

result.txt

@@ -63,9 +63,9 @@
     if rank != args.rank_group[0]:
         print = lambda *args, **kwargs: None
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write(f"Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}\n")
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write(f"Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}\n")
 
 setup_seed(args.seed)
 print(f"Using device={DEVICE}")
@@ -316,7 +316,7 @@
     if use_tp:
         dist.barrier()
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
-        file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
+#         file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))

run.sh

@@ -51,9 +51,9 @@
     if rank != args.rank_group[0]:
         print = lambda *args, **kwargs: None
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write(f"SnapKV-Selfspec: Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}; Draft budget:{args.draft_budget}\n")
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write(f"SnapKV-Selfspec: Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}; Draft budget:{args.draft_budget}\n")
 
 setup_seed(args.seed)
 print(f"Using device={DEVICE}")
@@ -232,7 +232,7 @@
     if use_tp:
         dist.barrier()
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
-        file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
+#         file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))

scripts/prepare_ruler.sh

@@ -60,9 +60,9 @@
     if rank != args.rank_group[0]:
         print = lambda *args, **kwargs: None
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write(f"Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}; Draft budget:{args.draft_budget}\n")
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write(f"Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}; Draft budget:{args.draft_budget}\n")
 
 setup_seed(args.seed)
 print(f"Using device={DEVICE}")
@@ -313,7 +313,7 @@
     if use_tp:
         dist.barrier()
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
-        file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
+#         file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))

scripts/run_selfspec.sh

@@ -48,9 +48,9 @@
     if rank != args.rank_group[0]:
         print = lambda *args, **kwargs: None
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write(f"Selfspec: Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}; Draft budget:{args.draft_budget}\n")
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write(f"Selfspec: Prefix:{args.prefix_len}; Bsz:{args.B}; Gamma:{args.gamma}; Draft budget:{args.draft_budget}\n")
 
 
 setup_seed(args.seed)
@@ -259,7 +259,7 @@
     if use_tp:
         dist.barrier()
 
-if rank == 0:
-    with open("result.txt", "a") as file:
-        file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
-        file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))
+# if rank == 0:
+#     with open("result.txt", "a") as file:
+#         file.write("total time :{:.5f}s, time per iter :{:.5f}s, decoding step: {}, large model step: {}, avg latency: {} \n".format(total_time, total_time / target_steps, num_gen_tokens, target_steps, total_time / num_gen_tokens * BATCH_SIZE))
+#         file.write("target time :{:.5f}s, draft time :{:.5f}s, verify loop : {}, avg generate len per sentence: {} \n".format(target_time/target_steps, draft_time / target_steps, verify_loop/target_steps, num_gen_tokens/target_steps/BATCH_SIZE))

tests/SnapKV/longspec_benchmark.py

@@ -95,7 +95,7 @@
         for i in range(BATCH_SIZE):
             print(tokenizer.decode(output[i, args.prefix_len:]))
     print(f"Tokens per second :{total_time/model_steps}")
-    if step < 10:
+    if step < 5:
         total_time = 0.0
         model_steps = 0
     if use_tp: