fix

huggingface · Dec 11, 2024 · 6738b85 · 6738b85
1 parent d5cd55c
commit 6738b85
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/optimum_benchmark/backends/tensorrt_llm/backend.py b/optimum_benchmark/backends/tensorrt_llm/backend.py
@@ -137,7 +137,7 @@ def trtllm_kwargs(self):
         if self.config.max_prompt_length is not None:
             kwargs["max_prompt_length"] = self.config.max_prompt_length
 
-        if self.config.tp is not None:
+        if self.config.max_new_tokens is not None:
             kwargs["max_new_tokens"] = self.config.max_new_tokens
 
         if self.config.max_beam_width is not None:

diff --git a/optimum_benchmark/backends/tensorrt_llm/config.py b/optimum_benchmark/backends/tensorrt_llm/config.py
@@ -4,7 +4,7 @@
 from ...import_utils import tesnorrt_llm_version
 from ..config import BackendConfig
 
-SUPPORTED_DTYPES = ["float16", "bfloat16", "float32"]
+SUPPORTED_DTYPES = [None, "float16", "bfloat16", "float32"]
 
 
 @dataclass
@@ -38,8 +38,13 @@ def __post_init__(self) -> None:
         if self.dtype not in SUPPORTED_DTYPES:
             raise ValueError(f"dtype must be one of float16, bfloat16, float32, got {self.dtype}")
 
-        if self.gpus_per_node != self.world_size:
+        if self.gpus_per_node is not None and self.world_size is not None and self.gpus_per_node != self.world_size:
             raise ValueError(f"gpus_per_node ({self.gpus_per_node}) != world_size ({self.world_size})")
 
-        if self.world_size != self.pp * self.tp:
+        if (
+            self.world_size is not None
+            and self.pp is not None
+            and self.tp is not None
+            and self.world_size != self.pp * self.tp
+        ):
             raise ValueError(f"world_size ({self.gpus_per_node}) != pp ({self.pp}) * tp ({self.tp})")