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

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

View File

@@ -27,6 +27,7 @@ class Controller:
dtype: str,
seed: int,
model_path: str,
use_dummy_weights: bool,
max_num_batched_tokens: int,
) -> None:
self.stage_id = stage_id
@@ -58,6 +59,7 @@ class Controller:
tensor_parallel_size=tensor_parallel_size,
pipeline_parallel_size=pipeline_parallel_size,
model_path=model_path,
use_dummy_weights=use_dummy_weights,
max_num_batched_tokens=max_num_batched_tokens,
)
self.workers.append(worker)