[model][utils] add extract_layer_index utility function (#10599)

Signed-off-by: youkaichao <youkaichao@gmail.com>

vllm/model_executor/models/qwen2_moe.py

@@ -53,7 +53,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils import print_warning_once
 
 from .interfaces import SupportsPP
-from .utils import (is_pp_missing_parameter,
+from .utils import (extract_layer_index, is_pp_missing_parameter,
                     make_empty_intermediate_tensors_factory, make_layers,
                     maybe_prefix)
 
@@ -244,7 +244,6 @@ class Qwen2MoeDecoderLayer(nn.Module):
     def __init__(
         self,
         config: PretrainedConfig,
-        layer_idx: int,
         cache_config: Optional[CacheConfig] = None,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
@@ -269,6 +268,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
 
         # Note: Qwen/Qwen2-57B-A14B-Instruct does not have
         # `mlp_only_layers` in the config.
+        layer_idx = extract_layer_index(prefix)
         mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else
                            config.mlp_only_layers)
         if (layer_idx not in mlp_only_layers) and (
@@ -337,8 +337,6 @@ class Qwen2MoeModel(nn.Module):
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
             lambda prefix: Qwen2MoeDecoderLayer(config=config,
-                                                layer_idx=int(
-                                                    prefix.split(".")[-1]),
                                                 cache_config=cache_config,
                                                 quant_config=quant_config,
                                                 prefix=prefix),