Simplify Tensor Parallel

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:

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")