Merge pull request #14 from TianyiQ/main

fix(evaluation): allow manual spec of tensor parallelization
PKU-Alignment · Nov 16, 2024 · c283c3c · c283c3c
2 parents a6a0448 + 2ee86f8
commit c283c3c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/abstractions/backends.py b/src/abstractions/backends.py
@@ -344,6 +344,10 @@ def vllm_process_batch(
                 min_gpus_per_instance = (
                     2 if model_size <= 30 else 4 if model_size <= 80 else 8
                 )
+
+                if os.environ.get("FORCE_TP"):
+                    min_gpus_per_instance = int(os.environ.get("FORCE_TP"))
+
                 assert num_gpus % min_gpus_per_instance == 0
                 args = [
                     "python",