diff --git a/README.md b/README.md index f0933a2..28fc0d3 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,6 @@ Development Image: ```runpod/worker-vllm:dev``` - Tensor Parallelism: Note that the more GPUs you split a model's weights accross, the slower it will be due to inter-GPU communication overhead. If you can fit the model on a single GPU, it is recommended to do so. - - `USE_TENSOR_PARALLEL`: Enable (`1`) or disable (`0`) Tensor Parallelism. (default: `0`) - `TENSOR_PARALLEL_SIZE`: Number of GPUs to shard the model across (default: `1`). - System Settings: diff --git a/src/engine.py b/src/engine.py index 2c9140c..5b53e3f 100644 --- a/src/engine.py +++ b/src/engine.py @@ -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")