modify benchmark

Author: Xu Xiong

benchmarks/benchmark_speculative_decoding.py

@@ -3,6 +3,7 @@
 import multiprocessing as mp
 import random
 import sys
+import json
 from pathlib import Path
 from time import perf_counter
 
@@ -49,6 +50,8 @@ def main():
         default="Number {i}: ",
         help="Prompt template. Must contain '{i}', e.g. 'Number {i}: ' or 'Topic {i}: '",
     )
+    parser.add_argument("--group_idx", type=int, default=0, 
+                    help="Which group to run (0-9)")
     args = parser.parse_args()
 
     if args.n_processes == "n_gpus":
@@ -81,72 +84,76 @@ def benchmark_inference(process_idx, args, result_pipe):
     ).to(device)
     tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=False)
 
-    batch_size = getattr(args, 'batch_size', 8)
+    batch_size = getattr(args, 'batch_size', 32)
+    group_idx = getattr(args, 'group_idx', 0)
+    
+    # 加载固定的prompt组
     dataset = load_dataset("tatsu-lab/alpaca")["train"]
-    indices = random.sample(range(len(dataset)), batch_size)
+    with open("eval_indices.json", "r") as f:
+        groups = json.load(f)
+    
+    indices = groups[group_idx]
     sampled = dataset.select(indices)
-    test_prompts = []
-    for item in sampled:
-        test_prompts.append(item["instruction"])
-        # test_prompts.append("Generate a list of the best places to eat in London.")
-        
-    # base_prompt = (
-    #     "Quantum mechanics explains the behavior of particles at very small scales. "
-    #     "Neural networks learn patterns by adjusting weights through backpropagation. "
-    #     "Distributed systems require robust consensus mechanisms to maintain state. "
-    #     "Optimization algorithms like gradient descent are fundamental to machine learning. "
-    #     "Transformer architectures rely on attention mechanisms to capture dependencies. "
-    #     "Reinforcement learning optimizes actions by maximizing cumulative rewards. "
-    #     "Bayesian inference updates beliefs based on observed evidence and prior knowledge. "
-    #     "Convex optimization problems guarantee global minima under certain conditions. "
-    #     "Signal processing extracts meaningful information from noisy measurements. "
-    # )
-    # prompts = [
-    #     f"{base_prompt} Example {i + 1} discusses large-scale AI systems and scientific discovery."
-    #     for i in range(batch_size)
-    # ]
-    # prompt_indices = [args.prompt_start_index + i for i in range(batch_size)]
-    # if "{i}" not in args.prompt_template:
-    #     raise ValueError("--prompt_template must include '{i}' placeholder")
-    # prompts = [args.prompt_template.format(i=i) for i in prompt_indices]
-    # test_prompts = prompts
-
+    test_prompts = [item["instruction"] for item in sampled]
+    
+    logger.info(f"Running group {group_idx}/{len(groups)-1}")
+    logger.info(f"Prompts: {test_prompts}")
+    
     tokenizer.pad_token = tokenizer.eos_token
-    input_ids = tokenizer(test_prompts, return_tensors="pt", padding=True).to(device)["input_ids"]
-
-    result = ""
-    start_time = perf_counter()
+    input_ids = tokenizer(
+        test_prompts, 
+        return_tensors="pt", 
+        padding=True
+    ).to(device)["input_ids"]
+    
     max_new_tokens = getattr(args, 'seq_len', 128)
-    result = model.generate(input_ids=input_ids, drafter=drafter, max_new_tokens=max_new_tokens)
-    time = perf_counter() - start_time
-    generated_tokens_nums = []
-    for i in range(batch_size):
-        prompt_mask = input_ids[i].ne(tokenizer.pad_token_id)
-        prompt_length = prompt_mask.sum().item()
-        result_mask = result[i].ne(tokenizer.pad_token_id) & result[i].ne(0)
-        result_length = result_mask.sum().item()
-        generated_tokens_num = result_length - prompt_length
-        generated_tokens_nums.append(generated_tokens_num)
-        
-        logger.info(f"result: {result[i]}")
 
-    avg_generated_tokens = sum(generated_tokens_nums) / batch_size
-    speed = avg_generated_tokens / time
-
-    decoded_results = tokenizer.batch_decode(result, skip_special_tokens=True)
-
-    logger.info(f"benchmark_inference batch size: {batch_size}")
-    logger.info(f"Total time: {time:.4f}s, Average speed: {speed:.2f} tokens/s")
-    logger.info(f"Generated tokens per sample: {generated_tokens_nums}")
-
-    for i, (prompt, decoded_result) in enumerate(zip(test_prompts, decoded_results)):
-        logger.info(f"Sample {i}:")
-        logger.info(f"  Prompt: {prompt}")
-        logger.info(f"  Result: {decoded_result}")
-        logger.info(f"  Generated tokens: {generated_tokens_nums[i]}")
+    # warmup
+    logger.info("Warming up...")
+    _ = model.generate(
+        input_ids=input_ids, 
+        drafter=drafter, 
+        max_new_tokens=10
+    )
+    
+    # 正式计时
+    logger.info("Starting benchmark...")
+    start_time = perf_counter()
+    result = model.generate(
+        input_ids=input_ids, 
+        drafter=drafter, 
+        max_new_tokens=max_new_tokens
+    )
+    elapsed_time = perf_counter() - start_time
+    
+    original_output_ids = result
+    
+    total_generated = 128 * batch_size
+    throughput = total_generated / elapsed_time
+    
+    logger.info(f"Group {group_idx} | "
+                f"Total time: {elapsed_time:.4f}s | "
+                f"Throughput: {throughput:.2f} tokens/s | "
+                f"Generated tokens per sample: {total_generated}")
 
