72d30108a0 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,1395 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# --------------------------------------------------------

4 +# Adapted from

5 +# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/internvl.py

6 +# under Apache-2.0 License

7 +# LICENSE is in root directory.

8 +# --------------------------------------------------------

9 +

10 +import copy

11 +import warnings

12 +from abc import ABC, abstractmethod

13 +from collections.abc import Iterable, Mapping, Sequence

14 +from typing import Annotated, Any, Literal, Optional, TypedDict, TypeVar, Union

15 +

16 +import numpy.typing as npt

17 +import torch

18 +import torch.nn as nn

19 +import torchvision.transforms as T

20 +from PIL import Image

21 +from transformers import (AutoModel, BatchEncoding, BatchFeature,

22 + PretrainedConfig, TensorType)

23 +

24 +from vllm.config import VllmConfig

25 +from vllm.model_executor.layers.activation import ReLUSquaredActivation

26 +from vllm.model_executor.layers.layernorm import RMSNorm

27 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

28 +from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,

29 + MultiModalEmbeddings,

30 + SupportsMultiModal)

31 +from vllm.model_executor.models.internvl import (calculate_internvl_targets,

32 + get_internvl_target_ratios)

33 +from vllm.model_executor.models.module_mapping import MultiModelKeys

34 +from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM

35 +from vllm.model_executor.models.utils import (flatten_bn,

36 + init_vllm_registered_model,

37 + maybe_prefix,

38 + merge_multimodal_embeddings)

39 +from vllm.model_executor.sampling_metadata import SamplingMetadata

40 +from vllm.multimodal import MULTIMODAL_REGISTRY

41 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

42 + MultiModalKwargs, MultiModalKwargsItems,

43 + NestedTensors)

44 +from vllm.multimodal.parse import (ImageEmbeddingItems, ImageProcessorItems,

45 + ImageSize, MultiModalDataItems)

46 +from vllm.multimodal.processing import (BaseMultiModalProcessor,

47 + BaseProcessingInfo, PromptReplacement,

48 + PromptUpdate, PromptUpdateDetails)

49 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

50 +from vllm.sequence import IntermediateTensors

51 +from vllm.transformers_utils.tokenizer import AnyTokenizer

52 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

53 +

54 +# Configure PIL to handle large images without warnings

55 +# This prevents DecompressionBombWarning for legitimate large images

56 +Image.MAX_IMAGE_PIXELS = None # Disable the limit entirely

57 +# Alternative: Set a specific higher limit

58 +# Image.MAX_IMAGE_PIXELS = 300000000 # ~300M pixels

59 +

60 +IMG_START = "<img>"

61 +IMG_END = "</img>"

62 +IMG_CONTEXT = "<image>"

63 +

64 +# Profiling

65 +MAX_FRAMES = 16

66 +

67 +

68 +class NanoNemotronVLImagePixelInputs(TypedDict):

69 + type: Literal["pixel_values"]

70 + pixel_values_flat: torch.Tensor

71 + """

72 + Shape:

73 + `(batch_size * num_images * (1 + num_patches), num_channels, height, width)`

74 + """

75 +

76 + num_patches: torch.Tensor

77 + """Shape: `(batch_size * num_images)`"""

78 +

79 +

80 +class NanoNemotronVLImageEmbeddinInputs(TypedDict):

81 + type: Literal["image_embeds"]

82 + data: Union[torch.Tensor, list[torch.Tensor]]

83 + """

84 + A tensor of shape `(num_images, total_image_feature_size, hidden_size)`

85 + or a list of tensors of shape `(total_image_feature_size, hidden_size)`

86 +

87 + `hidden_size` must match the hidden size of language model backbone.

88 + """

89 +

90 +

91 +NanoNemotronVLImageInputs = Union[NanoNemotronVLImagePixelInputs,

92 + NanoNemotronVLImageEmbeddinInputs]

93 +

94 +

95 +class NanoNemotronVLVideoPixelInputs(TensorSchema):

96 + """

97 + Dimensions:

98 + - bvf: Batch size * number of videos * num_frames

99 + - bn: Batch size * number of images

100 + - c: Number of channels (3)

101 + - h: Height of each video frame

102 + - w: Width of each video frame

103 + """

104 + type: Literal["pixel_values_videos"]

105 + pixel_values_flat: Annotated[torch.Tensor, TensorShape("bvf", 3, "h", "w")]

106 + num_patches: Annotated[torch.Tensor, TensorShape("bn")]

107 +

108 +

109 +class NanoNemotronVLVideoEmbeddingInputs(TensorSchema):

110 + """

111 + Dimensions:

112 + - n: Number of videos

113 + - f: Total video feature size

114 + - h: Hidden size (must match the hidden size of language model backbone)

115 + """

116 + type: Literal["video_embeds"]

117 + data: Annotated[Union[torch.Tensor, list[torch.Tensor]],

118 + TensorShape("n", "f", "h")]

119 +

120 +

121 +NanoNemotronVLVideoInputs = Union[NanoNemotronVLVideoPixelInputs,

122 + NanoNemotronVLVideoEmbeddingInputs]

123 +

124 +

125 +def input_conditioner(x, norm_mean, norm_std):

126 + y = (x - norm_mean) / norm_std

127 + return y

128 +

129 +

130 +def dynamic_preprocess(image,

131 + *,

132 + image_size=512,

133 + max_num_tiles=12,

134 + use_thumbnail=True,

135 + idx=0):

136 + orig_width, orig_height = image.size

137 +

138 + target_ratios = get_internvl_target_ratios(1, max_num_tiles)

139 +

140 + blocks, target_width, target_height = calculate_internvl_targets(

141 + orig_width=orig_width,

142 + orig_height=orig_height,

143 + target_ratios=target_ratios,

144 + image_size=image_size,

145 + use_thumbnail=False)

146 + # resize the image

147 + resized_img = image.resize((target_width, target_height))

148 + processed_images = []

149 + for i in range(blocks):

