@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 829bbd7882
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[New Model]mBART model (#22883)

Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>

汪志鹏 committed Aug 16, 2025 at 20:16 UTC 829bbd7882222c85c0ca5a17fbb2f70e543f50ac
6 files changed +716 -91
docs/models/supported_models.md
+4
@@ -330,6 +330,7 @@ th {
330 | `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ | ✅︎ |
331 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | |
332 | `BartForConditionalGeneration` | BART | `facebook/bart-base`, `facebook/bart-large-cnn`, etc. | | | |
333 +| `MBartForConditionalGeneration` | mBART | `facebook/mbart-large-en-ro`, `facebook/mbart-large-50`, etc. | | | |
334 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ | ✅︎ |
335 | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, etc. | ✅︎ | ✅︎ | ✅︎ |
336 | `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ | ✅︎ |
@@ -418,6 +419,9 @@ Some models are supported only via the [Transformers backend](#transformers). Th
419 !!! note
420 Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
421
422 +!!! note
423 + Some mBART models' config files do not have an `architecture` defined. Therefore, you need to use `--hf-overrides '{"architectures": ["MBartForConditionalGeneration"]}'` to explicitly specify the use of the `MBartForConditionalGeneration` architecture.
424 +
425 ### Pooling Models
426
427 See [this page](./pooling_models.md) for more information on how to use pooling models.
examples/offline_inference/encoder_decoder.py
+147 -86
@@ -2,9 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """
4 Demonstrate prompting of text-to-text
5 -encoder/decoder models, specifically BART
5 +encoder/decoder models, specifically BART and mBART.
6 +
7 +This script is refactored to allow model selection via command-line arguments.
8 """
9
10 +import argparse
11 +from typing import NamedTuple, Optional
12 +
13 from vllm import LLM, SamplingParams
14 from vllm.inputs import (
15 ExplicitEncoderDecoderPrompt,
@@ -14,119 +19,175 @@ from vllm.inputs import (
19 )
20
21
17 -def create_prompts(tokenizer):
18 - # Test prompts
19 - #
20 - # This section shows all of the valid ways to prompt an
21 - # encoder/decoder model.
22 - #
23 - # - Helpers for building prompts
24 - text_prompt_raw = "Hello, my name is"
25 - text_prompt = TextPrompt(prompt="The president of the United States is")
26 - tokens_prompt = TokensPrompt(
27 - prompt_token_ids=tokenizer.encode(prompt="The capital of France is")
22 +class ModelRequestData(NamedTuple):
23 + """
24 + Holds the configuration for a specific model, including its
25 + HuggingFace ID and the prompts to use for the demo.
26 + """
27 +
28 + model_id: str
29 + encoder_prompts: list
30 + decoder_prompts: list
31 + hf_overrides: Optional[dict] = None
32 +
33 +
34 +def get_bart_config() -> ModelRequestData:
35 + """
36 + Returns the configuration for facebook/bart-large-cnn.
37 + This uses the exact test cases from the original script.
38 + """
39 + encoder_prompts = [
40 + "Hello, my name is",
41 + "The president of the United States is",
42 + "The capital of France is",
43 + "An encoder prompt",
44 + ]
45 + decoder_prompts = [
46 + "A decoder prompt",
47 + "Another decoder prompt",
48 + ]
49 + return ModelRequestData(
50 + model_id="facebook/bart-large-cnn",
51 + encoder_prompts=encoder_prompts,
52 + decoder_prompts=decoder_prompts,
53 )
29 - # - Pass a single prompt to encoder/decoder model
30 - # (implicitly encoder input prompt);
31 - # decoder input prompt is assumed to be None
32 -
33 - single_text_prompt_raw = text_prompt_raw # Pass a string directly
34 - single_text_prompt = text_prompt # Pass a TextPrompt
35 - single_tokens_prompt = tokens_prompt # Pass a TokensPrompt
36 -
37 - # ruff: noqa: E501
38 - # - Pass explicit encoder and decoder input prompts within one data structure.
39 - # Encoder and decoder prompts can both independently be text or tokens, with
40 - # no requirement that they be the same prompt type. Some example prompt-type
41 - # combinations are shown below, note that these are not exhaustive.
42 -
43 - enc_dec_prompt1 = ExplicitEncoderDecoderPrompt(
44 - # Pass encoder prompt string directly, &
45 - # pass decoder prompt tokens
46 - encoder_prompt=single_text_prompt_raw,
47 - decoder_prompt=single_tokens_prompt,
48 - )
49 - enc_dec_prompt2 = ExplicitEncoderDecoderPrompt(
50 - # Pass TextPrompt to encoder, and
51 - # pass decoder prompt string directly
52 - encoder_prompt=single_text_prompt,
53 - decoder_prompt=single_text_prompt_raw,
54 +
55 +
56 +def get_mbart_config() -> ModelRequestData:
57 + """
58 + Returns the configuration for facebook/mbart-large-en-ro.
59 + This uses prompts suitable for an English-to-Romanian translation task.
60 + """
61 + encoder_prompts = [
62 + "The quick brown fox jumps over the lazy dog.",
63 + "How are you today?",
64 + ]
65 + decoder_prompts = ["", ""]
66 + hf_overrides = {"architectures": ["MBartForConditionalGeneration"]}
67 + return ModelRequestData(
68 + model_id="facebook/mbart-large-en-ro",
69 + encoder_prompts=encoder_prompts,
70 + decoder_prompts=decoder_prompts,
71 + hf_overrides=hf_overrides,
72 )
55 - enc_dec_prompt3 = ExplicitEncoderDecoderPrompt(
56 - # Pass encoder prompt tokens directly, and
57 - # pass TextPrompt to decoder
58 - encoder_prompt=single_tokens_prompt,
59 - decoder_prompt=single_text_prompt,
73 +
74 +
75 +MODEL_GETTERS = {
76 + "bart": get_bart_config,
77 + "mbart": get_mbart_config,
78 +}
79 +
80 +
81 +def create_all_prompt_types(
82 + encoder_prompts_raw: list,
83 + decoder_prompts_raw: list,
84 + tokenizer,
85 +) -> list:
86 + """
87 + Generates a list of diverse prompt types for demonstration.
88 + This function is generic and uses the provided raw prompts
89 + to create various vLLM input objects.
90 + """
91 + text_prompt_raw = encoder_prompts_raw[0]
92 + text_prompt = TextPrompt(prompt=encoder_prompts_raw[1 % len(encoder_prompts_raw)])
93 + tokens_prompt = TokensPrompt(
94 + prompt_token_ids=tokenizer.encode(
95 + encoder_prompts_raw[2 % len(encoder_prompts_raw)]
96 + )
97 )
98
62 - # - Finally, here's a useful helper function for zipping encoder and
63 - # decoder prompts together into a list of ExplicitEncoderDecoderPrompt
64 - # instances
99 + decoder_tokens_prompt = TokensPrompt(
100 + prompt_token_ids=tokenizer.encode(decoder_prompts_raw[0])
101 + )
102 + single_prompt_examples = [
103 + text_prompt_raw,
104 + text_prompt,
105 + tokens_prompt,
106 + ]
107 + explicit_pair_examples = [
108 + ExplicitEncoderDecoderPrompt(
109 + encoder_prompt=text_prompt_raw,
110 + decoder_prompt=decoder_tokens_prompt,
111 + ),
112 + ExplicitEncoderDecoderPrompt(
113 + encoder_prompt=text_prompt,
114 + decoder_prompt=decoder_prompts_raw[1 % len(decoder_prompts_raw)],
115 + ),
116 + ExplicitEncoderDecoderPrompt(
117 + encoder_prompt=tokens_prompt,
118 + decoder_prompt=text_prompt,
119 + ),
120 + ]
121 zipped_prompt_list = zip_enc_dec_prompts(
66 - ["An encoder prompt", "Another encoder prompt"],
67 - ["A decoder prompt", "Another decoder prompt"],
122 + encoder_prompts_raw,
123 + decoder_prompts_raw,
124 )
125 + return single_prompt_examples + explicit_pair_examples + zipped_prompt_list
126
70 - # - Let's put all of the above example prompts together into one list
71 - # which we will pass to the encoder/decoder LLM.
72 - return [
73 - single_text_prompt_raw,
74 - single_text_prompt,
75 - single_tokens_prompt,
76 - enc_dec_prompt1,
77 - enc_dec_prompt2,
78 - enc_dec_prompt3,
79 - ] + zipped_prompt_list
127
81 -
82 -# Create a sampling params object.
83 -def create_sampling_params():
128 +def create_sampling_params() -> SamplingParams:
129 + """Create a sampling params object."""
130 return SamplingParams(
131 temperature=0,
132 top_p=1.0,
133 min_tokens=0,
88 - max_tokens=20,
134 + max_tokens=30,
135 )
136
137
92 -# Print the outputs.
93 -def print_outputs(outputs):
94 - print("-" * 50)
138 +def print_outputs(outputs: list):
139 + """Formats and prints the generation outputs."""
140 + print("-" * 80)
141 for i, output in enumerate(outputs):
142 prompt = output.prompt
143 encoder_prompt = output.encoder_prompt
144 generated_text = output.outputs[0].text
145 print(f"Output {i + 1}:")
100 - print(
101 - f"Encoder prompt: {encoder_prompt!r}\n"
102 - f"Decoder prompt: {prompt!r}\n"
103 - f"Generated text: {generated_text!r}"
146 + print(f"Encoder Prompt: {encoder_prompt!r}")
147 + print(f"Decoder Prompt: {prompt!r}")
148 + print(f"Generated Text: {generated_text!r}")
149 + print("-" * 80)
150 +
151 +
152 +def main(args):
153 + """Main execution function."""
154 + model_key = args.model
155 + if model_key not in MODEL_GETTERS:
156 + raise ValueError(
157 + f"Unknown model: {model_key}. "
158 + f"Available models: {list(MODEL_GETTERS.keys())}"
159 )
105 - print("-" * 50)
106 -
107 -
108 -def main():
109 - dtype = "float"
160 + config_getter = MODEL_GETTERS[model_key]
161 + model_config = config_getter()
162
111 - # Create a BART encoder/decoder model instance
163 + print(f"🚀 Running demo for model: {model_config.model_id}")
164 llm = LLM(
113 - model="facebook/bart-large-cnn",
114 - dtype=dtype,
165 + model=model_config.model_id,
166 + dtype="float",
167 + hf_overrides=model_config.hf_overrides,
168 )
116 -
117 - # Get BART tokenizer
169 tokenizer = llm.llm_engine.get_tokenizer_group()
119 -
120 - prompts = create_prompts(tokenizer)
170 + prompts = create_all_prompt_types(
171 + encoder_prompts_raw=model_config.encoder_prompts,
172 + decoder_prompts_raw=model_config.decoder_prompts,
173 + tokenizer=tokenizer,
174 + )
175 sampling_params = create_sampling_params()
122 -
123 - # Generate output tokens from the prompts. The output is a list of
124 - # RequestOutput objects that contain the prompt, generated
125 - # text, and other information.
176 outputs = llm.generate(prompts, sampling_params)
127 -
177 print_outputs(outputs)
178
179
180 if __name__ == "__main__":
132 - main()
181 + parser = argparse.ArgumentParser(
182 + description="A flexible demo for vLLM encoder-decoder models."
183 + )
184 + parser.add_argument(
185 + "--model",
186 + "-m",
187 + type=str,
188 + default="bart",
189 + choices=MODEL_GETTERS.keys(),
190 + help="The short name of the model to run.",
191 + )
192 + args = parser.parse_args()
193 + main(args)
tests/models/language/generation/test_mbart.py new
+123
@@ -0,0 +1,123 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Optional
4 +
5 +import pytest
6 +from transformers import AutoModelForSeq2SeqLM
7 +
8 +from vllm.sequence import SampleLogprobs
9 +
10 +from ....conftest import DecoderPromptType, HfRunner, VllmRunner
11 +from ...utils import check_logprobs_close
12 +
13 +
14 +def vllm_to_hf_output(
15 + vllm_output: tuple[list[int], str, Optional[SampleLogprobs]],
16 + decoder_prompt_type: DecoderPromptType,
17 +):
18 + """Sanitize vllm output to be comparable with hf output."""
19 + output_ids, output_str, out_logprobs = vllm_output
20 + hf_output_str = output_str + "</s>"
21 + return output_ids, hf_output_str, out_logprobs
22 +
23 +
24 +def run_test(
25 + hf_runner: type[HfRunner],
26 + vllm_runner: type[VllmRunner],
27 + prompts: list[dict[str, str]],
28 + decoder_prompt_type: DecoderPromptType,
29 + model: str,
30 + *,
31 + dtype: str,
32 + max_tokens: int,
33 + num_logprobs: int,
34 + tensor_parallel_size: int,
35 + distributed_executor_backend: Optional[str] = None,
36 +) -> None:
37 + '''
38 + Test the vLLM mBART model by validating it against HuggingFace (HF).
39 + (Docstring content is omitted for brevity)
40 + '''
41 +
42 + vllm_prompts = prompts
43 + if decoder_prompt_type == DecoderPromptType.NONE:
44 + vllm_prompts = [{
45 + "encoder_prompt": p['encoder_prompt'],
46 + "decoder_prompt": ""
47 + } for p in prompts]
48 +
49 + vllm_kwargs = {
50 + "hf_overrides": {
51 + "architectures": ["MBartForConditionalGeneration"]
52 + }
53 + }
54 +
55 + with vllm_runner(model,
56 + dtype=dtype,
57 + tensor_parallel_size=tensor_parallel_size,
58 + distributed_executor_backend=distributed_executor_backend,
59 + enforce_eager=True,
60 + **vllm_kwargs) as vllm_model: # type: ignore
61 + vllm_outputs = vllm_model.generate_encoder_decoder_greedy_logprobs(
62 + vllm_prompts, max_tokens, num_logprobs)
63 +
64 + hf_kwargs = {
65 + "top_k": None,
66 + "num_beams": 1,
67 + "repetition_penalty": 1.0,
68 + "top_p": 1.0,
69 + "length_penalty": 1.0,
70 + "early_stopping": False,
71 + "no_repeat_ngram_size": None,
72 + "min_length": 0
73 + }
74 +
75 + with hf_runner(model, dtype=dtype,
76 + auto_cls=AutoModelForSeq2SeqLM) as hf_model:
77 + hf_kwargs["decoder_start_token_id"] = (
78 + hf_model.tokenizer.lang_code_to_id["ro_RO"])
79 +
80 + hf_outputs = (
81 + hf_model.generate_encoder_decoder_greedy_logprobs_limit(
82 + prompts, # HF runner still uses the original prompts
83 + max_tokens,
84 + num_logprobs,
85 + **hf_kwargs,
86 + ))
87 +
88 + hf_skip_tokens = 0
89 +
90 + check_logprobs_close(
91 + outputs_0_lst=hf_outputs,
92 + outputs_1_lst=[
93 + vllm_to_hf_output(vllm_output, decoder_prompt_type)
94 + for vllm_output in vllm_outputs
95 + ],
96 + name_0="hf",
97 + name_1="vllm",
98 + num_outputs_0_skip_tokens=hf_skip_tokens,
99 + )
100 +
101 +
102 +@pytest.mark.parametrize(
103 + "model",
104 + [pytest.param("facebook/mbart-large-en-ro")],
105 +)
106 +@pytest.mark.parametrize("dtype", ["float", "bfloat16"])
107 +@pytest.mark.parametrize("max_tokens", [64])
108 +@pytest.mark.parametrize("num_logprobs", [5])
109 +@pytest.mark.parametrize("decoder_prompt_type", list(DecoderPromptType))
110 +def test_models(hf_runner, vllm_runner, example_encoder_decoder_prompts, model,
111 + dtype, max_tokens, num_logprobs, decoder_prompt_type) -> None:
112 +
113 + run_test(
114 + hf_runner,
115 + vllm_runner,
116 + example_encoder_decoder_prompts[decoder_prompt_type],
117 + decoder_prompt_type,
118 + model,
119 + dtype=dtype,
120 + max_tokens=max_tokens,
121 + num_logprobs=num_logprobs,
122 + tensor_parallel_size=1,
123 + )
tests/models/registry.py
+2
@@ -316,6 +316,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
316 # [Encoder-decoder]
317 "BartModel": _HfExamplesInfo("facebook/bart-base"),
318 "BartForConditionalGeneration": _HfExamplesInfo("facebook/bart-large-cnn"),
319 + "MBartForConditionalGeneration": _HfExamplesInfo("facebook/mbart-large-en-ro", # noqa: E501
320 + hf_overrides={"architectures": ["MBartForConditionalGeneration"]}), # noqa: E501
321 }
322
323 _EMBEDDING_EXAMPLE_MODELS = {
vllm/model_executor/models/bart.py
+439 -5
@@ -46,7 +46,8 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
46 from vllm.sequence import IntermediateTensors
47
48 from .interfaces import SupportsQuant, SupportsV0Only
49 -from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
49 +from .utils import (AutoWeightsLoader, WeightsMapper, cast_overflow_tensors,
50 + maybe_prefix)
51
52 logger = logging.get_logger(__name__)
53
@@ -422,10 +423,7 @@ class BartEncoderLayer(nn.Module):
423 if hidden_states.dtype == torch.float16 and (
424 torch.isinf(hidden_states).any()
425 or torch.isnan(hidden_states).any()):
425 - clamp_value = torch.finfo(hidden_states.dtype).max - 1000
426 - hidden_states = torch.clamp(hidden_states,
427 - min=-clamp_value,
428 - max=clamp_value)
426 + hidden_states = cast_overflow_tensors(hidden_states)
427
428 return hidden_states
429
@@ -906,3 +904,439 @@ class BartForConditionalGeneration(nn.Module, SupportsV0Only, SupportsQuant):
904 })
905
906 return loaded_params
907 +
908 +
909 +class MBartEncoderLayer(BartEncoderLayer):
910 +
911 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
912 + r"""
913 + Args:
914 + hidden_states
915 + torch.Tensor of *encoder* input embeddings.
916 + Returns:
917 + Encoder layer output torch.Tensor
918 + """
919 + residual = hidden_states
920 + hidden_states = self.self_attn_layer_norm(hidden_states)
921 + hidden_states = self.self_attn(hidden_states=hidden_states)
922 +
923 + hidden_states = residual + hidden_states
924 +
925 + residual = hidden_states
926 + hidden_states = self.final_layer_norm(hidden_states)
927 + fc1_out, _ = self.fc1(hidden_states)
928 + hidden_states = self.activation_fn(fc1_out)
929 +
930 + hidden_states, _ = self.fc2(hidden_states)
931 +
932 + hidden_states = residual + hidden_states
933 +
934 + if hidden_states.dtype == torch.float16 and (
935 + torch.isinf(hidden_states).any()
936 + or torch.isnan(hidden_states).any()):
937 + hidden_states = cast_overflow_tensors(hidden_states)
938 +
939 + return hidden_states
940 +
941 +
942 +class MBartDecoderLayer(BartDecoderLayer):
943 +
944 + def forward(
945 + self,
946 + decoder_hidden_states: torch.Tensor,
947 + encoder_hidden_states: Optional[torch.Tensor] = None,
948 + ) -> torch.Tensor:
949 + residual = decoder_hidden_states
950 + hidden_states = self.self_attn_layer_norm(decoder_hidden_states)
951 +
952 + # Self Attention
953 + hidden_states = self.self_attn(hidden_states=hidden_states)
954 +
955 + hidden_states = residual + hidden_states
956 +
957 + # Cross-Attention Block
958 +
959 + residual = hidden_states
960 + hidden_states = self.encoder_attn_layer_norm(hidden_states)
961 +
962 + hidden_states = self.encoder_attn(
963 + decoder_hidden_states=hidden_states,
964 + encoder_hidden_states=encoder_hidden_states,
965 + )
966 +
967 + hidden_states = residual + hidden_states
968 +
969 + # Fully Connected
970 + residual = hidden_states
971 + hidden_states = self.final_layer_norm(hidden_states)
972 + fc1_out, _ = self.fc1(hidden_states)
973 + hidden_states = self.activation_fn(fc1_out)
974 +
975 + hidden_states, _ = self.fc2(hidden_states)
976 +
977 + hidden_states = residual + hidden_states
978 +
979 + return hidden_states
980 +
981 +
982 +class MBartEncoder(nn.Module):
983 + """
984 + Transformer encoder consisting of *config.encoder_layers*
985 + self attention layers. Each layer is a [`BartEncoderLayer`].
986 + Args:
987 + config: BartConfig
988 + embed_tokens (nn.Embedding): output embedding
989 + """
990 +
991 + def __init__(self,
992 + config: BartConfig,
993 + cache_config: Optional[CacheConfig] = None,
994 + quant_config: Optional[QuantizationConfig] = None,
995 + lora_config: Optional[LoRAConfig] = None,
996 + embed_tokens: Optional[nn.Embedding] = None,
997 + prefix: str = ""):
998 + super().__init__()
999 +
1000 + self.cache_config = cache_config
1001 + self.quant_config = quant_config
1002 + self.lora_config = lora_config
1003 + embed_dim = config.d_model
1004 + self.max_source_positions = config.max_position_embeddings
1005 + embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
1006 +
1007 + self.embed_tokens = BartScaledWordEmbedding(config.vocab_size,
1008 + embed_dim,
1009 + embed_scale=embed_scale)
1010 +
1011 + if embed_tokens is not None:
1012 + self.embed_tokens.weight = embed_tokens.weight
1013 +
1014 + self.embed_positions = BartLearnedPositionalEmbedding(
1015 + config.max_position_embeddings,
1016 + embed_dim,
1017 + )
1018 + self.layers = nn.ModuleList([
1019 + MBartEncoderLayer(config,
1020 + cache_config,
1021 + quant_config,
1022 + prefix=f"{prefix}.layers.{layer_idx}")
1023 + for layer_idx in range(config.encoder_layers)
1024 + ])
1025 +
1026 + self.layernorm_embedding = nn.LayerNorm(embed_dim)
1027 + self.layer_norm = nn.LayerNorm(config.d_model) # 改动
1028 +
1029 + def forward(
1030 + self,
1031 + input_ids: torch.Tensor,
1032 + positions: torch.Tensor,
1033 + inputs_embeds: Optional[torch.Tensor] = None,
1034 + ) -> torch.Tensor:
1035 + r"""
1036 + Args:
1037 + input_ids
1038 + Indices of *encoder* input sequence tokens in the vocabulary.
1039 + Padding will be ignored by default should you
1040 + provide it.
1041 + positions
1042 + Positions of *encoder* input sequence tokens.
1043 + Returns:
1044 + Decoder output torch.Tensor
1045 + """
1046 + # retrieve input_ids and inputs_embeds
1047 + if inputs_embeds is None:
1048 + inputs_embeds = self.embed_tokens(input_ids)
1049 +
1050 + embed_pos = self.embed_positions(positions)
1051 + embed_pos = embed_pos.to(inputs_embeds.device)
1052 +
1053 + hidden_states = inputs_embeds + embed_pos
1054 + hidden_states = self.layernorm_embedding(hidden_states)
1055 +
1056 + for encoder_layer in self.layers:
1057 + hidden_states = encoder_layer(hidden_states=hidden_states)
1058 +
1059 + hidden_states = self.layer_norm(hidden_states)
1060 + return hidden_states
1061 +
1062 +
1063 +class MBartDecoder(nn.Module):
1064 + """
1065 + Transformer decoder consisting of *config.decoder_layers* layers.
1066 + Each layer is a [`BartDecoderLayer`]
1067 + Args:
1068 + config: BartConfig
1069 + embed_tokens (nn.Embedding): output embedding
1070 + """
1071 +
1072 + def __init__(
1073 + self,
1074 + config: BartConfig,
1075 + cache_config: Optional[CacheConfig] = None,
1076 + quant_config: Optional[QuantizationConfig] = None,
1077 + lora_config: Optional[LoRAConfig] = None,
1078 + embed_tokens: Optional[nn.Embedding] = None,
1079 + prefix: str = "",
1080 + ):
1081 + super().__init__()
1082 + self.cache_config = cache_config
1083 + self.quant_config = quant_config
1084 + self.lora_config = lora_config
1085 + self.max_target_positions = config.max_position_embeddings
1086 + embed_scale = math.sqrt(
1087 + config.d_model) if config.scale_embedding else 1.0
1088 +
1089 + self.embed_tokens = BartScaledWordEmbedding(config.vocab_size,
1090 + config.d_model,
1091 + embed_scale=embed_scale)
1092 +
1093 + if embed_tokens is not None:
1094 + self.embed_tokens.weight = embed_tokens.weight
1095 +
1096 + self.embed_positions = BartLearnedPositionalEmbedding(
1097 + config.max_position_embeddings,
1098 + config.d_model,
1099 + )
1100 +
1101 + self.layers = nn.ModuleList(
1102 + [MBartDecoderLayer(config, cache_config, quant_config,
1103 + prefix=f"{prefix}.layers.{layer_idx}") \
1104 + for layer_idx in range(config.decoder_layers)])
1105 +
1106 + self.layernorm_embedding = nn.LayerNorm(config.d_model)
1107 + self.layer_norm = nn.LayerNorm(config.d_model)
1108 +
1109 + def forward(
1110 + self,
1111 + decoder_input_ids: torch.Tensor,
1112 + decoder_positions: torch.Tensor,
1113 + encoder_hidden_states: Optional[torch.Tensor],
1114 + inputs_embeds: Optional[torch.Tensor] = None,
1115 + ) -> torch.Tensor:
1116 + r"""
1117 + Args:
1118 + decoder_input_ids
1119 + Indices of *decoder* input sequence tokens in the vocabulary.
1120 + Padding will be ignored by default should you
1121 + provide it.
1122 + decoder_positions
1123 + Positions of *decoder* input sequence tokens.
1124 + encoder_hidden_states:
1125 + Tensor of encoder output embeddings
1126 + Returns:
1127 + Decoder output torch.Tensor
1128 + """
1129 + if inputs_embeds is None:
1130 + inputs_embeds = self.embed_tokens(decoder_input_ids)
1131 + else:
1132 + decoder_positions = inputs_embeds[:, -1]
1133 +
1134 + # embed positions
1135 + embed_pos = self.embed_positions(decoder_positions)
1136 + embed_pos = embed_pos.to(inputs_embeds.device)
1137 +
1138 + hidden_states = inputs_embeds + embed_pos
1139 + hidden_states = self.layernorm_embedding(hidden_states)
1140 +
1141 + # decoder layers
1142 +
1143 + for decoder_layer in self.layers:
1144 + hidden_states = decoder_layer(
1145 + decoder_hidden_states=hidden_states,
1146 + encoder_hidden_states=encoder_hidden_states,
1147 + )
1148 +
1149 + hidden_states = self.layer_norm(hidden_states)
1150 + return hidden_states
1151 +
1152 +
1153 +class MBartModel(nn.Module, SupportsQuant):
1154 + _tied_weights_keys = [
1155 + "encoder.embed_tokens.weight", "decoder.embed_tokens.weight"
1156 + ]
1157 +
1158 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1159 + super().__init__()
1160 +
1161 + config = vllm_config.model_config.hf_config
1162 + cache_config = vllm_config.cache_config
1163 + quant_config = vllm_config.quant_config
1164 + lora_config = vllm_config.lora_config
1165 +
1166 + self.config = config
1167 +
1168 + lora_vocab = (lora_config.lora_extra_vocab_size *
1169 + (lora_config.max_loras or 1)) if lora_config else 0
1170 + self.vocab_size = config.vocab_size + lora_vocab
1171 + self.org_vocab_size = config.vocab_size
1172 +
1173 + self.encoder = MBartEncoder(config,
1174 + cache_config,
1175 + quant_config=quant_config,
1176 + prefix=f"{prefix}.encoder")
1177 + self.decoder = MBartDecoder(config,
1178 + cache_config,
1179 + quant_config=quant_config,
1180 + prefix=f"{prefix}.decoder")
1181 +
1182 + def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,
1183 + encoder_input_ids: torch.Tensor,
1184 + encoder_positions: torch.Tensor) -> torch.Tensor:
1185 + r"""
1186 + Args:
1187 + input_ids
1188 + Indices of *decoder* input sequence tokens in the vocabulary.
1189 + Padding will be ignored by default should you
1190 + provide it.
1191 + positions
1192 + Positions of *decoder* input sequence tokens.
1193 + encoder_input_ids
1194 + Indices of *encoder* input sequence tokens in the vocabulary.
1195 + encoder_positions:
1196 + Positions of *encoder* input sequence tokens.
1197 + Returns:
1198 + Model output torch.Tensor
1199 + """
1200 +
1201 + encoder_hidden_states = None
1202 +
1203 + if encoder_input_ids.numel() > 0:
1204 + # Run encoder attention if a non-zero number of encoder tokens
1205 + # are provided as input
1206 + encoder_hidden_states = self.encoder(input_ids=encoder_input_ids,
1207 + positions=encoder_positions)
1208 +
1209 + # decoder outputs consists of
1210 + # (dec_features, past_key_value, dec_hidden, dec_attn)
1211 + decoder_outputs = self.decoder(
1212 + decoder_input_ids=input_ids,
1213 + decoder_positions=positions,
1214 + encoder_hidden_states=encoder_hidden_states)
1215 +
1216 + return decoder_outputs
1217 +
1218 +
1219 +class MBartForConditionalGeneration(nn.Module, SupportsV0Only, SupportsQuant):
1220 + base_model_prefix = "model"
1221 +
1222 + hf_to_vllm_mapper = WeightsMapper(
1223 + orig_to_new_prefix={
1224 + "decoder.": "model.decoder.",
1225 + "encoder.": "model.encoder.",
1226 + "shared.": "model.shared."
1227 + },
1228 + orig_to_new_substr={
1229 + "beta": "bias",
1230 + "gamma": "weight",
1231 + "LayerNorm": "layernorm",
1232 + },
1233 + )
1234 +
1235 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1236 + super().__init__()
1237 + config = vllm_config.model_config.hf_config
1238 + lora_config = vllm_config.lora_config
1239 + assert config.tie_word_embeddings
1240 + self.config = config
1241 + self.model = MBartModel(vllm_config=vllm_config,
1242 + prefix=maybe_prefix(prefix, "model"))
1243 +
1244 + self.unpadded_vocab_size = config.vocab_size
1245 + if lora_config:
1246 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
1247 +
1248 + embed_scale = math.sqrt(
1249 + config.d_model) if config.scale_embedding else 1.0
1250 +
1251 + self.lm_head = BartParallelLMHead(config.vocab_size,
1252 + config.d_model,
1253 + embed_scale=embed_scale)
1254 +
1255 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
1256 + config.vocab_size)
1257 +
1258 + def forward(
1259 + self,
1260 + input_ids: torch.Tensor,
1261 + positions: torch.Tensor,
1262 + intermediate_tensors: Optional[IntermediateTensors] = None,
1263 + *,
1264 + encoder_input_ids: torch.Tensor,
1265 + encoder_positions: torch.Tensor,
1266 + **kwargs,
1267 + ) -> torch.Tensor:
1268 + return self.model(input_ids, positions, encoder_input_ids,
1269 + encoder_positions)
1270 +
1271 + def compute_logits(
1272 + self,
1273 + hidden_states: torch.Tensor,
1274 + sampling_metadata: SamplingMetadata,
1275 + ) -> Optional[torch.Tensor]:
1276 + logits = self.logits_processor(self.lm_head, hidden_states,
1277 + sampling_metadata)
1278 + return logits
1279 +
1280 + def load_weights(self, weights: Iterable[tuple[str,
1281 + torch.Tensor]]) -> set[str]:
1282 + stacked_params_mapping = [
1283 + ("qkv_proj", "q_proj", "q"),
1284 + ("qkv_proj", "k_proj", "k"),
1285 + ("qkv_proj", "v_proj", "v"),
1286 + ]
1287 + model_params_dict = dict(self.named_parameters())
1288 + loaded_params = set()
1289 + remaining_weights = []
1290 + shared_embedding_weight = None
1291 +
1292 + for name, loaded_weight in weights:
1293 + if any(skip in name
1294 + for skip in ["cls.", "pooler.", "final_logits_bias"]):
1295 + continue
1296 + if any(embed_name in name for embed_name in [
1297 + 'shared.weight', 'encoder.embed_tokens.weight',
1298 + 'decoder.embed_tokens.weight'
1299 + ]):
1300 + if shared_embedding_weight is None:
1301 + shared_embedding_weight = loaded_weight
1302 + continue
1303 + is_stacked = False
1304 + for param_name, weight_name, shard_id in stacked_params_mapping:
1305 + if weight_name not in name:
1306 + continue
1307 + vllm_name = name
1308 + for src, dst in self.hf_to_vllm_mapper.orig_to_new_substr.items(
1309 + ):
1310 + vllm_name = vllm_name.replace(src, dst)
1311 + for src, dst in self.hf_to_vllm_mapper.orig_to_new_prefix.items(
1312 + ):
1313 + if vllm_name.startswith(src):
1314 + vllm_name = dst + vllm_name[len(src):]
1315 + break
1316 + vllm_name = vllm_name.replace(weight_name, param_name)
1317 + if vllm_name in model_params_dict:
1318 + param = model_params_dict[vllm_name]
1319 + weight_loader = getattr(param, "weight_loader",
1320 + default_weight_loader)
1321 + weight_loader(param, loaded_weight, shard_id)
1322 + loaded_params.add(vllm_name)
1323 + is_stacked = True
1324 + break
1325 + if not is_stacked:
1326 + remaining_weights.append((name, loaded_weight))
1327 + loader = AutoWeightsLoader(self, skip_prefixes=["cls.", "pooler."])
1328 + auto_loaded_params = loader.load_weights(remaining_weights,
1329 + mapper=self.hf_to_vllm_mapper)
1330 + loaded_params.update(auto_loaded_params)
1331 + if shared_embedding_weight is not None:
1332 + lm_head_param = self.lm_head.weight
1333 + weight_loader = getattr(lm_head_param, "weight_loader",
1334 + default_weight_loader)
1335 + weight_loader(lm_head_param, shared_embedding_weight)
1336 + self.model.encoder.embed_tokens.weight = self.lm_head.weight
1337 + self.model.decoder.embed_tokens.weight = self.lm_head.weight
1338 + loaded_params.update({
1339 + 'model.encoder.embed_tokens.weight', 'lm_head.weight',
1340 + 'model.decoder.embed_tokens.weight'
1341 + })
1342 + return loaded_params
vllm/model_executor/models/registry.py
+1
@@ -141,6 +141,7 @@ _TEXT_GENERATION_MODELS = {
141 # [Encoder-decoder]
142 "BartModel": ("bart", "BartForConditionalGeneration"),
143 "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
144 + "MBartForConditionalGeneration": ("bart", "MBartForConditionalGeneration"),
145 }
146
147 _EMBEDDING_MODELS = {