owen/worker-vllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Simplify Tensor Parallel

Browse Source
This commit is contained in:
alpayariyak
2024-01-31 03:28:02 +00:00
parent 12d6f0778e
commit 46eee12819
2 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine.py
+5 -6
View File
@@ -205,13 +205,12 @@ class vLLMEngine:
return model, download_dir
def _get_num_gpu_shard(self):
final_num_gpu_shard = 1
if bool(int(os.getenv("USE_TENSOR_PARALLEL", 0))):
env_num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
if num_gpu_shard > 1:
num_gpu_available = device_count()
final_num_gpu_shard = min(env_num_gpu_shard, num_gpu_available)
logging.info("Using %s GPU shards", final_num_gpu_shard)
return final_num_gpu_shard
num_gpu_shard = min(num_gpu_shard, num_gpu_available)
logging.info("Using %s GPU shards", num_gpu_shard)
return num_gpu_shard
def _get_max_model_len(self):
max_model_len = os.getenv("MAX_MODEL_LEN")