biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add an option to use dummy model weights (#33)

Browse Source
This commit is contained in:
Woosuk Kwon
2023-04-08 23:36:12 -07:00
committed by GitHub
parent c267b1a02c
commit ee88a7e5f3
9 changed files with 36 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
cacheflow/worker/worker.py
View File

@@ -29,6 +29,7 @@ class Worker:
rank: int,
world_size: int,
model_path: str,
use_dummy_weights: bool,
max_num_batched_tokens: int,
tensor_parallel_size: int = 1,
pipeline_parallel_size: int = 1,
@@ -43,8 +44,8 @@ class Worker:
set_random_seed(seed)
# Initialize the model.
self.model, self.dtype = get_model(model_name, dtype=dtype, path=model_path)
self.model = self.model.cuda()
self.model, self.dtype = get_model(
model_name, dtype=dtype, path=model_path, use_dummy_weights=use_dummy_weights)
tensor_model_parallel_world_size = (
get_tensor_model_parallel_world_size())
initialize_all_reduce_launcher(