| |
| |
| |
| |
| @@ -110,10 +110,23 @@ class Glm5NextMultiTokenPredictorLayer(nn.Module): |
| return hidden_states, hidden_states |
| |
| |
| +def _mtp_hf_config(vllm_config: VllmConfig): |
| + """The MTP head's own config. With a standalone draft checkpoint |
| + (``--speculative-config '{"model": ...}'``) the draft config carries |
| + num_nextn_predict_layers / n_routed_experts, while vllm_config.model_config |
| + is still the *target* (which may have num_nextn_predict_layers == 0, e.g. an |
| + checkpoint that ships without an MTP layer). Fall back to the |
| + target config for in-checkpoint MTP layers.""" |
| + spec = vllm_config.speculative_config |
| + if spec is not None and spec.draft_model_config is not None: |
| + return spec.draft_model_config.hf_config |
| + return vllm_config.model_config.hf_config |
| + |
| + |
| class Glm5NextMultiTokenPredictor(nn.Module): |
| def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| super().__init__() |
| - config = vllm_config.model_config.hf_config |
| + config = _mtp_hf_config(vllm_config) |
| self.mtp_start_layer_idx = config.num_hidden_layers |
| self.num_mtp_layers = config.num_nextn_predict_layers |
| self.layers = torch.nn.ModuleDict( |
| @@ -210,7 +223,7 @@ class Glm5NextMultiTokenPredictor(nn.Module): |
| class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts): |
| def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| super().__init__() |
| - self.config = vllm_config.model_config.hf_config |
| + self.config = _mtp_hf_config(vllm_config) |
| self.quant_config = vllm_config.quant_config |
| self.model = Glm5NextMultiTokenPredictor( |
| vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
|
|
