@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ee21291825
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 10, 17:45 UTC

[Model] Nemotron Parse 1.1 Support (#30864)

Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

amitz-nv committed Jan 5, 2026 at 23:00 UTC ee212918250ab5aa33a2968c027d71ffd6048ae4
13 files changed +1117 -31
requirements/test.in
+2 -1
@@ -9,6 +9,7 @@ pytest-timeout
9 pytest-cov
10
11 # testing utils
12 +albumentations # required for Nemotron Parse in test_common.py
13 backoff # required for phi4mm test
14 blobfile # required for kimi-vl test
15 einops # required for MPT, qwen-vl
@@ -31,7 +32,7 @@ transformers_stream_generator # required for qwen-vl test
32 matplotlib # required for qwen-vl test
33 mistral_common[image,audio] >= 1.8.8 # required for voxtral test
34 num2words # required for smolvlm test
34 -open_clip_torch==2.32.0 # Required for nemotron_vl test
35 +open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
36 opencv-python-headless >= 4.11.0 # required for video test
37 datamodel_code_generator # required for minicpm3 test
38 lm-eval[api]>=0.4.9.2 # required for model evaluation test
requirements/test.txt
+3 -1
@@ -27,7 +27,9 @@ aiosignal==1.4.0
27 albucore==0.0.16
28 # via terratorch
29 albumentations==1.4.6
30 - # via terratorch
30 + # via
31 + # -r requirements/test.in
32 + # terratorch
33 alembic==1.16.4
34 # via mlflow
35 annotated-types==0.7.0
tests/conftest.py
+2 -1
@@ -685,6 +685,7 @@ class HfRunner:
685 images: PromptImageInput | None = None,
686 audios: PromptAudioInput | None = None,
687 videos: PromptVideoInput | None = None,
688 + use_cache: bool = True,
689 **kwargs: Any,
690 ) -> list[TokensTextLogprobs]:
691 all_inputs = self.get_inputs(
@@ -698,7 +699,7 @@ class HfRunner:
699 for inputs in all_inputs:
700 output: "GenerateOutput" = self.model.generate(
701 **self.wrap_device(inputs),
701 - use_cache=True,
702 + use_cache=use_cache,
703 do_sample=False,
704 max_new_tokens=max_tokens,
705 output_hidden_states=True,
tests/models/multimodal/generation/test_nemotron_parse.py new
+89
@@ -0,0 +1,89 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Sequence
5 +
6 +import pytest
7 +from transformers import AutoModel
8 +
9 +from tests.models.utils import check_logprobs_close
10 +from vllm.assets.image import ImageAsset
11 +
12 +from ....conftest import HfRunner, PromptImageInput, VllmRunner
13 +from ....utils import create_new_process_for_each_test
14 +
15 +IMAGE = ImageAsset("paper-11").pil_image_ext(ext="png").convert("RGB")
16 +PROMPT = "</s><s><predict_bbox><predict_classes><output_markdown>"
17 +
18 +
19 +def run_test(
20 + hf_runner: type[HfRunner],
21 + vllm_runner: type[VllmRunner],
22 + inputs: Sequence[tuple[list[str], PromptImageInput]],
23 + model: str,
24 + *,
25 + dtype: str,
26 + max_tokens: int,
27 + num_logprobs: int,
28 +) -> None:
29 + """Verify that the inference result is the same between hf and vllm."""
30 + with vllm_runner(
31 + model,
32 + dtype=dtype,
33 + max_num_seqs=64,
34 + limit_mm_per_prompt={"image": 1},
35 + trust_remote_code=True,
36 + ) as vllm_model:
37 + vllm_outputs_per_case = [
38 + vllm_model.generate_greedy_logprobs(
39 + prompts,
40 + max_tokens,
41 + num_logprobs=num_logprobs,
42 + images=images,
43 + )
44 + for prompts, images in inputs
45 + ]
46 +
47 + with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
48 + hf_outputs_per_case = [
49 + hf_model.generate_greedy_logprobs_limit(
50 + prompts,
51 + max_tokens,
52 + num_logprobs=num_logprobs,
53 + images=images,
54 + use_cache=False, # HF Nemotron Parse crashes here without this
55 + )
56 + for prompts, images in inputs
57 + ]
58 +
59 + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
60 + check_logprobs_close(
61 + outputs_0_lst=hf_outputs,
62 + outputs_1_lst=vllm_outputs,
63 + name_0="hf",
64 + name_1="vllm",
65 + )
66 +
67 +
68 +@pytest.mark.core_model
69 +@pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.1"])
70 +@pytest.mark.parametrize("dtype", ["bfloat16"])
71 +@pytest.mark.parametrize("num_logprobs", [5])
72 +@create_new_process_for_each_test("spawn")
73 +def test_models(
74 + hf_runner, vllm_runner, model: str, dtype: str, num_logprobs: int
75 +) -> None:
76 + run_test(
77 + hf_runner,
78 + vllm_runner,
79 + inputs=[
80 + (
81 + [PROMPT] * 10,
82 + [IMAGE] * 10,
83 + ),
84 + ],
85 + model=model,
86 + dtype=dtype,
87 + max_tokens=100,
88 + num_logprobs=num_logprobs,
89 + )
tests/models/multimodal/pooling/test_radio.py
+10 -8
@@ -40,15 +40,15 @@ def run_radio_test(
40 for image in images
41 ]
42
43 - config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
43 + hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
44
45 # RADIO model on HF does not properly handle torch_dtype argument
46 # And relies on args["dtype"] which we have to patch manually:
47 - config.args["dtype"] = torch_dtype
47 + hf_config.args["dtype"] = torch_dtype
48
49 hf_model = AutoModel.from_pretrained(
50 model_id,
51 - config=config,
51 + config=hf_config,
52 dtype=torch_dtype,
53 trust_remote_code=True,
54 ).to("cuda")
@@ -62,13 +62,14 @@ def run_radio_test(
62 hf_model.make_preprocessor_external()
63
64 hf_outputs_per_image = [
65 - hf_model(pixel_value.to("cuda")).features for pixel_value in pixel_values
65 + hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values
66 ]
67
68 - radio_config = RadioConfig(
69 - model_name=config.args["model"], reg_tokens=config.args["register_multiple"]
68 + vllm_config = RadioConfig(
69 + model_name=hf_config.args["model"],
70 + **hf_config.args,
71 )
71 - vllm_model = RadioModel(radio_config)
72 + vllm_model = RadioModel(vllm_config)
73 vllm_model.load_weights(hf_model.state_dict())
74 vllm_model = vllm_model.to("cuda", torch_dtype)
75
@@ -80,7 +81,8 @@ def run_radio_test(
81
82 cos_similar = nn.CosineSimilarity(dim=-1)
83 for vllm_output, hf_output in zip(vllm_outputs_per_image, hf_outputs_per_image):
83 - assert cos_similar(vllm_output, hf_output).mean() > 0.99
84 + assert cos_similar(vllm_output[0], hf_output[0]).mean() > 0.99
85 + assert cos_similar(vllm_output[1], hf_output[1]).mean() > 0.99
86
87
88 @pytest.mark.parametrize(
tests/models/multimodal/processing/test_common.py
+1
@@ -102,6 +102,7 @@ def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
102 # incorrect token ids. So we need use `add_special_tokens=False` here
103 # to leave bos_token to be added by the processor.
104 _ADD_SPECIAL_TOKENS_OVERRIDES = {
105 + "nemotron_parse": False,
106 "ovis": False,
107 "ovis2_5": False,
108 "paligemma": False,
tests/models/registry.py
+3
@@ -907,6 +907,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
907 is_available_online=False,
908 ),
909 # [Encoder-decoder]
910 + "NemotronParseForConditionalGeneration": _HfExamplesInfo(
911 + "nvidia/NVIDIA-Nemotron-Parse-v1.1", trust_remote_code=True
912 + ),
913 "WhisperForConditionalGeneration": _HfExamplesInfo(
914 "openai/whisper-large-v3-turbo",
915 extras={"v3": "openai/whisper-large-v3"},
vllm/assets/image.py
+5 -2
@@ -42,8 +42,11 @@ class ImageAsset:
42 )
43
44 @property
45 - def pil_image(self, ext="jpg") -> Image.Image:
46 - image_path = self.get_path(ext)
45 + def pil_image(self) -> Image.Image:
46 + return self.pil_image_ext(ext="jpg")
47 +
48 + def pil_image_ext(self, ext: str) -> Image.Image:
49 + image_path = self.get_path(ext=ext)
50 return Image.open(image_path)
51
52 @property
vllm/model_executor/models/nano_nemotron_vl.py
+2 -7
@@ -1220,7 +1220,7 @@ class NemotronH_Nano_VL_V2(
1220 n = pixel_values.shape[0]
1221 vit_embeds_list = []
1222 for i in range(0, n, micro_batch_size):
1223 - vit_embeds = self.vision_model(pixel_values[i : i + micro_batch_size])
1223 + _, vit_embeds = self.vision_model(pixel_values[i : i + micro_batch_size])
1224 vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
1225 h = w = int(vit_embeds.shape[1] ** 0.5)
1226 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
@@ -1695,12 +1695,7 @@ class NemotronH_Nano_VL_V2(
1695 patch_size=patch_size,
1696 norm_mean=hf_config.norm_mean,
1697 norm_std=hf_config.norm_std,
1698 - reg_tokens=(
1699 - hf_config_vision.args.get("register_multiple")
1700 - if hasattr(hf_config_vision, "args")
1701 - and isinstance(hf_config_vision.args, dict)
1702 - else None
1703 - ),
1698 + **hf_config_vision.args,
1699 )
1700
1701 return RadioModel(config=radio_config)
vllm/model_executor/models/nemotron_parse.py new
+958
@@ -0,0 +1,958 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Adapted from https://github.com/amalad/vllm/blob/nemotron_parse/vllm/model_executor/models/nemotron_parse.py
5 +# that's based on https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.1/blob/main/hf_nemotron_parse_modeling.py
6 +#
7 +# Bart classes based on old vLLM codebase:
8 +# https://github.com/vllm-project/vllm/blob/v0.10.2/vllm/model_executor/models/bart.py
9 +
10 +import math
11 +from collections.abc import Iterable, Mapping, Sequence
12 +from typing import Annotated, Literal
13 +
14 +import cv2
15 +import numpy as np
16 +import torch
17 +import torch.nn as nn
18 +from einops import rearrange
19 +from PIL import Image
20 +from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
21 +from torchvision import transforms as T
22 +from transformers import (
23 + BartConfig,
24 + BatchFeature,
25 + PretrainedConfig,
26 + TensorType,
27 +)
28 +
29 +from vllm.attention.backends.abstract import AttentionType
30 +from vllm.config import CacheConfig, VllmConfig
31 +from vllm.config.lora import LoRAConfig
32 +from vllm.config.multimodal import BaseDummyOptions
33 +from vllm.logger import init_logger
34 +from vllm.model_executor.layers.activation import get_act_fn
35 +from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
36 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
37 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
38 +from vllm.model_executor.layers.vocab_parallel_embedding import (
39 + ParallelLMHead,
40 + VocabParallelEmbedding,
41 +)
42 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43 +from vllm.model_executor.models.interfaces import (
44 + MultiModalEmbeddings,
45 + SupportsMultiModal,
46 +)
47 +from vllm.model_executor.models.radio import RadioModel
48 +from vllm.model_executor.models.whisper import WhisperAttention, WhisperCrossAttention
49 +from vllm.multimodal import MULTIMODAL_REGISTRY
50 +from vllm.multimodal.inputs import (
51 + MultiModalDataDict,
52 + MultiModalFieldConfig,
53 + MultiModalKwargsItems,
54 +)
55 +from vllm.multimodal.parse import MultiModalDataItems
56 +from vllm.multimodal.processing import (
57 + BaseProcessingInfo,
58 + EncDecMultiModalProcessor,
59 + PromptReplacement,
60 + PromptUpdate,
61 +)
62 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
63 +from vllm.transformers_utils.configs.radio import RadioConfig
64 +from vllm.transformers_utils.tokenizer import AnyTokenizer
65 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
66 +
67 +logger = init_logger(__name__)
68 +DEFAULT_FINAL_IMAGE_SIZE = (2048, 1648)
69 +
70 +
71 +class BartScaledWordEmbedding(VocabParallelEmbedding):
72 + """
73 + This module overrides VocabParallelEmbedding's
74 + forward by multiplying with embeddings scale.
75 + """
76 +
77 + def __init__(
78 + self, num_embeddings: int, embedding_dim: int, embed_scale: float = 1.0
79 + ):
80 + super().__init__(num_embeddings, embedding_dim)
81 + self.embed_scale = embed_scale
82 +
83 + def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
84 + return super().forward(input_ids) * self.embed_scale
85 +
86 +
87 +class BartParallelLMHead(ParallelLMHead):
88 + """
89 + This module overrides ParallelLMHead's
90 + forward by dividing by embeddings scale,
91 + yielding effectively the inverse of
92 + BartScaledWordEmbedding
93 + """
94 +
95 + def __init__(
96 + self, num_embeddings: int, embedding_dim: int, embed_scale: float = 1.0
97 + ):
98 + super().__init__(num_embeddings, embedding_dim)
99 + self.embed_scale = embed_scale
100 +
101 + def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
102 + return super().forward(input_ids) / self.embed_scale
103 +
104 +
105 +class BartDecoderLayer(nn.Module):
106 + def __init__(
107 + self,
108 + config: BartConfig,
109 + cache_config: CacheConfig | None = None,
110 + quant_config: QuantizationConfig | None = None,
111 + prefix: str = "",
112 + ):
113 + super().__init__()
114 + self.embed_dim = config.d_model
115 +
116 + self.self_attn = WhisperAttention(
117 + embed_dim=self.embed_dim,
118 + num_heads=config.decoder_attention_heads,
119 + attn_type=AttentionType.DECODER,
120 + cache_config=cache_config,
121 + quant_config=quant_config,
122 + prefix=f"{prefix}.self_attn",
123 + )
124 + self.activation_fn = get_act_fn(config.activation_function)
125 +
126 + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
127 + """
128 + afeldman-nm: personally I would call this "cross-attention",
129 + however I left the name as "encoder_attn" to maintain consistency
130 + with the name of the pretrained weights.
131 + """
132 + self.encoder_attn = WhisperCrossAttention(
133 + self.embed_dim,
134 + config.decoder_attention_heads,
135 + cache_config=cache_config,
136 + quant_config=quant_config,
137 + prefix=f"{prefix}.encoder_attn",
138 + )
139 + self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
140 +
141 + ffn_hidden_size = self.embed_dim
142 + ffn_intermediate_size = config.encoder_ffn_dim
143 + ffn_has_bias = True
144 + self.fc1 = ColumnParallelLinear(
145 + ffn_hidden_size,
146 + ffn_intermediate_size,
147 + bias=ffn_has_bias,
148 + quant_config=quant_config,
149 + prefix=f"{prefix}.fc1",
150 + )
151 + self.fc2 = RowParallelLinear(
152 + ffn_intermediate_size,
153 + ffn_hidden_size,
154 + bias=ffn_has_bias,
155 + quant_config=quant_config,
156 + prefix=f"{prefix}.fc2",
157 + )
158 +
159 + self.final_layer_norm = nn.LayerNorm(self.embed_dim)
160 +
161 + def forward(
162 + self,
163 + decoder_hidden_states: torch.Tensor,
164 + encoder_hidden_states: torch.Tensor | None = None,
165 + ) -> torch.Tensor:
166 + r"""
167 + Args:
168 + decoder_hidden_states: torch.Tensor of *decoder* input embeddings.
169 + encoder_hidden_states: torch.Tensor of *encoder* input embeddings.
170 + Returns:
171 + Decoder layer output torch.Tensor
172 + """
173 + residual = decoder_hidden_states
174 +
175 + # Self Attention
176 + hidden_states = self.self_attn(hidden_states=decoder_hidden_states)
177 +
178 + hidden_states = residual + hidden_states
179 + hidden_states = self.self_attn_layer_norm(hidden_states)
180 +
181 + # Cross-Attention Block
182 +
183 + residual = hidden_states
184 +
185 + hidden_states = self.encoder_attn(
186 + hidden_states=hidden_states,
187 + encoder_hidden_states=encoder_hidden_states,
188 + )
189 +
190 + hidden_states = residual + hidden_states
191 + hidden_states = self.encoder_attn_layer_norm(hidden_states)
192 +
193 + # Fully Connected
194 + residual = hidden_states
195 + fc1_out, _ = self.fc1(hidden_states)
196 + hidden_states = self.activation_fn(fc1_out)
197 +
198 + hidden_states, _ = self.fc2(hidden_states)
199 +
200 + hidden_states = residual + hidden_states
201 + hidden_states = self.final_layer_norm(hidden_states)
202 +
203 + return hidden_states
204 +
205 +
206 +class MBartDecoderLayer(BartDecoderLayer):
207 + def forward(
208 + self,
209 + decoder_hidden_states: torch.Tensor,
210 + encoder_hidden_states: torch.Tensor | None = None,
211 + ) -> torch.Tensor:
212 + residual = decoder_hidden_states
213 + hidden_states = self.self_attn_layer_norm(decoder_hidden_states)
214 +
215 + # Self Attention
216 + hidden_states = self.self_attn(hidden_states=hidden_states)
217 +
218 + hidden_states = residual + hidden_states
219 +
220 + # Cross-Attention Block
221 +
222 + residual = hidden_states
223 + hidden_states = self.encoder_attn_layer_norm(hidden_states)
224 +
225 + hidden_states = self.encoder_attn(
226 + hidden_states=hidden_states,
227 + encoder_hidden_states=encoder_hidden_states,
228 + )
229 +
230 + hidden_states = residual + hidden_states
231 +
232 + # Fully Connected
233 + residual = hidden_states
234 + hidden_states = self.final_layer_norm(hidden_states)
235 + fc1_out, _ = self.fc1(hidden_states)
236 + hidden_states = self.activation_fn(fc1_out)
237 +
238 + hidden_states, _ = self.fc2(hidden_states)
239 +
240 + hidden_states = residual + hidden_states
241 +
242 + return hidden_states
243 +
244 +
245 +class MBartDecoderNoPos(nn.Module):
246 + """
247 + Transformer decoder consisting of *config.decoder_layers* layers.
248 + Each layer is a [`BartDecoderLayer`]
249 + Args:
250 + config: BartConfig
251 + embed_tokens (nn.Embedding): output embedding
252 + """
253 +
254 + def __init__(
255 + self,
256 + config: BartConfig,
257 + cache_config: CacheConfig | None = None,
258 + quant_config: QuantizationConfig | None = None,
259 + lora_config: LoRAConfig | None = None,
260 + embed_tokens: nn.Embedding | None = None,
261 + prefix: str = "",
262 + ):
263 + super().__init__()
264 + self.cache_config = cache_config
265 + self.quant_config = quant_config
266 + self.lora_config = lora_config
267 + embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
268 +
269 + self.embed_tokens = BartScaledWordEmbedding(
270 + config.vocab_size, config.d_model, embed_scale=embed_scale
271 + )
272 +
273 + if embed_tokens is not None:
274 + self.embed_tokens.weight = embed_tokens.weight
275 +
276 + self.layers = nn.ModuleList(
277 + [
278 + MBartDecoderLayer(
279 + config,
280 + cache_config,
281 + quant_config,
282 + prefix=f"{prefix}.layers.{layer_idx}",
283 + )
284 + for layer_idx in range(config.decoder_layers)
285 + ]
286 + )
287 +
288 + self.layernorm_embedding = nn.LayerNorm(config.d_model)
289 + self.layer_norm = nn.LayerNorm(config.d_model)
290 +
291 + def forward(
292 + self,
293 + decoder_input_ids: torch.Tensor,
294 + *,
295 + encoder_hidden_states: torch.Tensor | None,
296 + inputs_embeds: torch.Tensor | None = None,
297 + **kwargs,
298 + ) -> torch.Tensor:
299 + r"""
300 + Args:
301 + decoder_input_ids: Indices of *decoder* input sequence tokens in the
302 + vocabulary. Padding will be ignored by default should you provide it.
303 + encoder_hidden_states: Tensor of encoder output embeddings
304 + Returns:
305 + Decoder output torch.Tensor
306 + """
307 + if inputs_embeds is None:
308 + inputs_embeds = self.embed_tokens(decoder_input_ids)
309 +
310 + hidden_states = self.layernorm_embedding(inputs_embeds)
311 +
312 + # decoder layers
313 +
314 + for decoder_layer in self.layers:
315 + hidden_states = decoder_layer(
316 + decoder_hidden_states=hidden_states,
317 + encoder_hidden_states=encoder_hidden_states,
318 + )
319 +
320 + hidden_states = self.layer_norm(hidden_states)
321 + return hidden_states
322 +
323 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
324 + stacked_params_mapping = [
325 + # (param_name, shard_name, shard_id)
326 + (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
327 + (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
328 + (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
329 + (".encoder_attn.kv_proj", ".encoder_attn.k_proj", "k"),
330 + (".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),
331 + ]
332 + params_dict = dict(self.named_parameters())
333 + loaded_params: set[str] = set()
334 + for name, loaded_weight in weights:
335 + if name.startswith("embed_positions"):
336 + continue
337 +
338 + for param_name, weight_name, shard_id in stacked_params_mapping:
339 + if weight_name not in name:
340 + continue
341 + name = name.replace(weight_name, param_name)
342 + # Skip loading extra bias for GPTQ models.
343 + if name.endswith(".bias") and name not in params_dict:
344 + continue
345 +
346 + param = params_dict[name]
347 + weight_loader = param.weight_loader
348 + weight_loader(param, loaded_weight, shard_id)
349 + break
350 + else:
351 + # Skip loading extra bias for GPTQ models.
352 + if name.endswith(".bias") and name not in params_dict:
353 + continue
354 +
355 + param = params_dict[name]
356 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
357 + weight_loader(param, loaded_weight)
358 + loaded_params.add(name)
359 + return loaded_params
360 +
361 +
362 +class NemotronParsePixelInputs(TensorSchema):
363 + """
364 + Dimensions:
365 + - b: Batch size
366 + - c: Number of channels (3)
367 + - h: Height
368 + - w: Width
369 + """
370 +
371 + type: Literal["pixel_values"]
372 + data: Annotated[torch.Tensor, TensorShape("b", 3, "h", "w")]
373 +
374 +
375 +class NemotronParseImageProcessor:
376 + """
377 + NemotronParse Image Processor
378 + """
379 +
380 + def __init__(
381 + self,
382 + final_size: tuple = DEFAULT_FINAL_IMAGE_SIZE,
383 + **kwargs,
384 + ):
385 + # Ensure final_size is properly formatted
386 + if isinstance(final_size, (list, tuple)) and len(final_size) >= 2:
387 + self.final_size = (int(final_size[0]), int(final_size[1]))
388 + elif isinstance(final_size, (int, float)):
389 + self.final_size = (int(final_size), int(final_size))
390 + else:
391 + self.final_size = DEFAULT_FINAL_IMAGE_SIZE # Default fallback
392 +
393 + self.norm_mean = torch.Tensor(OPENAI_CLIP_MEAN).reshape(1, 3, 1, 1)
394 + self.norm_std = torch.Tensor(OPENAI_CLIP_STD).reshape(1, 3, 1, 1)
395 +
396 + # Create transforms
397 + self._create_transforms()
398 +
399 + def _create_transforms(self):
400 + """Create transform objects."""
401 + try:
402 + import albumentations as A
403 + except ImportError as err:
404 + raise ImportError(
405 + "The package `albumentations` is required to use "
406 + "NemotronParse model. Please install it with `pip install "
407 + "albumentations`."
408 + ) from err
409 +
410 + # Ensure final_size is a tuple of integers
411 + if isinstance(self.final_size, (list, tuple)):
412 + self.target_height, self.target_width = (
413 + int(self.final_size[0]),
414 + int(self.final_size[1]),
415 + )
416 + else:
417 + self.target_height = self.target_width = int(self.final_size)
418 +
419 + self.transform = A.Compose(
420 + [
421 + A.PadIfNeeded(
422 + min_height=self.target_height,
423 + min_width=self.target_width,
424 + border_mode=cv2.BORDER_CONSTANT,
425 + fill=[255, 255, 255],
426 + p=1.0,
427 + ),
428 + ]
429 + )
430 +
431 + self.torch_transform = T.Compose(
432 + [
433 + T.ToTensor(),
434 + ]
435 + )
436 +
437 + def _resize_with_aspect_ratio(self, image: np.ndarray) -> np.ndarray:
438 + """Resize image maintaining aspect ratio (exact replica of original
439 + LongestMaxSizeHW)."""
440 + height, width = image.shape[:2]
441 + max_size_height = self.target_height
442 + max_size_width = self.target_width
443 +
444 + # Original LongestMaxSizeHW algorithm from custom_augmentations.py
445 + aspect_ratio = width / height
446 + new_height = height
447 + new_width = width
448 +
449 + # If height too big then scale image down
450 + if height > max_size_height:
451 + new_height = max_size_height
452 + new_width = int(new_height * aspect_ratio)
453 +
454 + # If width too big, scale image down further
455 + if new_width > max_size_width:
456 + new_width = max_size_width
457 + new_height = int(new_width / aspect_ratio)
458 +
459 + # Use cv2.INTER_LINEAR like the original
460 + return cv2.resize(
461 + image, (new_width, new_height), interpolation=cv2.INTER_LINEAR
462 + )
463 +
464 + def _pad_to_size(self, image: np.ndarray) -> np.ndarray:
465 + """Pad image to target size with white padding (matches A.PadIfNeeded
466 + behavior)."""
467 + h, w = image.shape[:2]
468 + min_height, min_width = self.target_height, self.target_width
469 +
470 + # Only pad if image is smaller than target (matches A.PadIfNeeded logic)
471 + pad_h = max(0, min_height - h)
472 + pad_w = max(0, min_width - w)
473 +
474 + if pad_h == 0 and pad_w == 0:
475 + return image
476 +
477 + # A.PadIfNeeded pads to bottom-right with constant value
478 + if len(image.shape) == 3:
479 + # Color image - pad bottom and right with white (255, 255, 255)
480 + padded = np.pad(
481 + image,
482 + ((0, pad_h), (0, pad_w), (0, 0)),
483 + mode="constant",
484 + constant_values=255,
485 + )
486 + else:
487 + # Grayscale image - pad with white (255)
488 + padded = np.pad(
489 + image, ((0, pad_h), (0, pad_w)), mode="constant", constant_values=255
490 + )
491 +
492 + return padded
493 +
494 + def preprocess(
495 + self,
496 + images: Image.Image | list[Image.Image],
497 + **kwargs,
498 + ) -> dict[str, torch.Tensor]:
499 + """
500 + Preprocess an image or batch of images for the NemotronParse model.
501 +
502 + Args:
503 + images: Input image(s)
504 + """
505 + # Ensure images is a list
506 + if not isinstance(images, list):
507 + images = [images]
508 +
509 + # Convert PIL images to numpy arrays if needed
510 + processed_images = []
511 + for image in images:
512 + if isinstance(image, Image.Image):
513 + image = np.asarray(image)
514 + processed_images.append(image)
515 +
516 + # Apply NemotronParse-specific transforms
517 + pixel_values = []
518 + for image in processed_images:
519 + # Manual resize with aspect ratio preservation
520 + # (replaces LongestMaxSizeHW)
521 + processed_image = self._resize_with_aspect_ratio(image)
522 +
523 + # Apply remaining albumentations transforms if available
524 + if self.transform is not None:
525 + transformed = self.transform(image=processed_image)
526 + processed_image = transformed["image"]
527 + else:
528 + # Fallback: just pad to target size
529 + processed_image = self._pad_to_size(processed_image)
530 +
531 + # Convert to tensor
532 + pixel_values_tensor = self.torch_transform(processed_image)
533 +
534 + # Handle grayscale images
535 + if pixel_values_tensor.shape[0] == 1:
536 + pixel_values_tensor = pixel_values_tensor.expand(3, -1, -1)
537 +
538 + pixel_values.append(pixel_values_tensor)
539 +
540 + # Stack into batch
541 + pixel_values = torch.stack(pixel_values)
542 +
543 + # Normalize pixel values
544 + normalized_values = (pixel_values - self.norm_mean) / self.norm_std
545 + return {"pixel_values": normalized_values}
546 +
547 + def __call__(
548 + self, images: Image.Image | list[Image.Image], **kwargs
549 + ) -> dict[str, torch.Tensor]:
550 + return self.preprocess(images, **kwargs)
551 +
552 +
553 +class NemotronParseProcessor:
554 + """
555 + NemotronParse Processor
556 + """
557 +
558 + def __init__(
559 + self,
560 + config: PretrainedConfig,
561 + tokenizer: AnyTokenizer,
562 + **kwargs,
563 + ) -> None:
564 + super().__init__()
565 +
566 + self.config = config
567 + self.tokenizer = tokenizer
568 +
569 + self.image_processor = NemotronParseImageProcessor(final_size=config.image_size)
570 +
571 + def _make_batch_input(self, input_item=None):
572 + if input_item is None:
573 + input_item = []
574 + if not isinstance(input_item, list):
575 + input_item = [input_item]
576 + return input_item
577 +
578 + def __call__(
579 + self,
580 + text: str | None = None,
581 + images: Image.Image | list[Image.Image] | None = None,
582 + return_tensors: str | TensorType | None = None,
583 + **kwargs,
584 + ) -> BatchFeature:
585 + text, images = [self._make_batch_input(x) for x in (text, images)]
586 + image_inputs = {} if len(images) == 0 else self.image_processor(images)
587 +
588 + text_inputs = self.tokenizer(text, add_special_tokens=False, **kwargs)
589 + combined_outputs = BatchFeature(
590 + data={**text_inputs, **image_inputs},
591 + tensor_type=return_tensors,
592 + )
593 + return combined_outputs
594 +
595 +
596 +class NemotronParseProcessingInfo(BaseProcessingInfo):
597 + def get_hf_config(self):
598 + return self.ctx.get_hf_config()
599 +
600 + def get_hf_processor(self, **kwargs) -> NemotronParseProcessor:
601 + return self.ctx.init_processor(
602 + NemotronParseProcessor,
603 + config=self.get_hf_config(),
604 + tokenizer=self.get_tokenizer(),
605 + **kwargs,
606 + )
607 +
608 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
609 + return {"image": 1}
610 +
611 + def get_num_image_tokens(self) -> int:
612 + config = self.get_hf_config()
613 + final_size = config.image_size
614 + patch_size = config.encoder.patch_size
615 +
616 + return (final_size[0] // patch_size) * ((final_size[1] // patch_size) // 4) + 1
617 +
618 + def get_mm_max_tokens_per_item(
619 + self,
620 + seq_len: int,
621 + mm_counts: Mapping[str, int],
622 + ) -> Mapping[str, int] | None:
623 + image_tokens = self.get_num_image_tokens()
624 + return {"image": image_tokens}
625 +
626 +
627 +class NemotronParseDummyInputsBuilder(
628 + BaseDummyInputsBuilder[NemotronParseProcessingInfo]
629 +):
630 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
631 + return ""
632 +
633 + def get_dummy_mm_data(
634 + self,
635 + seq_len: int,
636 + mm_counts: Mapping[str, int],
637 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
638 + ) -> MultiModalDataDict:
639 + num_images = mm_counts.get("image", 0)
640 +
641 + target_width, target_height = self.info.get_hf_config().image_size
642 +
643 + return {
644 + "image": self._get_dummy_images(
645 + width=target_width, height=target_height, num_images=num_images
646 + )
647 + }
648 +
649 +
650 +class NemotronParseMultiModalProcessor(
651 + EncDecMultiModalProcessor[NemotronParseProcessingInfo]
652 +):
653 + def create_encoder_prompt(
654 + self,
655 + prompt: str | list[int],
656 + mm_data: MultiModalDataDict,
657 + ) -> str | list[int]:
658 + return [0]
659 +
660 + @property
661 + def pad_dummy_encoder_prompt(self) -> bool:
662 + return True
663 +
664 + def _call_hf_processor(
665 + self,
666 + prompt: str,
667 + mm_data: Mapping[str, object],
668 + mm_kwargs: Mapping[str, object],
669 + tok_kwargs: Mapping[str, object],
670 + ) -> BatchFeature:
671 + if mm_data:
672 + processed_outputs = super()._call_hf_processor(
673 + prompt, mm_data, mm_kwargs, tok_kwargs
674 + )
675 + else:
676 + hf_processor = self.info.get_hf_processor()
677 + tokenizer = hf_processor.tokenizer
678 + processed_outputs = tokenizer(
679 + prompt, add_special_tokens=False, return_tensors="pt"
680 + )
681 + return processed_outputs
682 +
683 + def _get_mm_fields_config(
684 + self,
685 + hf_inputs: BatchFeature,
686 + hf_processor_mm_kwargs: Mapping[str, object],
687 + ) -> Mapping[str, MultiModalFieldConfig]:
688 + return dict(pixel_values=MultiModalFieldConfig.batched("image"))
689 +
690 + def _get_prompt_updates(
691 + self,
692 + mm_items: MultiModalDataItems,
693 + hf_processor_mm_kwargs: Mapping[str, object],
694 + out_mm_kwargs: MultiModalKwargsItems,
695 + ) -> Sequence[PromptUpdate]:
696 + num_image_tokens = self.info.get_num_image_tokens()
697 +
698 + return [
699 + PromptReplacement(
700 + modality="image",
701 + target=[0],
702 + replacement=[0] * num_image_tokens,
703 + )
704 + ]
705 +
706 +
707 +class RadioWithNeck(nn.Module):
708 + """Vision encoder using RADIO model with custom neck."""
709 +
710 + def __init__(
711 + self,
712 + config: PretrainedConfig,
713 + quant_config: QuantizationConfig | None = None,
714 + prefix: str = "",
715 + ):
716 + super().__init__()
717 + self.config = config.encoder
718 +
719 + self.model_encoder = self.get_vit_model_from_radio_config(
720 + config, quant_config=quant_config
721 + )
722 +
723 + # Neck components
724 + last_hidden_state = 1024
725 + self.conv1 = nn.Conv1d(1280, last_hidden_state, 1)
726 + self.layer_norm1 = nn.LayerNorm(
727 + last_hidden_state, eps=1e-06, elementwise_affine=True
728 + )
729 + self.conv2 = nn.Conv2d(
730 + last_hidden_state,
731 + last_hidden_state,
732 + kernel_size=(1, 4),
733 + stride=(1, 4),
734 + padding=0,
735 + bias=False,
736 + )
737 + self.layer_norm2 = nn.LayerNorm(
738 + last_hidden_state, eps=1e-06, elementwise_affine=True
739 + )
740 + self.sum_proj = ColumnParallelLinear(
741 + 3840,
742 + last_hidden_state,
743 + quant_config=quant_config,
744 + prefix=f"{prefix}.sum_proj",
745 + )
746 + self.layer_norm3 = nn.LayerNorm(
747 + last_hidden_state, eps=1e-06, elementwise_affine=True
748 + )
749 +
750 + def get_vit_model_from_radio_config(
751 + self,
752 + hf_config: PretrainedConfig,
753 + quant_config: QuantizationConfig | None = None,
754 + ) -> RadioModel:
755 + hf_config_vision = hf_config.encoder
756 + model_name = hf_config_vision.args.get("model")
757 + if model_name is None:
758 + raise ValueError(f"Unsupported vit model type: {model_name}")
759 +
760 + radio_config = RadioConfig(
761 + model_name=model_name,
762 + image_size=hf_config.image_size,
763 + **hf_config_vision.args,
764 + )
765 +
766 + return RadioModel(config=radio_config, quant_config=quant_config)
767 +
768 + def forward(self, pixel_values: torch.Tensor, **kwargs) -> torch.Tensor:
769 + summary, feature = self.model_encoder(pixel_values)
770 +
771 + output = self.conv1(feature.permute(0, 2, 1)).permute(0, 2, 1)
772 + output = self.layer_norm1(output)
773 +
774 + patch_size = self.config.patch_size
775 + output = rearrange(
776 + output,
777 + "b (h w) d -> b d h w",
778 + h=pixel_values.shape[-2] // patch_size,
779 + w=pixel_values.shape[-1] // patch_size,
780 + )
781 +
782 + output = self.conv2(output)
783 + output = rearrange(output, "b d h w -> b (h w) d")
784 + output = self.layer_norm2(output)
785 + summary = self.layer_norm3(self.sum_proj(summary)[0])
786 + output = torch.cat((output, summary.unsqueeze(1)), dim=1)
787 +
788 + return output
789 +
790 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
791 + model_encoder_weights = []
792 + adaptor_dict = {
793 + name: param
794 + for name, param in dict(self.named_parameters()).items()
795 + if not name.startswith("model_encoder")
796 + }
797 + for name, w in weights:
798 + if name.startswith("model_encoder"):
799 + model_encoder_weights.append((".".join(name.split(".")[1:]), w))
800 + else:
801 + param = adaptor_dict[name]
802 + with torch.no_grad():
803 + default_weight_loader(param, w)
804 +
805 + self.model_encoder.load_weights(model_encoder_weights)
806 +
807 +
808 +@MULTIMODAL_REGISTRY.register_processor(
809 + NemotronParseMultiModalProcessor,
810 + info=NemotronParseProcessingInfo,
811 + dummy_inputs=NemotronParseDummyInputsBuilder,
812 +)
813 +class NemotronParseForConditionalGeneration(nn.Module, SupportsMultiModal):
814 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
815 + super().__init__()
816 + config = vllm_config.model_config.hf_config
817 +
818 + self.config = config
819 + self.vision_config = config.encoder
820 + cache_config = vllm_config.cache_config
821 + quant_config = vllm_config.quant_config
822 +
823 + self.encoder = RadioWithNeck(
824 + config=config, quant_config=quant_config, prefix=f"{prefix}.encoder"
825 + )
826 +
827 + self.decoder = MBartDecoderNoPos(
828 + config.decoder,
829 + cache_config=cache_config,
830 + quant_config=quant_config,
831 + prefix=f"{prefix}.decoder",
832 + )
833 +
834 + self.vocab_size = config.decoder.vocab_size
835 + self.lm_head = ParallelLMHead(
836 + config.decoder.vocab_size, config.decoder.d_model, quant_config=quant_config
837 + )
838 + self.logits_processor = LogitsProcessor(
839 + self.vocab_size, config.decoder.vocab_size
840 + )
841 +
842 + @classmethod
843 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
844 + if modality.startswith("image"):
845 + return None
846 +
847 + raise ValueError("Only image modality is supported")
848 +
849 + def _parse_and_validate_image_input(
850 + self, **kwargs: object
851 + ) -> NemotronParsePixelInputs | None:
852 + pixel_values = kwargs.pop("pixel_values", None)
853 + image_embeds = kwargs.pop("image_embeds", None)
854 +
855 + if pixel_values is None and image_embeds is None:
856 + return None
857 +
858 + if pixel_values is not None and image_embeds is not None:
859 + raise ValueError("Both pixel values and image embeds are provided.")
860 +
861 + if pixel_values is not None:
862 + h, w = self.config.image_size
863 + return NemotronParsePixelInputs(
864 + type="pixel_values",
865 + data=pixel_values,
866 + resolve_bindings={
867 + "h": h,
868 + "w": w,
869 + },
870 + )
871 +
872 + if image_embeds is not None:
873 + raise NotImplementedError
874 +
875 + raise AssertionError("This line should be unreachable.")
876 +
877 + def _process_image_input(
878 + self, image_input: NemotronParsePixelInputs
879 + ) -> torch.Tensor:
880 + assert image_input["type"] == "pixel_values"
881 + pixel_values = image_input["data"]
882 + dtype = next(self.encoder.parameters()).dtype
883 + pixel_values = pixel_values.to(dtype)
884 + return self.encoder(pixel_values)
885 +
886 + def get_language_model(self) -> torch.nn.Module:
887 + return self.decoder
888 +
889 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
890 + image_input = self._parse_and_validate_image_input(**kwargs)
891 + if image_input is None:
892 + return None
893 + vision_embeddings = self._process_image_input(image_input)
894 + return vision_embeddings
895 +
896 + def forward(
897 + self,
898 + input_ids: torch.Tensor,
899 + positions: torch.Tensor,
900 + encoder_outputs: list[torch.Tensor] | None = None,
901 + **kwargs,
902 + ) -> torch.Tensor:
903 + r"""
904 + Args:
905 + input_ids: torch.Tensor of *decoder* input token ids.
906 + positions: torch.Tensor of *decoder* position indices.
907 + encoder_outputs: List of encoder output tensors (vision embeddings).
908 + During profiling, this may be None or empty.
909 + Returns:
910 + Output torch.Tensor
911 + """
912 + inputs_embeds = None
913 + if encoder_outputs:
914 + inputs_embeds = torch.cat(encoder_outputs, dim=0)
915 + hidden_states = self.decoder(
916 + decoder_input_ids=input_ids, encoder_hidden_states=inputs_embeds
917 + )
918 + return hidden_states
919 +
920 + def compute_logits(
921 + self,
922 + hidden_states: torch.Tensor,
923 + ) -> torch.Tensor | None:
924 + return self.logits_processor(self.lm_head, hidden_states)
925 +
926 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
927 + lm_head_dict = dict(self.lm_head.named_parameters())
928 +
929 + def is_encoder(name: str) -> bool:
930 + return name.startswith("encoder")
931 +
932 + def is_decoder(name: str) -> bool:
933 + return name.startswith("decoder")
934 +
935 + def is_lm_head(name: str):
936 + return name.startswith("lm_head")
937 +
938 + # Separate weights by component
939 + encoder_weights = []
940 + decoder_weights = []
941 +
942 + for name, w in weights:
943 + if is_encoder(name):
944 + encoder_weights.append((".".join(name.split(".")[1:]), w))
945 + elif is_decoder(name):
946 + decoder_weights.append((".".join(name.split(".")[1:]), w))
947 + elif is_lm_head(name):
948 + trimmed_name = ".".join(name.split(".")[1:])
949 + param = lm_head_dict[trimmed_name]
950 + with torch.no_grad():
951 + default_weight_loader(param, w)
952 + else:
953 + logger.info("Found unexpected weight: %s", name)
954 +
955 + # Load encoder weights
956 + self.encoder.load_weights(encoder_weights)
957 + # Load decoder weights
958 + self.decoder.load_weights(decoder_weights)
vllm/model_executor/models/radio.py
+23 -5
@@ -427,15 +427,17 @@ class RadioInternVisionModel(nn.Module):
427 to_2tuple(config.patch_size), config.image_size
428 )
429 max_img_size = int(
430 - round(config.max_img_size / config.patch_size) * config.patch_size
430 + round(config.cpe_max_size / config.patch_size) * config.patch_size
431 )
432 + unique_teachers = set(t["name"] for t in config.teachers)
433 self.patch_generator = ViTPatchGenerator(
434 config.patch_size,
435 config.hidden_size,
436 input_dims=self.img_size,
437 max_input_dims=max_img_size,
438 cls_token=True,
438 - register_multiple=config.reg_tokens,
439 + num_cls_tokens=len(unique_teachers) if config.cls_token_per_teacher else 1,
440 + register_multiple=config.register_multiple,
441 )
442
443 self.encoder = InternVisionEncoder(
@@ -489,11 +491,20 @@ class RadioModel(nn.Module):
491 prefix=prefix,
492 )
493
494 + summary_idxs = None
495 + if config.teachers:
496 + summary_idxs = torch.tensor(
497 + [i for i, t in enumerate(config.teachers) if t.get("use_summary", True)]
498 + )
499 + if summary_idxs.numel() > 0:
500 + self.register_buffer("summary_idxs", summary_idxs)
501 + self.summary_idxs = summary_idxs
502 +
503 def forward(
504 self,
505 pixel_values: torch.Tensor | None = None,
506 pixel_embeds: torch.Tensor | None = None,
496 - ) -> torch.FloatTensor:
507 + ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
508 y = self.model(pixel_values)
509 return self._extract_final(y)
510
@@ -546,10 +557,17 @@ class RadioModel(nn.Module):
557
558 return loaded_params
559
549 - def _extract_final(self, y: torch.Tensor):
560 + def _extract_final(
561 + self, y: torch.Tensor
562 + ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
563 # Remove CLS + REGISTERS tokens
564 patch_gen = getattr(self.model, "patch_generator", None)
565 if patch_gen is not None:
566 + all_summary = y[:, : patch_gen.num_cls_tokens]
567 + if self.summary_idxs is not None:
568 + bb_summary = all_summary[:, self.summary_idxs]
569 + else:
570 + bb_summary = all_summary
571 all_feat = y[:, patch_gen.num_skip :]
572
555 - return all_feat
573 + return bb_summary.flatten(1), all_feat
vllm/model_executor/models/registry.py
+4
@@ -428,6 +428,10 @@ _MULTIMODAL_MODELS = {
428 "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), # noqa: E501
429 "VoxtralStreamingGeneration": ("voxtral_streaming", "VoxtralStreamingGeneration"), # noqa: E501
430 # [Encoder-decoder]
431 + "NemotronParseForConditionalGeneration": (
432 + "nemotron_parse",
433 + "NemotronParseForConditionalGeneration",
434 + ),
435 "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"), # noqa: E501
436 }
437
vllm/transformers_utils/configs/radio.py
+15 -6
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Radio vision model configuration"""
4
5 +from typing import Any
6 +
7 from transformers.configuration_utils import PretrainedConfig
8 from transformers.utils import logging
9
@@ -36,12 +38,15 @@ class RadioConfig(PretrainedConfig):
38 layer_norm_eps: The epsilon used by the layer normalization layers.
39 initializer_factor: A factor for initializing all weight matrices.
40 hidden_act: The non-linear activation function in the encoder.
39 - max_img_size: Maximum image size for position embeddings.
41 + cpe_max_size: Maximum image size for position embeddings.
42 norm_mean: Mean values for image normalization (RGB channels).
43 Defaults to (0.48145466, 0.4578275, 0.40821073)).
44 norm_std: Standard deviation values for image normalization
45 (RGB channels). Defaults to (0.26862954, 0.26130258, 0.27577711)).
44 - reg_tokens: Number of register tokens to use.
46 + register_multiple: Number of register tokens to use.
47 + teachers: A list of teacher model configurations. Each teacher configuration is
48 + a dict with keys like "name" and some may have "use_summary".
49 + cls_token_per_teacher: Whether to use a separate CLS token for each teacher.
50 """
51
52 model_type = "radio"
@@ -57,10 +62,12 @@ class RadioConfig(PretrainedConfig):
62 layer_norm_eps: float = 1e-6,
63 initializer_factor: float = 1.0,
64 hidden_act: str = "gelu",
60 - max_img_size: int = 2048,
65 + cpe_max_size: int = 2048,
66 norm_mean: tuple[float, float, float] | list = OPENAI_CLIP_MEAN,
67 norm_std: tuple[float, float, float] | list = OPENAI_CLIP_STD,
63 - reg_tokens: int | None = None,
68 + register_multiple: int | None = None,
69 + teachers: list[dict[str, Any]] | None = None,
70 + cls_token_per_teacher: bool = False,
71 **kwargs,
72 ):
73 self.model_name = model_name
@@ -78,12 +85,14 @@ class RadioConfig(PretrainedConfig):
85 self.layer_norm_eps = layer_norm_eps
86 self.initializer_factor = initializer_factor
87 self.hidden_act = hidden_act
81 - self.max_img_size = max_img_size
88 + self.cpe_max_size = cpe_max_size
89 self.norm_mean = (
90 list(norm_mean) if isinstance(norm_mean, (tuple, list)) else norm_mean
91 )
92 self.norm_std = (
93 list(norm_std) if isinstance(norm_std, (tuple, list)) else norm_std
94 )
88 - self.reg_tokens = reg_tokens
95 + self.register_multiple = register_multiple
96 + self.teachers = teachers if teachers is not None else []
97 + self.cls_token_per_teacher = cls_token_per_teacher
98 super().__init__(**kwargs)