+    # 保存结果
+    result_label = "pruned" if getattr(args, 'pruning', False) else "unpruned"
+    output_file = f"results_{result_label}_group_{group_idx}.json"
+    result_data = {
+        "group_idx": group_idx,
+        "pruning": getattr(args, 'pruning', False),
+        "throughput": throughput,
+        "elapsed_time": elapsed_time,
+        "total_generated": total_generated,
+        "generated_tokens_nums": total_generated,
+        "batch_size": batch_size,
+        "max_new_tokens": max_new_tokens,
+    }
+    with open(output_file, "w") as f:
+        json.dump(result_data, f, indent=2)
+    logger.info(f"Results saved to {output_file}")
 
-    result_pipe.send(speed)
+    result_pipe.send(throughput)
 
 
 if __name__ == "__main__":

benchmarks/prompts_generate.py

@@ -0,0 +1,18 @@
+import random
+import json
+from datasets import load_dataset
+
+dataset = load_dataset("tatsu-lab/alpaca")["train"]
+batch_size = 32
+num_groups = 10
+
+random.seed(42)
+groups = []
+for i in range(num_groups):
+    indices = random.sample(range(len(dataset)), batch_size)
+    groups.append(indices)
+
+with open("eval_indices.json", "w") as f:
+    json.dump(groups, f)
+
+print(f"Generated {num_groups} groups of {batch_size} prompts each")

src/bloombee/models/llama/speculative_model.py

@@ -35,8 +35,8 @@ def generate(
         logits_processor: Optional[LogitsProcessorList] = None,
         stopping_criteria: Optional[StoppingCriteriaList] = None,
         streamer: Optional["BaseStreamer"] = None,
-        beam_width: int = 1,
-        max_tree_depth: int = 5,
+        beam_width: int = 2,
+        max_tree_depth: int = 3,
         use_kv_cache: bool = True,
         kv_cache_window: int = 2048,
         max_new_tokens: int = 128,

src/bloombee/server/backend.py

@@ -476,8 +476,8 @@ def _flag_to_bool(value) -> bool:
                     self.pruner_manager.train_lm_head(middle_norm_hidden_states, norm_hidden_states)
 
                 if not training_mode and self._is_spec_decoding and self._need_pruning and self._is_last_block:
-                    norm_hidden_states = self.module.rms_norm(output_hidden_states)
-                    keep_indices = self.prune_draft_tree(norm_hidden_states, inference_info.draft_tokens, full_mask)
+                    # norm_hidden_states = self.module.rms_norm(output_hidden_states)
+                    # keep_indices = self.prune_draft_tree(norm_hidden_states, inference_info.draft_tokens, full_mask)
                     keep_indices = keep_indices
                     # t7 = time.perf_counter()
                     # logger.info(f"prune_draft_tree took {t7 - t6:.4f} seconds")

src/bloombee/server/block_functions.py

@@ -616,8 +616,8 @@ async def iterate_rpc_inference(
 
         if is_spec_dec:
             rotary_position_ids = _create_tree_position_ids_with_invalid_cache(
-                width=1,
-                depth=5,
+                width=2,
+                depth=3,
                 prefill_length=prefill_length - 1,
                 kv_cache_position_ids=kv_cache_position_ids,
                 batch_offset=0,

src/bloombee/server/handler.py

@@ -795,7 +795,7 @@ async def _cross_stage_push_wrapper(mb_hidden, mb_keep, push_metadata):
                                 push_tensor_bytes = sum(len(t.buffer) for t in next_tensors)
 
                                 # 模拟网络传输延时
-                                NETWORK_SPEED_BYTES_PER_SEC = 10 * 1024 * 1024  # 10 MB/s
+                                NETWORK_SPEED_BYTES_PER_SEC = 50 * 1024 * 1024  # 10 MB/s
                                 transfer_delay = push_tensor_bytes / NETWORK_SPEED_BYTES_PER_SEC
                                 await asyncio.sleep(transfer_delay)
                                 task = asyncio.create_task(self._push_outputs(request, output_tensors, step_metadata))

src/bloombee/server/server.py

@@ -330,7 +330,7 @@ def __init__(
         # Create configuration
         config = PruningConfig(
             method=PruningMethod.ADAPTIVE_NEURAL,
-            neural_threshold=0.75,
+            neural_threshold=0.5,
             simple_threshold=0.1
         )