| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Adapted from |
| 5 | # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py |
| 6 | # Copyright 2025 Xiaomi Corporation. |
| 7 | # Copyright 2024 The Qwen team. |
| 8 | # Copyright 2023 The vLLM team. |
| 9 | # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. |
| 10 | # |
| 11 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 12 | # and OPT implementations in this library. It has been modified from its |
| 13 | # original forms to accommodate minor architectural differences compared |
| 14 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 15 | # |
| 16 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 17 | # you may not use this file except in compliance with the License. |
| 18 | # You may obtain a copy of the License at |
| 19 | # |
| 20 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 21 | # |
| 22 | # Unless required by applicable law or agreed to in writing, software |
| 23 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 24 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 25 | # See the License for the specific language governing permissions and |
| 26 | # limitations under the License. |
| 27 | """Inference-only MiMo model compatible with HuggingFace weights.""" |
| 28 | |
| 29 | from collections.abc import Iterable |
| 30 | from itertools import islice |
| 31 | |
| 32 | import torch |
| 33 | import torch.nn as nn |
| 34 | |
| 35 | from vllm.compilation.decorators import support_torch_compile |
| 36 | from vllm.config import VllmConfig |
| 37 | from vllm.distributed import get_pp_group |
| 38 | from vllm.logger import init_logger |
| 39 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 40 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 41 | from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2Model |
| 42 | from vllm.sequence import IntermediateTensors |
| 43 | |
| 44 | from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix |
| 45 | |
| 46 | logger = init_logger(__name__) |
| 47 | |
| 48 | |
| 49 | @support_torch_compile( |
| 50 | dynamic_arg_dims={ |
| 51 | "input_ids": 0, |
| 52 | "positions": -1, |
| 53 | "intermediate_tensors": 0, |
| 54 | "inputs_embeds": 0, |
| 55 | } |
| 56 | ) |
| 57 | class MiMoModel(Qwen2Model): |
| 58 | supports_aux_hidden_states_over_pp = False |
| 59 | |
| 60 | def forward( |
| 61 | self, |
| 62 | input_ids: torch.Tensor | None, |
| 63 | positions: torch.Tensor, |
| 64 | intermediate_tensors: IntermediateTensors | None = None, |
| 65 | inputs_embeds: torch.Tensor | None = None, |
| 66 | ) -> torch.Tensor | IntermediateTensors: |
| 67 | if get_pp_group().is_first_rank: |
| 68 | if inputs_embeds is not None: |
| 69 | hidden_states = inputs_embeds |
| 70 | else: |
| 71 | hidden_states = self.embed_input_ids(input_ids) |
| 72 | residual = None |
| 73 | else: |
| 74 | assert intermediate_tensors is not None |
| 75 | hidden_states = intermediate_tensors["hidden_states"] |
| 76 | residual = intermediate_tensors["residual"] |
| 77 | for layer in islice(self.layers, self.start_layer, self.end_layer): |
| 78 | hidden_states, residual = layer( |
| 79 | positions, |
| 80 | hidden_states, |
| 81 | residual, |
| 82 | ) |
| 83 | if not get_pp_group().is_last_rank: |
| 84 | return IntermediateTensors( |
| 85 | {"hidden_states": hidden_states, "residual": residual} |
| 86 | ) |
| 87 | hidden_states = hidden_states + residual |
| 88 | return hidden_states |
| 89 | |
| 90 | |
| 91 | class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module): |
| 92 | # MTP layers are loaded by the draft model, not the main model. |
| 93 | hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.mtp_layers.": None}) |
| 94 | |
| 95 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 96 | nn.Module.__init__(self) |
| 97 | config = vllm_config.model_config.hf_config |
| 98 | quant_config = vllm_config.quant_config |
| 99 | |
| 100 | self.config = config |
| 101 | |
| 102 | self.quant_config = quant_config |
| 103 | |
| 104 | self.model = MiMoModel( |
| 105 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 106 | ) |
| 107 | |
| 108 | if get_pp_group().is_last_rank: |
| 109 | self.lm_head = ParallelLMHead( |
| 110 | config.vocab_size, |
| 111 | config.hidden_size, |
| 112 | quant_config=quant_config, |
| 113 | prefix=maybe_prefix(prefix, "lm_head"), |
| 114 | ) |
| 115 | if config.tie_word_embeddings: |
| 116 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 117 | else: |
| 118 | self.lm_head = PPMissingLayer() |
| 119 | |
| 120 | self.logits_processor = LogitsProcessor(config.vocab_size) |
| 121 | |
| 122 | self.make_empty_intermediate_tensors = ( |
| 123 | self.model.make_empty_intermediate_tensors |
| 124 | ) |
| 125 | |
| 126 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 127 | loader = AutoWeightsLoader(self) |
| 128 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 129 | |
| 130 | def compute_logits( |
| 131 | self, |
| 132 | hidden_states: torch.Tensor, |
| 133 | ) -> torch.Tensor | None: |
| 134 | hidden_states = self.model.norm(hidden_states) |
| 135 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 136 | return logits |