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
README.md
-1
View File
@@ -70,7 +70,6 @@ Development Image: ```runpod/worker-vllm:dev```
- Tensor Parallelism: - 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. 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`). - `TENSOR_PARALLEL_SIZE`: Number of GPUs to shard the model across (default: `1`).
- System Settings: - System Settings:
src/engine.py
+5 -6
View File
@@ -205,13 +205,12 @@ class vLLMEngine:
return model, download_dir return model, download_dir
def _get_num_gpu_shard(self): def _get_num_gpu_shard(self):
final_num_gpu_shard = 1 num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
if bool(int(os.getenv("USE_TENSOR_PARALLEL", 0))): if num_gpu_shard > 1:
env_num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
num_gpu_available = device_count() num_gpu_available = device_count()
final_num_gpu_shard = min(env_num_gpu_shard, num_gpu_available) num_gpu_shard = min(num_gpu_shard, num_gpu_available)
logging.info("Using %s GPU shards", final_num_gpu_shard) logging.info("Using %s GPU shards", num_gpu_shard)
return final_num_gpu_shard return num_gpu_shard
def _get_max_model_len(self): def _get_max_model_len(self):
max_model_len = os.getenv("MAX_MODEL_LEN") max_model_len = os.getenv("MAX_MODEL_LEN")