diff --git a/docs/supported_models/multimodal_language_models.md b/docs/supported_models/multimodal_language_models.md index dfcde4e55..68faa423a 100644 --- a/docs/supported_models/multimodal_language_models.md +++ b/docs/supported_models/multimodal_language_models.md @@ -45,6 +45,7 @@ in the GitHub search bar. | **DotsVLM** (General/OCR) | `rednote-hilab/dots.vlm1.inst` | RedNote's vision-language model built on a 1.2B vision encoder and DeepSeek V3 LLM, featuring NaViT vision encoder trained from scratch with dynamic resolution support and enhanced OCR capabilities through structured image data training. | | | **DotsVLM-OCR** | `rednote-hilab/dots.ocr` | Specialized OCR variant of DotsVLM optimized for optical character recognition tasks with enhanced text extraction and document understanding capabilities. | Don't use `--trust-remote-code` | | **NVILA** (8B, 15B, Lite-2B, Lite-8B, Lite-15B) | `Efficient-Large-Model/NVILA-8B` | `chatml` | NVILA explores the full stack efficiency of multi-modal design, achieving cheaper training, faster deployment and better performance. | +| **JetVLM** | | JetVLM is an vision-language model designed for high-performance multimodal understanding and generation tasks built upon Jet-Nemotron. | Coming soon | ## Video Input Support @@ -56,6 +57,7 @@ SGLang supports video input for Vision-Language Models (VLMs), enabling temporal | **GLM-4v** (4.5V, 4.1V, MOE) | `zai-org/GLM-4.5V` | Video clips are read with Decord, converted to tensors, and passed to the model alongside metadata for rotary-position handling. | | **NVILA** (Full & Lite) | `Efficient-Large-Model/NVILA-8B` | The runtime samples eight frames per clip and attaches them to the multimodal request when `video_data` is present. | | **LLaVA video variants** (LLaVA-NeXT-Video, LLaVA-OneVision) | `lmms-lab/LLaVA-NeXT-Video-7B` | The processor routes video prompts to the LlavaVid video-enabled architecture, and the provided example shows how to query it with `sgl.video(...)` clips. | +| **JetVLM** | | The runtime samples eight frames per clip and attaches them to the multimodal request when `video_data` is present. | Use `sgl.video(path, num_frames)` when building prompts to attach clips from your SGLang programs. diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 690a1e3eb..623d91d73 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -7,6 +7,7 @@ from sglang.srt.configs.exaone import ExaoneConfig from sglang.srt.configs.falcon_h1 import FalconH1Config from sglang.srt.configs.janus_pro import MultiModalityConfig from sglang.srt.configs.jet_nemotron import JetNemotronConfig +from sglang.srt.configs.jet_vlm import JetVLMConfig from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.configs.kimi_vl import KimiVLConfig from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig @@ -40,4 +41,5 @@ __all__ = [ "FalconH1Config", "NemotronHConfig", "JetNemotronConfig", + "JetVLMConfig", ] diff --git a/python/sglang/srt/configs/jet_vlm.py b/python/sglang/srt/configs/jet_vlm.py new file mode 100644 index 000000000..9b8cba6e1 --- /dev/null +++ b/python/sglang/srt/configs/jet_vlm.py @@ -0,0 +1,53 @@ +from typing import Any + +from transformers.configuration_utils import PretrainedConfig +from transformers.models.siglip import SiglipVisionConfig + +from sglang.srt.configs.jet_nemotron import JetNemotronConfig +from sglang.srt.configs.mamba_utils import Mamba2CacheParams + + +class JetVLMConfig(PretrainedConfig): + model_type = "jet_vlm" + sub_configs = { + "text_config": JetNemotronConfig, + "vision_config": SiglipVisionConfig, + } + _auto_class = "AutoConfig" + + def __init__( + self, + *, + text_config: dict[str, Any] | None = None, + vision_config: dict[str, Any] | None = None, + image_token_id: int | None = None, + video_token_id: int | None = None, + **kwargs, + ): + self.text_config = ( + JetNemotronConfig(**text_config) + if text_config is not None + else JetNemotronConfig() + ) + self.vision_config = ( + SiglipVisionConfig(**vision_config) + if vision_config is not None + else SiglipVisionConfig() + ) + + self.image_token_id = image_token_id if image_token_id is not None else -1 + self.video_token_id = video_token_id if video_token_id is not None else -1 + + super().__init__(**kwargs) + + @property + def full_attention_layer_ids(self) -> list[int]: + return self.text_config.full_attention_layer_ids + + @property + def linear_layer_ids(self) -> list[int]: + return self.text_config.linear_layer_ids + + @property + def mamba2_cache_params(self) -> Mamba2CacheParams: + return self.text_config.mamba2_cache_params diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 2c12ad0f3..77c2d81cd 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -955,6 +955,7 @@ multimodal_model_archs = [ "NVILAForConditionalGeneration", "NVILALiteForConditionalGeneration", "DeepseekOCRForCausalLM", + "JetVLMForConditionalGeneration", ] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 76226cca9..950d3acfb 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -32,6 +32,7 @@ import torch.distributed as dist from sglang.srt.configs import ( FalconH1Config, JetNemotronConfig, + JetVLMConfig, KimiLinearConfig, NemotronHConfig, Qwen3NextConfig, @@ -1416,7 +1417,7 @@ class ModelRunner: @property def hybrid_gdn_config(self): config = self.model_config.hf_config - if isinstance(config, Qwen3NextConfig | JetNemotronConfig): + if isinstance(config, Qwen3NextConfig | JetNemotronConfig | JetVLMConfig): return config return None diff --git a/python/sglang/srt/models/jet_nemotron.py b/python/sglang/srt/models/jet_nemotron.py index 918cdc65b..513f2ce37 100644 --- a/python/sglang/srt/models/jet_nemotron.py +++ b/python/sglang/srt/models/jet_nemotron.py @@ -545,6 +545,9 @@ class JetNemotronForCausalLM(nn.Module): else: return self.pooler(hidden_states, forward_batch) + def get_input_embeddings(self) -> nn.Module: + return self.model.embed_tokens + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): stacked_params_mapping: list[tuple[str, str, str | int]] = [ # (param_name, shard_weight_name, shard_id) diff --git a/python/sglang/srt/models/jet_vlm.py b/python/sglang/srt/models/jet_vlm.py new file mode 100644 index 000000000..fc7c6b6c1 --- /dev/null +++ b/python/sglang/srt/models/jet_vlm.py @@ -0,0 +1,143 @@ +import math +from collections.abc import Iterable + +import einops +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from transformers.modeling_outputs import BaseModelOutputWithPooling +from transformers.models.siglip import SiglipVisionModel + +import sglang.srt.managers.mm_utils as mm_utils +import sglang.srt.model_loader.weight_utils as weight_utils +import sglang.srt.utils as utils +from sglang.srt.configs.jet_vlm import JetVLMConfig +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.jet_nemotron import JetNemotronForCausalLM + +MM_HIDDEN_SIZE = 1152 + + +class JetVLMDownSample2x2BlockFix(nn.Module): + def forward(self, x: Tensor) -> Tensor: + _, seq_len, _ = x.shape + + feat_size = math.isqrt(seq_len) + + features = einops.rearrange(x, "b (h w) d -> b h w d", h=feat_size, w=feat_size) + + if feat_size % 2 == 1: + features = F.pad(features, (0, 0, 0, 1, 0, 1)) + + features = einops.rearrange( + features, "b (h p1) (w p2) d -> b (h w) (p1 p2 d)", p1=2, p2=2 + ) + + return features + + +class JetVLMMultiModalProjector(nn.Module): + def __init__(self, config: JetVLMConfig) -> None: + super().__init__() + + self.layers = nn.Sequential( + JetVLMDownSample2x2BlockFix(), + nn.LayerNorm(MM_HIDDEN_SIZE * 4), + nn.Linear(MM_HIDDEN_SIZE * 4, config.text_config.hidden_size), + nn.GELU(), + nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.layers(x) + + +class JetVLMForConditionalGeneration(nn.Module): + def __init__( + self, + config: JetVLMConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + + self.vision_tower = SiglipVisionModel(config.vision_config) + self.mm_projector = JetVLMMultiModalProjector(config) + self.llm = JetNemotronForCausalLM( + config=config.text_config, + quant_config=quant_config, + prefix=utils.add_prefix("llm", prefix), + ) + + def forward( + self, + input_ids: Tensor, + positions: Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + ) -> LogitsProcessorOutput: + output = mm_utils.general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.llm, + data_embedding_funcs={ + Modality.IMAGE: self.get_image_feature, + Modality.VIDEO: self.get_image_feature, + }, + get_embedding=get_embedding, + positions=positions, + ) + + assert isinstance(output, LogitsProcessorOutput) + + return output + + def get_image_feature(self, mm_input: list[MultimodalDataItem]) -> Tensor: + pixel_values = torch.cat([torch.tensor(x.feature) for x in mm_input], dim=0) + + vision_tower_output: BaseModelOutputWithPooling = self.vision_tower( + pixel_values, + output_hidden_states=True, + ) + assert vision_tower_output.hidden_states is not None + + vision_features = vision_tower_output.hidden_states[-2] + + vision_features = self.mm_projector(vision_features) + + vision_features = einops.rearrange(vision_features, "n p d -> (n p) d") + + return vision_features + + def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> None: + params_dict = dict(self.named_parameters()) + + for name, loaded_weight in weights: + if name.startswith("llm."): + self.llm.load_weights([(name[len("llm.") :], loaded_weight)]) + else: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", weight_utils.default_weight_loader + ) + weight_loader(param, loaded_weight) + + def pad_input_ids( + self, input_ids: list[int], mm_inputs: MultimodalInputs + ) -> list[int]: + pattern = MultiModalityDataPaddingPatternMultimodalTokens() + return pattern.pad_input_tokens(input_ids, mm_inputs) + + +EntryClass = [JetVLMForConditionalGeneration] diff --git a/python/sglang/srt/multimodal/processors/nvila.py b/python/sglang/srt/multimodal/processors/nvila.py index fde974862..f34d600b3 100644 --- a/python/sglang/srt/multimodal/processors/nvila.py +++ b/python/sglang/srt/multimodal/processors/nvila.py @@ -6,6 +6,7 @@ from transformers.processing_utils import ProcessorMixin from transformers.tokenization_utils_base import PreTrainedTokenizerBase from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.models.jet_vlm import JetVLMForConditionalGeneration from sglang.srt.models.nvila import NVILAForConditionalGeneration from sglang.srt.models.nvila_lite import NVILALiteForConditionalGeneration from sglang.srt.multimodal.processors.base_processor import ( @@ -21,6 +22,7 @@ class NVILAMultimodalProcessor(BaseMultimodalProcessor): models: list[type[nn.Module]] = [ NVILAForConditionalGeneration, NVILALiteForConditionalGeneration, + JetVLMForConditionalGeneration, ] def __init__( diff --git a/python/sglang/srt/utils/hf_transformers_utils.py b/python/sglang/srt/utils/hf_transformers_utils.py index fbc34e906..aa1c46820 100644 --- a/python/sglang/srt/utils/hf_transformers_utils.py +++ b/python/sglang/srt/utils/hf_transformers_utils.py @@ -45,6 +45,7 @@ from sglang.srt.configs import ( ExaoneConfig, FalconH1Config, JetNemotronConfig, + JetVLMConfig, KimiLinearConfig, KimiVLConfig, LongcatFlashConfig, @@ -79,6 +80,7 @@ _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [ NemotronHConfig, DeepseekVLV2Config, JetNemotronConfig, + JetVLMConfig, ] _CONFIG_REGISTRY = {