From 875afe38ab3b622ea2e84162799d5368cfb05257 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 12 Oct 2023 01:05:37 -0700 Subject: [PATCH] Add blacklist in model checkpoint (#1325) --- vllm/model_executor/weight_utils.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/weight_utils.py b/vllm/model_executor/weight_utils.py index 74de968422962..027e741c5f6bf 100644 --- a/vllm/model_executor/weight_utils.py +++ b/vllm/model_executor/weight_utils.py @@ -144,8 +144,18 @@ def prepare_hf_model_weights( for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) if not use_safetensors: + # Exclude files that are not needed for inference. + # https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 + blacklist = [ + "training_args.bin", + "optimizer.bin", + "optimizer.pt", + "scheduler.pt", + "scaler.pt", + ] hf_weights_files = [ - x for x in hf_weights_files if not x.endswith("training_args.bin") + f for f in hf_weights_files + if not any(f.endswith(x) for x in blacklist) ] if len(hf_weights_files) == 0 and use_safetensors and fall_back_to_pt: