From 8bb96edebfb806d921396ba5ad2bbb211479a7ba Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 20 Feb 2024 20:15:54 +0000 Subject: [PATCH] validate num_shards in engine creation --- serve/mlc_serve/engine/base.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/serve/mlc_serve/engine/base.py b/serve/mlc_serve/engine/base.py index 835874d72f..de1fba57c0 100644 --- a/serve/mlc_serve/engine/base.py +++ b/serve/mlc_serve/engine/base.py @@ -59,6 +59,11 @@ def get_engine_config(dict_config): assert (engine_config.min_decode_steps > 0) and (engine_config.max_decode_steps > 0) assert engine_config.max_decode_steps > engine_config.min_decode_steps + if engine_config.model_type == "torch": + assert ( + engine_config.num_shards is not None + ), "num_shards in MLCServeEngineConfig needs to be provided for PT models." + return engine_config