diff --git a/test/bench/backends/transformers.py b/test/bench/backends/transformers.py index 7464d731..2958d9da 100644 --- a/test/bench/backends/transformers.py +++ b/test/bench/backends/transformers.py @@ -8,7 +8,13 @@ class TransformersBenchmark(BaseBenchmark): """Hugging Face Transformers backend.""" - def __init__(self, model_dir_path, device_type_str="cpu", benchmark="ceval"): + def __init__( + self, + model_dir_path, + device_type_str="cpu", + tensor_parallel_size=1, + benchmark="ceval", + ): import torch import transformers @@ -20,6 +26,7 @@ def __init__(self, model_dir_path, device_type_str="cpu", benchmark="ceval"): f"Transformers backend unsupported device type: {device_type_str}" ) self.device = torch.device(device_type_str) + self.tensor_parallel_size = tensor_parallel_size with open(os.path.join(model_dir_path, "config.json"), "r") as f: self.config_dict = json.load(f) @@ -29,12 +36,36 @@ def __init__(self, model_dir_path, device_type_str="cpu", benchmark="ceval"): ) print("Loading model with Transformers backend...") + load_kwargs = { + "dtype": "auto", + "trust_remote_code": True, + } + if tensor_parallel_size > 1: + if device_type_str != "cuda": + raise ValueError( + "Transformers multi-GPU evaluation requires a CUDA device. " + f"Got device_type_str={device_type_str!r}." + ) + available_devices = torch.cuda.device_count() + if available_devices < tensor_parallel_size: + raise ValueError( + f"Requested tp={tensor_parallel_size}, but only " + f"{available_devices} CUDA devices are visible." + ) + load_kwargs["device_map"] = "auto" + print( + "Transformers multi-GPU device_map=auto enabled " + f"for {tensor_parallel_size} visible CUDA devices" + ) + self.model = transformers.AutoModelForCausalLM.from_pretrained( model_dir_path, - dtype="auto", - trust_remote_code=True, - ).to(self.device) + **load_kwargs, + ) + if tensor_parallel_size <= 1: + self.model = self.model.to(self.device) self.model.eval() + self.input_device = self.model.get_input_embeddings().weight.device print("Transformers model loaded successfully") eos_token_id = self.config_dict.get("eos_token_id") @@ -56,7 +87,7 @@ def generate(self, *args, max_steps=500, topp_=1.0, topk_=1, temperature_=1.0): prompt = self.render_input_content(*args) print(prompt, end="", flush=True) tokens = self.encode_text(prompt) - input_ids = torch.tensor([tokens], device=self.device) + input_ids = torch.tensor([tokens], device=self.input_device) self._synchronize() start_time = time.perf_counter() diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index 894726ac..b77b11cd 100644 --- a/test/bench/test_benchmark.py +++ b/test/bench/test_benchmark.py @@ -601,10 +601,7 @@ def main(): device_str = cfg.get_device_str(device_type_str) if cfg.backend in {"transformers", "torch"}: - assert cfg.tp == 1, ( - "Transformers backend only supports single-device evaluation" - ) - model = TransformersBenchmark(cfg.model, device_str, cfg.bench) + model = TransformersBenchmark(cfg.model, device_str, cfg.tp, cfg.bench) elif cfg.backend == "vllm": model = VLLMBenchmark(cfg.model, device_str, cfg.tp, cfg.bench) elif cfg.backend in {"infinilm", "cpp", "python"}: