From 180c15e3ab2400dd27e75be82d43d6af698c97a4 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Fri, 10 Jul 2026 20:20:39 +0800 Subject: [PATCH] fix: fix test benchmark --- test/bench/test_benchmark.py | 157 ++++++++++------------------------- 1 file changed, 44 insertions(+), 113 deletions(-) diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index 60e63dc66..23ca26506 100644 --- a/test/bench/test_benchmark.py +++ b/test/bench/test_benchmark.py @@ -1,15 +1,13 @@ -import sys -import os -import time -import re -import csv import argparse +import csv import json -import numpy as np -from datasets import load_dataset, Dataset +import os +import re +import time from abc import ABC, abstractmethod + +from datasets import Dataset, load_dataset from infinilm.base_config import BaseConfig -from infinilm.processors import AutoInfinilmProcessor TOTAL_TOKENS = 0 TOTAL_TIME = 0.0 @@ -43,7 +41,7 @@ def _generate_step(self, tokens, max_steps, topp_, topk_, temperature_): class InfiniLMBenchmark(BaseBenchmark): - """Wrapper class for InfiniLM cpp backend for benchmark evaluation""" + """InfiniLM cpp backend using the scheduler-backed high-level API.""" def __init__( self, @@ -56,18 +54,10 @@ def __init__( enable_graph=False, attn_backend="default", ): - import transformers - import infinicore - from infinilm.modeling_utils import load_model_state_dict_by_file - from infinilm.distributed import DistConfig - from infinilm.cache import StaticKVCacheConfig, PagedKVCacheConfig - from infinilm.infer_engine import InferEngine + from infinilm import LLM self.benchmark = benchmark - # Map device type string to infinicore device - # Note: These map to the Python device type strings used by infinicore.device() - # which correspond to _TORCH_DEVICE_MAP values in InfiniCore/python/infinicore/device.py device_map = { "cpu": "cpu", "nvidia": "cuda", @@ -79,21 +69,16 @@ def __init__( "kunlun": "cuda", "hygon": "cuda", "ali": "cuda", + "cuda": "cuda", + "mlu": "mlu", + "musa": "musa", + "npu": "npu", } - device_name = device_map.get(device_type_str.lower(), "cpu") - # CUDA_VISIBLE_DEVICES is automatically respected by CUDA runtime API - # When CUDA_VISIBLE_DEVICES=5 is set, CUDA only sees device 5 as device 0 - # So device index 0 will automatically map to the first visible device - self.device = infinicore.device(device_name, 0) - # Load config and tokenizer with open(os.path.join(model_dir_path, "config.json"), "r") as f: self.config_dict = json.load(f) - self.processor = AutoInfinilmProcessor.from_pretrained(model_dir_path) - self.tokenizer = self.processor.get_tokenizer() - eos_token_id = self.config_dict.get("eos_token_id") self.eos_token_id = ( [eos_token_id] if isinstance(eos_token_id, int) else eos_token_id @@ -105,113 +90,60 @@ def __init__( if enable_paged_attn and attn_backend == "default": attn_backend = "paged-attn" - # Create model with cpp backend print("Loading model with cpp backend...") print(f"Graph compilation: {'enabled' if enable_graph else 'disabled'}") print(f"Attention backend: {attn_backend}") - self.model = InferEngine( - model_dir_path, - device=self.device, - distributed_config=DistConfig(ndev), - cache_config=( - PagedKVCacheConfig(128) if enable_paged_attn else StaticKVCacheConfig() - ), - enable_graph_compiling=enable_graph, - attention_backend=attn_backend, - ) - - # Enable KV cache for generation - self.model.use_cache = True - - # Load weights - print("Loading model weights...") - load_model_state_dict_by_file( - self.model, - model_dir_path, - dtype=self.model.dtype, + self.model = LLM( + model_path=model_dir_path, + device=device_name, + tensor_parallel_size=ndev, + cache_type="paged" if enable_paged_attn else "static", + max_batch_size=1, + num_blocks=128, + block_size=256, + enable_graph=enable_graph, + attn_backend=attn_backend, ) + self.processor = self.model.engine.processor + self.tokenizer = self.processor.get_tokenizer() print("Model loaded successfully") def max_context_len(self): return self.config_dict.get("max_position_embeddings", 2048) def render_input_content(self, *args, **kwargs): - """Render input content based on benchmark type""" if self.benchmark == "ceval": return render_ceval(self.processor, *args, **kwargs) - elif self.benchmark == "mmlu": + if self.benchmark == "mmlu": return render_mmlu(self.processor, *args, **kwargs) - else: - raise ValueError(f"Unknown benchmark: {self.benchmark}") + raise ValueError(f"Unknown benchmark: {self.benchmark}") def generate(self, *args, max_steps=500, topp_=1.0, topk_=1, temperature_=1.0): - """Generate response based on benchmark type""" - # Render input content input_content = self.render_input_content(*args) print(input_content, end="", flush=True) + return self._generate_step(input_content, max_steps, topp_, topk_, temperature_) - # Encode input - tokens = self.encode_text(input_content) - - # Delegate to backend-specific generation implementation - output_content = self._generate_step( - tokens, max_steps, topp_, topk_, temperature_ - ) - - return output_content - - def _generate_step(self, tokens, max_steps, topp_, topk_, temperature_): - """ - InfiniLM cpp backend-specific generation implementation - - NOTE: Validation confirmed input configs are identical between backends. - The issue was that manual generation loop called InferEngine.generate() which - doesn't maintain KV cache. Solution: Use model's built-in generate() method - which properly handles KV cache through GenerationMixin. - """ - # Convert tokens to infinicore format - import infinicore - from infinilm.infer_engine import GenerationConfig - - input_ids_list = [tokens] - input_ids = infinicore.from_list(input_ids_list, dtype=infinicore.int64) + def _generate_step(self, prompt, max_steps, topp_, topk_, temperature_): + from infinilm import SamplingParams start_time = time.perf_counter() - - # For cpp backend, reset cache before generation if use_cache is enabled - if ( - self.model.use_cache - and hasattr(self.model, "_model") - and hasattr(self.model._model, "reset_cache") - ): - batch_size = input_ids.shape[0] - seq_len = input_ids.shape[1] - max_cache_len = max_steps + seq_len - self.model.reset_cache( - batch_size=batch_size, initial_capacity=max_cache_len - ) - - # Use model's built-in generate() method - output_ids = self.model.generate( - input_ids=input_ids, - generation_config=GenerationConfig( - max_new_tokens=max_steps, + request_output = self.model.generate( + prompts=prompt, + sampling_params=SamplingParams( + max_tokens=max_steps, temperature=temperature_, top_k=topk_, top_p=topp_, ), - ) - + use_tqdm=False, + )[0] end_time = time.perf_counter() - # ---- post process ---- - generated_ids = np.array([output_id.to_numpy()[0] for output_id in output_ids]) - output_text = self.tokenizer.decode(generated_ids) - - # ---- stats ---- - input_tokens = len(tokens) - new_tokens = generated_ids.size + completion = request_output.outputs[0] + output_text = completion.text + input_tokens = len(request_output.prompt_token_ids or []) + new_tokens = len(completion.token_ids) total_tokens = input_tokens + new_tokens total_time = end_time - start_time @@ -226,11 +158,9 @@ def _generate_step(self, tokens, max_steps, topp_, topk_, temperature_): global TOTAL_TOKENS, TOTAL_TIME TOTAL_TOKENS += total_tokens TOTAL_TIME += total_time - return output_text def destroy_model_instance(self): - # Cleanup if needed del self.model print("Model destroyed") @@ -298,9 +228,10 @@ def render_input_content(self, *args, **kwargs): raise ValueError(f"Unknown benchmark: {self.benchmark}") def _generate_step(self, tokens, max_steps, topp_, topk_, temperature_): - import torch import time + import torch + input_ids = torch.tensor([tokens], device=self.device) if self.device.type == "cuda": @@ -763,7 +694,7 @@ def load_one(subj): continue if not all_samples: raise FileNotFoundError( - f"No MMLU cached data found for any subject. Please ensure datasets are cached." + "No MMLU cached data found for any subject. Please ensure datasets are cached." ) return all_samples, "all" @@ -1005,7 +936,7 @@ def _load_mmlu_subject(subj): continue if not samples: raise FileNotFoundError( - f"No MMLU data found for any subject in the list" + "No MMLU data found for any subject in the list" ) return samples, "all" else: @@ -1119,7 +1050,7 @@ def main(): else: # cpp backend model = InfiniLMBenchmark( cfg.model, - device_type_str, + cfg.get_device_str(device_type_str), cfg.tp, cfg.backend, cfg.bench,