From 4cf256ae7f8b0be8f06f6b85821e55d4f5bdaa13 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Mon, 15 Jul 2024 10:32:35 -0700 Subject: [PATCH] [misc][distributed] fix pp missing layer condition (#6446) --- vllm/model_executor/models/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index a0d2a0286ff67..c135b20352203 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -83,7 +83,10 @@ def get_pp_missing_layer_names(model: torch.nn.Module) -> List[str]: missing_layer_names = [] for name, module in model.named_modules(): if isinstance(module, PPMissingLayer): - missing_layer_names.append(name) + # NOTE: the trailing dot is used to match the prefix of the layer. + # without the dot, we could match a layer that is not missing, + # e.g., 'encoder.layer.1' would match 'encoder.layer.11' + missing_layer_names.append(name + '.') _model_to_pp_missing_layer_names[model_id] = missing_layer_names return missing_layer_names