150 + box = (

151 + (i % (target_width // image_size)) * image_size,

152 + (i // (target_width // image_size)) * image_size,

153 + ((i % (target_width // image_size)) + 1) * image_size,

154 + ((i // (target_width // image_size)) + 1) * image_size,

155 + )

156 + # split the image

157 + split_img = resized_img.crop(box)

158 + processed_images.append(split_img)

159 + assert len(processed_images) == blocks

160 + if use_thumbnail and len(processed_images) != 1:

161 + thumbnail_img = image.resize((image_size, image_size))

162 + processed_images.append(thumbnail_img)

163 +

164 + processed_images = [

165 + img.convert("RGB") if img.mode != "RGB" else img

166 + for img in processed_images

167 + ]

168 + processed_images = [

169 + T.Resize((image_size, image_size),

170 + interpolation=T.InterpolationMode.BICUBIC)(img)

171 + for img in processed_images

172 + ]

173 + processed_images = [T.ToTensor()(img) for img in processed_images]

174 + return processed_images

175 +

176 +

177 +def image_to_pixel_values(

178 + image: Image.Image,

179 + *,

180 + input_size: int,

181 + max_num: int,

182 + use_thumbnail: bool,

183 + idx: int,

184 +) -> torch.Tensor:

185 + images = dynamic_preprocess(

186 + image,

187 + image_size=input_size,

188 + max_num_tiles=max_num,

189 + use_thumbnail=use_thumbnail,

190 + idx=idx,

191 + )

192 +

193 + pixel_values = torch.stack(images)

194 + return pixel_values

195 +

196 +

197 +def video_to_pixel_values(

198 + video: npt.NDArray,

199 + *,

200 + input_size: int,

201 + max_num_tiles: int = 1,

202 + use_thumbnail: bool,

203 +) -> torch.Tensor:

204 + # Convert each frame to a single resized tile tensor consistent

205 + # with image path

206 + frames_tensors: list[torch.Tensor] = []

207 + for frame in video:

208 + pil_frame = dynamic_preprocess(

209 + Image.fromarray(frame, mode="RGB"),

210 + image_size=input_size,

211 + max_num_tiles=max_num_tiles,

212 + use_thumbnail=use_thumbnail,

213 + idx=0,

214 + )

215 + # dynamic_preprocess returns tensors already; take the single tile

216 + assert len(pil_frame) >= 1

217 + frames_tensors.append(pil_frame[0])

218 +

219 + return torch.stack(frames_tensors)

220 +

221 +

222 +class BaseNanoNemotronVLProcessor(ABC):

223 + """

224 + This model doesn't define its own HF processor,

225 + so we implement our own one here.

226 +

227 + The code to insert image tokens is based on:

228 + https://huggingface.co/OpenGVLab/InternVL2-1B/blob/main/modeling_internvl_chat.py#L252

229 + """

230 +

231 + def __init__(self, config: PretrainedConfig, tokenizer: AnyTokenizer,

232 + *args, **kwargs) -> None:

233 + super().__init__()

234 +

235 + self.config = config

236 + self.tokenizer = tokenizer

237 +

238 + image_size: int = config.force_image_size

239 + patch_size: int = config.patch_size

240 +

241 + self.num_image_token = int(

242 + (image_size // patch_size)**2 * (config.downsample_ratio**2))

243 + self.image_size = image_size

244 + self.use_thumbnail: bool = config.use_thumbnail

245 + self.norm_mean = torch.Tensor(config.norm_mean).reshape(1, 3, 1, 1)

246 + self.norm_std = torch.Tensor(config.norm_std).reshape(1, 3, 1, 1)

247 +

248 + @property

249 + @abstractmethod

250 + def image_token_id(self) -> int:

251 + raise NotImplementedError

252 +

253 + @abstractmethod

254 + def get_image_repl(

255 + self,

256 + feature_size: int,

257 + num_patches: Optional[int],

258 + ) -> PromptUpdateDetails[str]:

259 + raise NotImplementedError

260 +

261 + def get_num_image_tokens(

262 + self,

263 + *,

264 + image_width: int,

265 + image_height: int,

266 + max_num_tiles: int,

267 + ) -> int:

268 + target_ratios = get_internvl_target_ratios(1, max_num_tiles)

269 +

270 + num_patches, _, _ = calculate_internvl_targets(

271 + orig_width=image_width,

272 + orig_height=image_height,

273 + target_ratios=target_ratios,

274 + image_size=self.image_size,

275 + use_thumbnail=self.use_thumbnail,

276 + )

277 +

278 + return num_patches * self.num_image_token

279 +

280 + def _images_to_pixel_values_lst(

281 + self,

282 + images: list[Image.Image],

283 + max_num_tiles: int,

284 + ) -> list[torch.Tensor]:

285 + return [

286 + image_to_pixel_values(

287 + image,

288 + input_size=self.image_size,

289 + max_num=max_num_tiles,

290 + use_thumbnail=self.use_thumbnail,

291 + idx=idx,

292 + ) for idx, image in enumerate(images)

293 + ]

294 +

295 + def _preprocess_image(

296 + self,

297 + text: list[str],

298 + images: list[Image.Image],

299 + max_num_tiles: int,

300 + ) -> tuple[list[str], dict[str, torch.Tensor]]:

301 + if len(images) == 0:

302 + image_inputs = {}

303 + else:

304 + pixel_values_lst = self._images_to_pixel_values_lst(

305 + images, max_num_tiles)

306 + image_inputs: dict[str, NestedTensors] = {

307 + "pixel_values_flat":

308 + input_conditioner(torch.cat(pixel_values_lst), self.norm_mean,

309 + self.norm_std),

310 + "image_num_patches":

311 + torch.tensor([len(item) for item in pixel_values_lst]),

312 + }

313 +

314 + for pixel_values in pixel_values_lst:

315 + num_patches = pixel_values.shape[0]

316 + feature_size = num_patches * self.num_image_token

317 + image_repl = self.get_image_repl(feature_size, num_patches)

318 + text = [t.replace('<image>', image_repl.full, 1) for t in text]

319 + return text, image_inputs

320 +

321 + def _make_batch_input(self,

322 + input_item: Optional[Union[Any, list[Any]]] = None):

323 + if input_item is None:

324 + input_item = []

325 + if not isinstance(input_item, list):

326 + input_item = [input_item]

327 + return input_item

328 +

329 + def __call__(

330 + self,

331 + text: Optional[Union[str, list[str]]] = None,

332 + images: Optional[Union[Image.Image, list[Image.Image]]] = None,

333 + return_tensors: Optional[Union[str, TensorType]] = None,

334 + max_num_tiles: Optional[int] = None,

335 + ) -> Mapping[str, NestedTensors]:

336 + # Use default if not provided

337 + if max_num_tiles is None:

338 + max_num_tiles = 12

339 +

340 + text, images = [self._make_batch_input(x) for x in (text, images)]

341 +

342 + text, image_inputs = self._preprocess_image(

343 + text=text,

344 + images=images,

345 + max_num_tiles=max_num_tiles,

346 + )

347 +

348 + text_inputs = self.tokenizer(text, add_special_tokens=False)

349 +

350 + return {

351 + **BatchEncoding(text_inputs, tensor_type=return_tensors),

352 + **image_inputs,

353 + }

354 +

355 +

356 +class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

357 + """

358 + HF Processor with extended video processing logic.

359 + Code for video processing is adapted from video example:

360 + https://huggingface.co/OpenGVLab/InternVL3-1B#inference-with-transformers

361 + """

362 +

363 + def __init__(

364 + self,

365 + config: PretrainedConfig,

366 + tokenizer: AnyTokenizer,

367 + *,

368 + min_dynamic_patch: Optional[int] = None,

369 + max_dynamic_patch: Optional[int] = None,

370 + dynamic_image_size: Optional[bool] = None,

371 + video_token: Optional[str] = None,

372 + ) -> None:

373 + super().__init__(

374 + config=config,

375 + tokenizer=tokenizer,

376 + min_dynamic_patch=min_dynamic_patch,

377 + max_dynamic_patch=max_dynamic_patch,

378 + dynamic_image_size=dynamic_image_size,

379 + )

380 + # add extra video token for video processing

381 + self.video_token = video_token

382 +

383 + @property

384 + def supports_video(self) -> bool:

385 + return self.video_token_id is not None

386 +

387 + @property

388 + def video_token_id(self) -> Optional[int]:

389 + if self.video_token is None:

390 + return None

391 + return self.tokenizer.get_vocab().get(self.video_token, None)

392 +

393 + @property

394 + def image_token_id(self) -> int:

395 + return self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT)

396 +

397 + def _videos_to_pixel_values_lst(

398 + self,

399 + videos: list[npt.NDArray],

400 + max_num_tiles: int,

401 + dynamic_image_size: Optional[bool] = None,

402 + ) -> list[torch.Tensor]:

403 +

404 + return [

405 + video_to_pixel_values(

406 + video,

407 + input_size=self.image_size,

408 + max_num_tiles=max_num_tiles,

409 + use_thumbnail=self.use_thumbnail,

410 + ) for video in videos

411 + ]

412 +

413 + def _preprocess_video(

414 + self,

415 + text: list[str],

416 + videos: list[npt.NDArray],

417 + max_num_tiles: int,

418 + dynamic_image_size: Optional[bool] = None,

419 + ):

420 + if len(videos) == 0 or not self.supports_video:

421 + video_inputs = {}

422 + else:

423 + pixel_values_lst_video = self._videos_to_pixel_values_lst(

424 + videos,

425 + max_num_tiles=max_num_tiles,

426 + dynamic_image_size=dynamic_image_size,

427 + )

428 +

429 + video_inputs: dict[str, NestedTensors] = {

430 + "pixel_values_flat_video":

431 + input_conditioner(torch.cat(pixel_values_lst_video),

432 + self.norm_mean, self.norm_std),

433 + "video_num_patches":

434 + torch.tensor([len(item) for item in pixel_values_lst_video]),

435 + }

436 +

437 + for pixel_values in pixel_values_lst_video:

438 + num_patches = pixel_values.shape[0]

439 +

440 + video_repl = self.get_video_repl(self.num_image_token,

441 + num_patches, self.video_token)

442 + text = [t.replace('<video>', video_repl.full, 1) for t in text]

443 + return text, video_inputs

444 +

445 + def __call__(

446 + self,

447 + text: Optional[Union[str, list[str]]] = None,

448 + images: Optional[Union[Image.Image, list[Image.Image]]] = None,

449 + videos: Optional[Union[npt.NDArray, list[npt.NDArray]]] = None,

450 + return_tensors: Optional[Union[str, TensorType]] = None,

451 + max_num_tiles: Optional[int] = None,

452 + dynamic_image_size: Optional[bool] = None,

453 + ) -> Mapping[str, NestedTensors]:

454 + # Use default if not provided

455 + if max_num_tiles is None:

456 + max_num_tiles = 12

457 +

458 + text, images, videos = [

459 + self._make_batch_input(x) for x in (text, images, videos)

460 + ]

461 +

462 + text, image_inputs = self._preprocess_image(

463 + text=text,

464 + images=images,

465 + max_num_tiles=max_num_tiles,

466 + )

467 +

468 + text, video_inputs = self._preprocess_video(

469 + text=text,

470 + videos=videos,

471 + max_num_tiles=max_num_tiles,

472 + dynamic_image_size=dynamic_image_size,

473 + )

474 +

475 + text_inputs = self.tokenizer(text, add_special_tokens=False)

476 +

477 + return BatchFeature({

478 + **BatchEncoding(text_inputs, tensor_type=return_tensors),

479 + **image_inputs,

480 + **video_inputs,

481 + })

482 +

483 + def get_image_repl(

484 + self,

485 + feature_size: int,

486 + num_patches: Optional[int],

487 + ) -> PromptUpdateDetails[str]:

488 + repl_features = IMG_CONTEXT * feature_size

489 + repl_full = IMG_START + repl_features + IMG_END

490 +

491 + return PromptUpdateDetails.select_text(repl_full, IMG_CONTEXT)

492 +

493 + def get_video_repl(

494 + self,

495 + feature_size: int,

496 + num_patches: Optional[int] = None,

497 + video_context_token: str = IMG_CONTEXT,

498 + ) -> PromptUpdateDetails[str]:

499 + repl_features = video_context_token * self.num_image_token

500 + repl_features_with_sep = IMG_START + repl_features + IMG_END

501 + # num_patches is equal to num_frames

502 + repl_full = ''.join([

503 + f'Frame{i+1}: {repl_features_with_sep}' for i in range(num_patches)

504 + ])

505 +

506 + return PromptUpdateDetails.select_text(repl_full, video_context_token)

507 +

508 +

509 +class BaseNanoNemotronVLProcessingInfo(BaseProcessingInfo):

510 + """Basic image-only ProcessingInfo for InternVL-style models."""

511 +

512 + @abstractmethod

513 + def get_hf_processor(

514 + self,

515 + **kwargs: object,

516 + ) -> BaseNanoNemotronVLProcessor:

517 + raise NotImplementedError

518 +

519 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

520 + return {"image": None}

521 +

522 + def get_num_image_tokens(

523 + self,

524 + *,

525 + image_width: int,

526 + image_height: int,

527 + max_num_tiles: int,

528 + processor: Optional[BaseNanoNemotronVLProcessor],

529 + ) -> int:

530 + if processor is None:

531 + processor = self.get_hf_processor()

532 +

533 + return processor.get_num_image_tokens(

534 + image_width=image_width,

535 + image_height=image_height,

536 + max_num_tiles=max_num_tiles,

537 + )

538 +

539 + def get_image_size_with_most_features(self,

540 + max_num_tiles: int) -> ImageSize:

541 + processor = self.get_hf_processor()

542 +

543 + base_size = processor.image_size

544 + target_ratios = get_internvl_target_ratios(1, max_num_tiles)

545 +

546 + largest_feature_size, largest_feature_pinpoint = 0, None

547 + for wr, hr in target_ratios:

548 + width, height = base_size * wr, base_size * hr

549 +

550 + feat_size = self.get_num_image_tokens(

551 + image_width=width,

552 + image_height=height,

553 + max_num_tiles=max_num_tiles,

554 + processor=processor,

555 + )

556 + if feat_size > largest_feature_size:

557 + largest_feature_size = feat_size

558 + largest_feature_pinpoint = ImageSize(width=width,

559 + height=height)

560 +

561 + if largest_feature_size == 0 or largest_feature_pinpoint is None:

562 + raise ValueError("Cannot have a largest feature size of 0!")

563 +

564 + return largest_feature_pinpoint

565 +

566 + def get_max_image_tokens(self) -> int:

567 + processor = self.get_hf_processor()

568 + # Use default max_num_tiles for max tokens calculation

569 + max_num_tiles = 12

570 + target_width, target_height = self.get_image_size_with_most_features(

571 + max_num_tiles)

572 +

573 + return self.get_num_image_tokens(

574 + image_width=target_width,

575 + image_height=target_height,

576 + max_num_tiles=max_num_tiles,

577 + processor=processor,

578 + )

579 +

580 +

581 +_I = TypeVar("_I", bound=BaseNanoNemotronVLProcessingInfo)

582 +

583 +

584 +class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):

585 + """ ProcessingInfo extended for video processing"""

586 +

587 + @property

588 + def supports_video(self):

589 + return self.get_hf_processor().supports_video

590 +

591 + def get_supported_mm_limits(self):

592 + video_limit = {"video": None} if self.supports_video else {}

593 + return {**super().get_supported_mm_limits(), **video_limit}

594 +

595 + def get_video_token(self) -> Optional[str]:

596 + return IMG_CONTEXT

597 +

598 + def get_num_frames_with_most_features(

599 + self,

600 + seq_len: int,

601 + mm_counts: Mapping[str, int],

602 + ) -> int:

603 + max_images = mm_counts.get("image", 0)

604 + max_videos = mm_counts.get("video", 0)

605 +

606 + processor = self.get_hf_processor() # we get the CustomProcessor here

607 +

608 + max_image_tokens = self.get_max_image_tokens() * max_images

609 + max_total_frames = (seq_len -

610 + max_image_tokens) // processor.num_image_token

611 + max_frames_per_video = max_total_frames // max(max_videos, 1)

612 +

613 + max_frames_per_video = min(max_frames_per_video, MAX_FRAMES)

614 + return max(max_frames_per_video, 1)

615 +

616 + def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor:

617 + return self.ctx.init_processor(

618 + NanoNemotronVLProcessor,

619 + config=self.get_hf_config(),

620 + tokenizer=self.get_tokenizer(),

621 + video_token=self.get_video_token(),

622 + **kwargs,

623 + )

624 +

625 +

626 +class NanoNemotronBaseVLMultiModalProcessor(BaseMultiModalProcessor[_I]):

627 + """Basic image-only MultiModalProcessor for InternVL-style models."""

628 +

629 + def _call_hf_processor(

630 + self,

631 + prompt: str,

632 + mm_data: Mapping[str, object],

633 + mm_kwargs: Mapping[str, object],

634 + tok_kwargs: Mapping[str, object],

635 + ) -> Mapping[str, NestedTensors]:

636 + processed_outputs = super()._call_hf_processor(

637 + prompt=prompt,

638 + mm_data=mm_data,

639 + mm_kwargs=mm_kwargs,

640 + tok_kwargs=tok_kwargs,

641 + )

642 +

643 + hf_processor = self.info.get_hf_processor(**mm_kwargs)

644 + image_token_id = hf_processor.image_token_id

645 +

646 + # Since there may be extra tokens in the feature placeholders,

647 + # we need to pass the image token ID to the model to select the

648 + # tokens to merge from the vision encoder outputs

649 + processed_outputs["image_token_id"] = torch.tensor(image_token_id)

650 +

651 + return processed_outputs

652 +

653 + def _get_mm_fields_config(

654 + self,

655 + hf_inputs: Mapping[str, NestedTensors],

656 + hf_processor_mm_kwargs: Mapping[str, object],

657 + ) -> Mapping[str, MultiModalFieldConfig]:

658 + image_num_patches = hf_inputs.get("image_num_patches", torch.empty(0))

659 + num_images = len(image_num_patches)

660 +

661 + return dict(

662 + pixel_values_flat=MultiModalFieldConfig.flat_from_sizes(

663 + "image", image_num_patches),

664 + image_num_patches=MultiModalFieldConfig.batched("image"),

665 + image_embeds=MultiModalFieldConfig.batched("image"),

666 + image_token_id=MultiModalFieldConfig.shared("image", num_images),

667 + )

668 +

669 + def _get_prompt_updates(

670 + self,

671 + mm_items: MultiModalDataItems,

672 + hf_processor_mm_kwargs: Mapping[str, object],

673 + out_mm_kwargs: MultiModalKwargs,

674 + ) -> Sequence[PromptUpdate]:

675 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

676 +

677 + if "image_num_patches" in out_mm_kwargs:

678 + image_num_patches = out_mm_kwargs["image_num_patches"]

679 + assert isinstance(image_num_patches, torch.Tensor)

680 + image_num_patches = image_num_patches.tolist()

681 + elif "image_embeds" in out_mm_kwargs:

682 + # to compute num_patches (similar to Qwen2-VL)

683 + image_num_patches = [None] * len(out_mm_kwargs["image_embeds"])

684 + else:

685 + image_num_patches = []

686 +

687 + def get_replacement_custom(item_idx: int):

688 + images = mm_items.get_items(

689 + "image", (ImageEmbeddingItems, ImageProcessorItems))

690 +

691 + if isinstance(images, ImageEmbeddingItems):

692 + feature_size = images.get_feature_size(item_idx)

693 + else:

694 + image_size = images.get_image_size(item_idx)

695 + # Extract max_num_tiles from kwargs, default to 12

696 + max_num_tiles = hf_processor_mm_kwargs.get("max_num_tiles", 12)

697 + feature_size = self.info.get_num_image_tokens(

698 + image_width=image_size.width,

699 + image_height=image_size.height,

700 + max_num_tiles=max_num_tiles,

701 + processor=hf_processor,

702 + )

703 +

704 + num_patches = None

705 + local_image_num_patches = image_num_patches

706 + if isinstance(local_image_num_patches, torch.Tensor):

707 + local_image_num_patches = local_image_num_patches.tolist()

708 + if isinstance(

709 + local_image_num_patches,

710 + (list, tuple)) and item_idx < len(local_image_num_patches):

711 + num_patches = int(local_image_num_patches[item_idx])

712 +

713 + return hf_processor.get_image_repl(feature_size, num_patches)

714 +

715 + return [

716 + PromptReplacement(

717 + modality="image",

718 + target="<image>",

719 + replacement=get_replacement_custom,

720 + )

721 + ]

722 +

723 +

724 +class NanoNemotronVLMultiModalProcessor(

725 + NanoNemotronBaseVLMultiModalProcessor[NanoNemotronVLProcessingInfo]):

726 + """MultiModalProcessor extended for video support"""

727 +

728 + def _call_hf_processor(

729 + self,

730 + prompt: str,

731 + mm_data: Mapping[str, object],

732 + mm_kwargs: Mapping[str, object],

733 + tok_kwargs: Mapping[str, object],

734 + ) -> Mapping[str, NestedTensors]:

735 + processed_outputs = super()._call_hf_processor(prompt, mm_data,

736 + mm_kwargs, tok_kwargs)

737 +

738 + hf_processor = self.info.get_hf_processor(**mm_kwargs)

739 + if self.info.supports_video and (

740 + video_token_id := hf_processor.video_token_id) is not None:

741 + processed_outputs["video_token_id"] = torch.tensor(video_token_id)

742 + return processed_outputs

743 +

744 + def _get_mm_fields_config(

745 + self,

746 + hf_inputs: Mapping[str, NestedTensors],

747 + hf_processor_mm_kwargs: Mapping[str, object],

748 + ) -> Mapping[str, MultiModalFieldConfig]:

749 + image_fields = super()._get_mm_fields_config(hf_inputs,

750 + hf_processor_mm_kwargs)

751 + if self.info.supports_video:

752 + video_num_patches = hf_inputs.get("video_num_patches",

753 + torch.empty(0))

754 + num_videos = len(video_num_patches)

755 + video_fields = dict(

756 + pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes(

757 + "video", video_num_patches),

758 + video_num_patches=MultiModalFieldConfig.batched("video"),

759 + video_token_id=MultiModalFieldConfig.shared(

760 + "video", num_videos))

761 + else:

762 + video_fields = {}

763 +

764 + return image_fields | video_fields

765 +

766 + def _get_prompt_updates(

767 + self,

768 + mm_items: MultiModalDataItems,

769 + hf_processor_mm_kwargs: Mapping[str, object],

770 + out_mm_kwargs: MultiModalKwargsItems,

771 + ) -> Sequence[PromptUpdate]:

772 + prompt_repl = super()._get_prompt_updates(

773 + mm_items=mm_items,

774 + hf_processor_mm_kwargs=hf_processor_mm_kwargs,

775 + out_mm_kwargs=out_mm_kwargs,

776 + )

777 +

778 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

779 +

780 + out_mm_data = out_mm_kwargs.get_data()

781 + if "video_num_patches" in out_mm_data:

782 + video_num_patches = out_mm_data["video_num_patches"]

783 + assert isinstance(video_num_patches, torch.Tensor)

784 + video_num_patches = video_num_patches.tolist()

785 + else:

786 + video_num_patches = []

787 +

788 + def get_video_replacement_internvl(item_idx: int):

789 + feature_size = hf_processor.num_image_token

790 + num_patches = video_num_patches[item_idx]

791 + if num_patches is not None:

792 + assert isinstance(num_patches, int)

793 +

794 + return hf_processor.get_video_repl(

795 + feature_size,

796 + num_patches,

797 + video_context_token=hf_processor.video_token)

798 +

799 + if self.info.supports_video:

800 + prompt_repl = [

801 + *prompt_repl,

802 + PromptReplacement(

803 + modality="video",

804 + target="<video>",

805 + replacement=get_video_replacement_internvl,

806 + )

807 + ]

808 +

809 + return prompt_repl

810 +

811 +

812 +class NanoNemotronVLDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

813 + """Basic image-only DummyInputsBuilder for InternVL-style models."""

814 +

815 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

816 + num_images = mm_counts.get("image", 0)

817 +

818 + return "<image>" * num_images

819 +

820 + def get_dummy_mm_data(

821 + self,

822 + seq_len: int,

823 + mm_counts: Mapping[str, int],

824 + ) -> MultiModalDataDict:

825 + # Use default max_num_tiles for dummy data generation

826 + max_num_tiles = 12

827 + target_width, target_height = (

828 + self.info.get_image_size_with_most_features(max_num_tiles))

829 + num_images = mm_counts.get("image", 0)

830 +

831 + return {

832 + "image":

833 + self._get_dummy_images(width=target_width,

834 + height=target_height,

835 + num_images=num_images)

836 + }

837 +

838 +

839 +class NanoNemotronVLDummyInputsBuilder(

840 + NanoNemotronVLDummyInputsBuilder[NanoNemotronVLProcessingInfo]):

841 + """DummyInputsBuilder extended for video support"""

842 +

843 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

844 + num_videos = mm_counts.get("video", 0)

845 +

846 + return super().get_dummy_text(mm_counts) + "<video>" * num_videos

847 +

848 + def get_dummy_mm_data(

849 + self,

850 + seq_len: int,

851 + mm_counts: Mapping[str, int],

852 + ) -> MultiModalDataDict:

853 + dummy_image = super().get_dummy_mm_data(seq_len=seq_len,

854 + mm_counts=mm_counts)

855 + if self.info.supports_video:

856 + config = self.info.get_hf_config()

857 + image_size: int = config.force_image_size

858 + target_num_frames = \

859 + self.info.get_num_frames_with_most_features(seq_len, mm_counts)

860 + num_videos = mm_counts.get("video", 0)

861 + dummy_video = {

862 + "video":

863 + self._get_dummy_videos(width=image_size,

864 + height=image_size,

865 + num_frames=target_num_frames,

866 + num_videos=num_videos)

867 + }

868 + else:

869 + dummy_video = {}

870 + return {**dummy_image, **dummy_video}

871 +

872 +

873 +@MULTIMODAL_REGISTRY.register_processor(

874 + NanoNemotronVLMultiModalProcessor,

875 + info=NanoNemotronVLProcessingInfo,

876 + dummy_inputs=NanoNemotronVLDummyInputsBuilder,

877 +)

878 +class NemotronH_Nano_VL(nn.Module, HasInnerState, IsHybrid,

879 + SupportsMultiModal):

880 +

881 + @classmethod

882 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

883 + if modality.startswith("image"):

884 + return "<image>"

885 + if modality.startswith("video"):

886 + return "<video>"

887 + return None

888 +

889 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

890 + super().__init__()

891 + config = vllm_config.model_config.hf_config

892 +

893 + image_size = config.force_image_size

894 + patch_size = config.patch_size

895 + self.patch_size = patch_size

896 + self.template = config.template

897 + self.num_image_token = int(

898 + (image_size // patch_size)**2 * (config.downsample_ratio**2))

899 + self.downsample_ratio = config.downsample_ratio

900 + self.ps_version = config.ps_version

901 + self.image_tag_type = config.image_tag_type

902 +

903 + self.language_model = init_vllm_registered_model(

904 + vllm_config=vllm_config,

905 + hf_config=config.text_config,

906 + prefix=maybe_prefix(prefix, "language_model"),

907 + )

908 + self.vision_model = AutoModel.from_config(config.vision_config,

909 + trust_remote_code=True)

910 + self.vision_model.model._initialize_weights = (

911 + self.vision_model.model._init_weights)

912 + # Move input normalization to processor to mirror original HF

913 + # implementation where normalization is done in fp32

914 + self.vision_model.radio_model.make_preprocessor_external()

915 + self.vision_model = self.vision_model.to(

916 + self.language_model.config.torch_dtype)

917 +

918 + self.drop_vision_class_token = True

919 +

920 + # Construct the vision projection.

921 + vit_hidden_size = config.vit_hidden_size

922 + vision_projection_hidden_size = config.projector_hidden_size

923 + llm_hidden_size = config.text_config.hidden_size

924 +

925 + self.mlp1 = nn.Sequential(

926 + RMSNorm(hidden_size=vit_hidden_size *

927 + int(1 / self.downsample_ratio)**2,

928 + eps=1e-5),

929 + nn.Linear(

930 + vit_hidden_size * int(1 / self.downsample_ratio)**2,

931 + vision_projection_hidden_size,

932 + bias=False,

933 + ),

934 + ReLUSquaredActivation(),

935 + nn.Linear(vision_projection_hidden_size,

936 + llm_hidden_size,

937 + bias=False),

938 + )

939 + self.mlp1 = self.mlp1.to(self.language_model.config.torch_dtype)

940 +

941 + self.img_context_token_id = None

942 + self.video_context_token_id = None

943 + self.config = config

944 +

945 + def pixel_shuffle(self, x, scale_factor=0.5):

946 + n, w, h, c = x.size()

947 + # N, W, H, C --> N, W, H * scale, C // scale

948 + x = x.view(

949 + n,

950 + w,

951 + int(h * scale_factor),

952 + int(c / scale_factor),

953 + )

954 + # N, W, H * scale, C // scale --> N, H * scale, W, C // scale

955 + x = x.permute(0, 2, 1, 3).contiguous()

956 + # N, H * scale, W, C // scale -->

957 + # N, H * scale, W * scale, C // (scale ** 2)

958 + x = x.view(

959 + n,

960 + int(h * scale_factor),

961 + int(w * scale_factor),

962 + int(c / (scale_factor * scale_factor)),

963 + )

964 + if self.ps_version == "v1":

965 + warnings.warn(

966 + "In ps_version 'v1', the height and width have not "

967 + "been swapped back, which results in a transposed image.",

968 + stacklevel=2,

969 + )

970 + else:

971 + x = x.permute(0, 2, 1, 3).contiguous()

972 + return x

973 +

974 + def extract_feature(self, pixel_values):

975 + vit_embeds = self.vision_model(pixel_values).features

976 + vit_embeds = vit_embeds.to(dtype=torch.bfloat16)

977 + h = w = int(vit_embeds.shape[1]**0.5)

978 + vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)

979 + vit_embeds = self.pixel_shuffle(vit_embeds,

980 + scale_factor=self.downsample_ratio)

981 + vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1,

982 + vit_embeds.shape[-1])

983 + vit_embeds = self.mlp1(vit_embeds)

984 + return vit_embeds

985 +

986 + def _parse_and_validate_image_input(

987 + self, **kwargs: object) -> Optional[NanoNemotronVLImageInputs]:

988 + pixel_values_flat = kwargs.pop("pixel_values_flat", None)

989 + image_num_patches = kwargs.pop("image_num_patches", None)

990 + image_embeds = kwargs.pop("image_embeds", None)

991 +

992 + if pixel_values_flat is None and image_embeds is None:

993 + return None

994 +

995 + if image_embeds is not None:

996 + if not isinstance(image_embeds, (torch.Tensor, list)):

997 + raise ValueError("Incorrect type of image embeddings. "

998 + f"Got type: {type(image_embeds)}")

999 +

1000 + return NanoNemotronVLImageEmbeddinInputs(

1001 + type="image_embeds",

1002 + data=flatten_bn(image_embeds),

1003 + )

1004 +

1005 + image_token_id = kwargs["image_token_id"]

1006 + assert isinstance(image_token_id, torch.Tensor)

1007 + self.img_context_token_id = image_token_id.flatten().unique().item()

1008 +

1009 + if pixel_values_flat is not None:

1010 + if not isinstance(pixel_values_flat, (torch.Tensor, list)):

1011 + raise ValueError("Incorrect type of pixel values. "

1012 + f"Got type: {type(pixel_values_flat)}")

1013 +

1014 + if not isinstance(image_num_patches, (torch.Tensor, list)):

1015 + raise ValueError("Incorrect type of image_num_patches. "

1016 + f"Got type: {type(image_num_patches)}")

1017 +

1018 + pixel_values_flat = flatten_bn(pixel_values_flat, concat=True)

1019 + image_num_patches = flatten_bn(image_num_patches, concat=True)

1020 +

1021 + return NanoNemotronVLImagePixelInputs(

1022 + type="pixel_values",

1023 + pixel_values_flat=pixel_values_flat,

1024 + num_patches=image_num_patches,

1025 + )

1026 +

1027 + raise AssertionError("This line should be unreachable.")

1028 +

1029 + def _process_image_input(

1030 + self, image_input: NanoNemotronVLImageInputs) -> torch.Tensor:

1031 + if image_input["type"] == "image_embeds":

1032 + return image_input["data"]

1033 +

1034 + assert self.vision_model is not None

1035 +

1036 + image_embeds = self.extract_feature(image_input["pixel_values_flat"])

1037 + num_patches = image_input["num_patches"]

1038 +

1039 + # Only one image in the current batch

1040 + if len(num_patches) == 1:

1041 + return (image_embeds.view(-1,

1042 + self.config.text_config.hidden_size), )

1043 +

1044 + # NOTE: Image embeddings are split into separate tensors for each image

1045 + # by the size of each embedding.

1046 + feature_size = image_embeds.shape[1]

1047 + image_embeds = image_embeds.view(-1,

1048 + self.config.text_config.hidden_size)

1049 + image_feature_sizes = [

1050 + num_patches * feature_size for num_patches in num_patches

1051 + ]

1052 + return image_embeds.split(image_feature_sizes)

1053 +

1054 + def _parse_and_validate_video_input(

1055 + self,

1056 + **kwargs: object) -> Optional[NanoNemotronVLVideoPixelInputs]:

1057 + pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None)

1058 + video_num_patches = kwargs.pop("video_num_patches", None)

1059 + video_embeds = kwargs.pop("video_embeds", None)

1060 +

1061 + if pixel_values_flat_video is None and video_embeds is None:

1062 + return None

1063 +

1064 + if video_embeds is not None:

1065 + return NanoNemotronVLVideoEmbeddingInputs(

1066 + type="video_embeds",

1067 + data=flatten_bn(video_embeds),

1068 + )

1069 +

1070 + video_token_id = kwargs["video_token_id"]

1071 + assert isinstance(video_token_id, torch.Tensor)

1072 + self.video_context_token_id = video_token_id.flatten().unique().item()

1073 +

1074 + if pixel_values_flat_video is not None:

1075 + if not isinstance(pixel_values_flat_video, (torch.Tensor, list)):

1076 + raise ValueError("Incorrect type of pixel values. "

1077 + f"Got type: {type(pixel_values_flat_video)}")

1078 +

1079 + if not isinstance(video_num_patches, (torch.Tensor, list)):

1080 + raise ValueError("Incorrect type of image_num_patches. "

1081 + f"Got type: {type(video_num_patches)}")

1082 +

1083 + pixel_values_flat_video = flatten_bn(pixel_values_flat_video,

1084 + concat=True)

1085 + video_num_patches = flatten_bn(video_num_patches, concat=True)

1086 + expected_h = expected_w = self.config.force_image_size

1087 + resolve_bindings = {"h": expected_h, "w": expected_w}

1088 +

1089 + return NanoNemotronVLVideoPixelInputs(

1090 + type="pixel_values_videos",

1091 + pixel_values_flat=pixel_values_flat_video,

1092 + num_patches=video_num_patches,

1093 + resolve_bindings=resolve_bindings,

1094 + )

1095 +

1096 + raise AssertionError("This line should be unreachable.")

1097 +

1098 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1099 + modalities = {}

1100 + # Preserve the order of modalities if there are multiple of them

1101 + # from the order of kwargs.

1102 + for input_key in kwargs:

1103 + if input_key in ("pixel_values_flat",

1104 + "image_embeds") and "images" not in modalities:

1105 + modalities["images"] = self._parse_and_validate_image_input(

1106 + **kwargs)

1107 + if input_key in ("pixel_values_flat_video",

1108 + ) and "videos" not in modalities:

1109 + modalities["videos"] = self._parse_and_validate_video_input(

1110 + **kwargs)

1111 +

1112 + return modalities

1113 +

1114 + def get_multimodal_embeddings(

1115 + self, **kwargs: object) -> Optional[MultiModalEmbeddings]:

1116 + # Validate the multimodal input keyword arguments

1117 + modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

1118 + if modalities is None:

1119 + return []

1120 +

1121 + # # The result multimodal_embeddings is tuple of tensors, with each

1122 + # tensor correspoending to a multimodal data item (image or video).

1123 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()

1124 +

1125 + # NOTE: It is important to iterate over the keys in this dictionary

1126 + # to preserve the order of the modalities.

1127 + for modality in modalities:

1128 + if modality == "images":

1129 + image_input = modalities["images"]

1130 + vision_embeddings = self._process_image_input(image_input)

1131 + multimodal_embeddings += vision_embeddings

1132 + if modality == "videos":

1133 + video_input = modalities["videos"]

1134 + video_embeddings = self._process_image_input(video_input)

1135 + multimodal_embeddings += video_embeddings

1136 +

1137 + return multimodal_embeddings

1138 +

1139 + def get_input_embeddings(

1140 + self,

1141 + input_ids: torch.Tensor,

1142 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1143 + ) -> torch.Tensor:

1144 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)

1145 +

1146 + if (multimodal_embeddings is not None

1147 + and len(multimodal_embeddings) != 0):

1148 + context_token_ids = [

1149 + token_id for token_id in (self.img_context_token_id,

1150 + self.video_context_token_id)

1151 + if token_id is not None

1152 + ]

1153 + assert len(context_token_ids) >= 1

1154 + inputs_embeds = merge_multimodal_embeddings(

1155 + input_ids,

1156 + inputs_embeds,

1157 + multimodal_embeddings,

1158 + context_token_ids,

1159 + )

1160 +

1161 + return inputs_embeds

1162 +

1163 + def get_language_model(self) -> torch.nn.Module:

1164 + return self.language_model

1165 +

1166 + def forward(

1167 + self,

1168 + input_ids: torch.Tensor,

1169 + positions: torch.Tensor,

1170 + intermediate_tensors: Optional[IntermediateTensors] = None,

1171 + inputs_embeds: Optional[torch.Tensor] = None,

1172 + **kwargs: object,

1173 + ) -> Union[torch.Tensor, IntermediateTensors]:

1174 + if intermediate_tensors is not None:

1175 + input_ids = None

1176 + inputs_embeds = None

1177 +

1178 + # NOTE: In v1, inputs_embeds is always generated at model runner, this

1179 + # condition is for v0 compatibility.

1180 + elif inputs_embeds is None:

1181 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1182 + inputs_embeds = self.get_input_embeddings(input_ids,

1183 + vision_embeddings)

1184 + input_ids = None

1185 +

1186 + hidden_states = self.language_model(

1187 + input_ids=input_ids,

1188 + positions=positions,

1189 + intermediate_tensors=intermediate_tensors,

1190 + inputs_embeds=inputs_embeds,

1191 + **kwargs,

1192 + )

1193 +

1194 + return hidden_states

1195 +

1196 + def get_mm_mapping(self) -> MultiModelKeys:

1197 + """

1198 + Get the module prefix in multimodal models

1199 + """

1200 + return MultiModelKeys.from_string_field(

1201 + language_model="language_model",

1202 + connector="mlp1",

1203 + tower_model="vision_model",

1204 + )

1205 +

1206 + def compute_logits(

1207 + self,

1208 + hidden_states: torch.Tensor,

1209 + sampling_metadata: SamplingMetadata,

1210 + ) -> Optional[torch.Tensor]:

1211 + return self.language_model.compute_logits(hidden_states,

1212 + sampling_metadata)

1213 +

1214 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

1215 +

1216 + def is_vision_model_weights(weight: tuple[str, torch.Tensor]):

1217 + return weight[0].startswith("vision_model")

1218 +

1219 + def is_adapter_weights(weight: tuple[str, torch.Tensor]):

1220 + return weight[0].startswith("mlp1")

1221 +

1222 + # Get references to parameters for direct loading

1223 + vision_model_dict = dict(self.vision_model.named_parameters())

1224 + vision_model_buffers = dict(self.vision_model.named_buffers())

1225 + adapter_dict = dict(self.mlp1.named_parameters())

1226 +

1227 + def llm_weights_generator():

1228 + # Single pass over weights

1229 + for name, w in weights:

1230 + if is_vision_model_weights((name, w)):

1231 + # Load vision encoder weights directly

1232 + trimmed_name = ".".join(name.split(".")[1:])

1233 + if "input_conditioner" in trimmed_name:

1234 + continue

1235 + if trimmed_name in vision_model_buffers:

1236 + param = vision_model_buffers[trimmed_name]

1237 + else:

1238 + param = vision_model_dict[trimmed_name]

1239 + with torch.no_grad():

1240 + default_weight_loader(param, w)

1241 + elif is_adapter_weights((name, w)):

1242 + # Load vision-language adapter weights directly

1243 + trimmed_name = ".".join(name.split(".")[1:])

1244 + param = adapter_dict[trimmed_name]

1245 + with torch.no_grad():

1246 + default_weight_loader(param, w)

1247 + else:

1248 + # LLM weights: yield them to be loaded

1249 + # by language_model.load_weights

1250 + assert name.startswith("language_model")

1251 + trimmed_name = ".".join(name.split(".")[1:])

1252 + yield (trimmed_name, w)

1253 +

1254 + # Now we call the language model load with the generator

1255 + self.language_model.load_weights(llm_weights_generator())

1256 +

1257 + def print_architecture(self,

1258 + detailed: bool = True,

1259 + save_to_file: str = None):

1260 + """

1261 + Print model architecture with parameter names, shapes, and sizes.

1262 +

1263 + Args:

1264 + detailed: If True, show detailed parameter breakdown

1265 + save_to_file: If provided, save output to this file path

1266 + """

1267 + import sys

1268 + from io import StringIO

1269 +

1270 + # Capture output if saving to file

1271 + original_stdout = sys.stdout

1272 + if save_to_file:

1273 + sys.stdout = StringIO()

1274 +

1275 + try:

1276 + print("=" * 100)

1277 + print("NemotronH_Nano_VL Model Architecture")

1278 + print("=" * 100)

1279 +

1280 + total_params = 0

1281 + param_groups = {

1282 + "language_model": [],

1283 + "vision_model": [],

1284 + "mlp1": [],

1285 + "other": [],

1286 + }

1287 +

1288 + for name, param in self.named_parameters():

1289 + param_size = param.numel()

1290 + total_params += param_size

1291 +

1292 + # Group parameters by main component

1293 + if name.startswith("language_model"):

1294 + param_groups["language_model"].append(

1295 + (name, param.shape, param_size, param.dtype))

1296 + elif name.startswith("vision_model"):

1297 + param_groups["vision_model"].append(

1298 + (name, param.shape, param_size, param.dtype))

1299 + elif name.startswith("mlp1"):

1300 + param_groups["mlp1"].append(

1301 + (name, param.shape, param_size, param.dtype))

1302 + else:

1303 + param_groups["other"].append(

1304 + (name, param.shape, param_size, param.dtype))

1305 +

1306 + if detailed:

1307 + print(f"{name:<70} | Shape: {str(param.shape):<25} | "

1308 + f"Size: {param_size:>12,} | Dtype: {param.dtype}")

1309 +

1310 + print("=" * 100)

1311 + print("Summary by Component:")

1312 + print("-" * 60)

1313 +

1314 + for component, params in param_groups.items():

1315 + if params: # Only show components that have parameters

1316 + component_total = sum(size for _, _, size, _ in params)

1317 + percentage = ((component_total / total_params) *

1318 + 100 if total_params > 0 else 0)

1319 + print(f"{component:<20} | Parameters: {len(params):>4} | "

1320 + f"Total Size: {component_total:>15,} | "

1321 + f"{percentage:>6.2f}%")

1322 +

1323 + print("-" * 60)

1324 + print(f"{'Total Parameters':<20} | {total_params:>15,}")

1325 +

1326 + # Estimate memory usage (assuming bfloat16 = 2 bytes per parameter)

1327 + memory_mb = total_params * 2 / (1024**2)

1328 + memory_gb = memory_mb / 1024

1329 + print(f"{'Est. Memory (MB)':<20} | {memory_mb:>15.2f}")

1330 + print(f"{'Est. Memory (GB)':<20} | {memory_gb:>15.2f}")

1331 + print("=" * 100)

1332 +

1333 + # Save to file if requested

1334 + if save_to_file:

1335 + output = sys.stdout.getvalue()

1336 + sys.stdout = original_stdout

1337 + with open(save_to_file, "w") as f:

1338 + f.write(output)

1339 + print(f"Architecture saved to: {save_to_file}")

1340 + print(output) # Also print to console

1341 +

1342 + finally:

1343 + if save_to_file and sys.stdout != original_stdout:

1344 + sys.stdout = original_stdout

1345 +

1346 + def get_model_info(self):

1347 + """

1348 + Get basic model information as a dictionary.

1349 + """

1350 + total_params = sum(p.numel() for p in self.parameters())

1351 +

1352 + component_info = {}

1353 + for name, param in self.named_parameters():

1354 + component = name.split(".")[0]

1355 + if component not in component_info:

1356 + component_info[component] = {"params": 0, "size": 0}

1357 + component_info[component]["params"] += 1

1358 + component_info[component]["size"] += param.numel()

1359 +

1360 + return {

1361 + "model_name": "NemotronH_Nano_VL",

1362 + "total_parameters": total_params,

1363 + "memory_estimate_mb": total_params * 2 / (1024**2), # bfloat16

1364 + "components": component_info,

1365 + "config": {

1366 + "image_size": getattr(self.config, "force_image_size", None),

1367 + "patch_size": getattr(self.config, "patch_size", None),

1368 + "num_image_token": self.num_image_token,

1369 + "downsample_ratio": self.downsample_ratio,

1370 + },

1371 + }

1372 +

1373 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

1374 + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(

1375 + input_buffers, **kwargs)

1376 +

1377 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

1378 + return (self.language_model.mamba_cache.

1379 + get_seqlen_agnostic_capture_inputs(batch_size))

1380 +

1381 + @classmethod

1382 + def get_mamba_state_shape_from_config(cls, vllm_config: "VllmConfig"):

1383 + text_config = vllm_config.model_config.hf_config.text_config

1384 + temp_vllm_config = copy.deepcopy(vllm_config)

1385 + temp_vllm_config.model_config.hf_config = text_config

1386 + return NemotronHForCausalLM.get_mamba_state_shape_from_config(

1387 + temp_vllm_config)

1388 +

1389 + @classmethod

1390 + def get_mamba_state_dtype_from_config(cls, vllm_config: "VllmConfig"):

1391 + text_config = vllm_config.model_config.hf_config.text_config

1392 + temp_vllm_config = copy.deepcopy(vllm_config)

1393 + temp_vllm_config.model_config.hf_config = text_config

1394 + return NemotronHForCausalLM.get_mamba_state_dtype_from_config(

1395 + temp_vllm_config)

Support for NemotronH Nano VLM (#23644)