From 0769de9b0f8ce031eb844dda940d8738871e0e8a Mon Sep 17 00:00:00 2001 From: Shivam jindal <73436052+yes-its-shivam@users.noreply.github.com> Date: Thu, 29 Jan 2026 20:33:41 +0530 Subject: [PATCH] Support LightOnOCR-2-1B (#17806) --- python/sglang/srt/configs/model_config.py | 1 + python/sglang/srt/models/lightonocr.py | 298 ++++++++++++++++++ .../srt/multimodal/processors/lightonocr.py | 110 +++++++ 3 files changed, 409 insertions(+) create mode 100644 python/sglang/srt/models/lightonocr.py create mode 100644 python/sglang/srt/multimodal/processors/lightonocr.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 359aa811e..68d0388ce 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1147,6 +1147,7 @@ multimodal_model_archs = [ "LlavaQwenForCausalLM", "LlavaForConditionalGeneration", "LlavaVidForCausalLM", + "LightOnOCRForConditionalGeneration", "MiniCPMO", "MiniCPMV", "Mistral3ForConditionalGeneration", diff --git a/python/sglang/srt/models/lightonocr.py b/python/sglang/srt/models/lightonocr.py new file mode 100644 index 000000000..fe854f603 --- /dev/null +++ b/python/sglang/srt/models/lightonocr.py @@ -0,0 +1,298 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +""" +Support for lightonai/LightOnOCR-2-1B. + +LightOnOCR is a vision-language OCR model that combines: +- Pixtral vision encoder (24 layers, 1024 hidden dim) +- Spatial merge projection with RMSNorm + PatchMerger (2x2 = 4x token reduction) +- Qwen3 language decoder (28 layers, 1024 hidden dim) + +Key differences from PixtralForConditionalGeneration: +- Uses Qwen3ForCausalLM instead of MistralLarge3ForCausalLM as the language model +- Has an RMSNorm applied to vision encoder output before patch merging +- Does not use image break/end tokens (single contiguous image token range) +- HuggingFace checkpoint uses a vision_projection namespace for norm, patch_merger, + and adapter weights + +References: +- https://huggingface.co/lightonai/LightOnOCR-2-1B +""" + +from dataclasses import fields +from typing import Iterable, List, Tuple + +import torch +import torch.nn as nn + +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.pixtral import ( + PATCH_MERGE, + PatchMerger, + PixtralHFVisionModel, + VisionEncoderArgs, + VisionLanguageAdapter, +) +from sglang.srt.models.qwen3 import Qwen3ForCausalLM + + +class LightOnOCRForConditionalGeneration(nn.Module): + """ + LightOnOCR model for SGLang inference. + + Architecture: + - Pixtral-based vision encoder (PixtralHFVisionModel, 24 layers) + - RMSNorm on vision encoder output + - Spatial merge via PatchMerger (2x2 = 4x token reduction) + - VisionLanguageAdapter projection to text hidden size + - Qwen3-based decoder (28 layers) with QK norms + """ + + merge_by_field_config = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return None + raise ValueError("Only image modality is supported") + + def __init__(self, *, config, prefix: str = "", **kwargs): + super().__init__() + self.config = config + quant_config = kwargs.get("quant_config") + + # Build VisionEncoderArgs from config + vision_config = config.vision_config + dataclass_fields = {field.name for field in fields(VisionEncoderArgs)} + vision_args = { + key: value + for key, value in vision_config.to_dict().items() + if key in dataclass_fields + } + # LightOnOCR stores these at the top-level config + if "image_token_id" not in vision_args: + vision_args["image_token_id"] = getattr(config, "image_token_id", 151655) + if "spatial_merge_size" not in vision_args: + vision_args["spatial_merge_size"] = getattr(config, "spatial_merge_size", 2) + if "adapter_bias" not in vision_args: + vision_args["adapter_bias"] = getattr( + config, "multimodal_projector_bias", True + ) + # LightOnOCR uses patch merging for spatial merge + vision_args["mm_projector_id"] = PATCH_MERGE + self.vision_args = VisionEncoderArgs(**vision_args) + + # Vision encoder (Pixtral HF variant with SGLang parallel layers) + self.vision_encoder = PixtralHFVisionModel(vision_config, quant_config=None) + + # RMSNorm applied to vision encoder output before patch merging + self.vision_projection_norm = RMSNorm(self.vision_args.hidden_size, eps=1e-5) + + # Patch merger for spatial token reduction + self.patch_merger = PatchMerger( + vision_encoder_dim=self.vision_args.hidden_size, + spatial_merge_size=self.vision_args.spatial_merge_size, + ) + + # Vision-to-language projection adapter + self.vision_language_adapter = VisionLanguageAdapter( + self.vision_args, dim=config.text_config.hidden_size + ) + + # Language model + self.language_model = Qwen3ForCausalLM( + config=config.text_config, + quant_config=quant_config, + ) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + pattern = MultiModalityDataPaddingPatternMultimodalTokens() + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + """Process images through vision encoder and projection pipeline.""" + images = [item.feature for item in items] + + # Extract image sizes from model-specific data or infer from tensor shape + image_sizes_list = [] + for item in items: + if item.model_specific_data and "image_sizes" in item.model_specific_data: + sizes_tensor = item.model_specific_data["image_sizes"] + for size in sizes_tensor: + image_sizes_list.append((int(size[0]), int(size[1]))) + else: + img = item.feature + for _ in range(img.shape[0]): + image_sizes_list.append((img.shape[-2], img.shape[-1])) + + # Stack pixel values + if len(images) > 1: + pixel_values = torch.cat(images, dim=0) + else: + pixel_values = images[0] + + # Vision encoder forward + image_features = self.vision_encoder(pixel_values, image_sizes=image_sizes_list) + image_features = image_features.view(-1, image_features.shape[-1]) + + # Norm before patch merge (matches HF Mistral3MultiModalProjector order) + image_features = self.vision_projection_norm(image_features) + + # Spatial merge via patch merger — use actual image sizes (not padded tensor + # shape) because PixtralHFVisionModel crops embeddings to real dimensions. + patch_size = self.vision_args.patch_size + img_patch_dims = [ + (h // patch_size, w // patch_size) for (h, w) in image_sizes_list + ] + image_features = self.patch_merger(image_features, image_sizes=img_patch_dims) + + # Project to language model dimension + image_embeds = self.vision_language_adapter(image_features) + return image_embeds + + def get_language_model(self) -> torch.nn.Module: + return self.language_model + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ): + return general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.language_model, + multimodal_model=self, + positions=positions, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def get_embed_and_head(self): + return self.language_model.get_embed_and_head() + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + """Load weights from HuggingFace checkpoint. + + HF checkpoint weight layout (after stripping ``model.`` prefix): + - ``vision_encoder.*`` -> self.vision_encoder + - ``vision_projection.norm.*`` -> self.vision_projection_norm + - ``vision_projection.patch_merger.*`` -> self.patch_merger + - ``vision_projection.linear_1.*`` -> self.vision_language_adapter.w_in + - ``vision_projection.linear_2.*`` -> self.vision_language_adapter.w_out + - ``language_model.*`` -> self.language_model (Qwen3ForCausalLM) + """ + vision_encoder_dict = dict(self.vision_encoder.named_parameters()) + patch_merger_dict = dict(self.patch_merger.named_parameters()) + norm_dict = dict(self.vision_projection_norm.named_parameters()) + adapter_dict = dict(self.vision_language_adapter.named_parameters()) + + # PixtralHFVisionModel uses SGLang parallel layers with stacked params + stacked_params_mapping = [ + (".attention.qkv_proj", ".attention.q_proj", "q"), + (".attention.qkv_proj", ".attention.k_proj", "k"), + (".attention.qkv_proj", ".attention.v_proj", "v"), + (".feed_forward.gate_up_proj", ".feed_forward.gate_proj", 0), + (".feed_forward.gate_up_proj", ".feed_forward.up_proj", 1), + ] + + def llm_weights_generator(): + for name, w in weights: + # HF checkpoint prefixes all weights with model. + if name.startswith("model."): + name = name[len("model.") :] + + if name.startswith("vision_encoder."): + trimmed = name[len("vision_encoder.") :] + + # Handle stacked params (QKV, gate/up) + loaded = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name in trimmed: + transformed = trimmed.replace(weight_name, param_name) + if transformed in vision_encoder_dict: + param = vision_encoder_dict[transformed] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + with torch.no_grad(): + weight_loader(param, w, shard_id) + loaded = True + break + + if not loaded: + # Handle o_proj -> proj rename + if ".attention.o_proj" in trimmed: + trimmed = trimmed.replace( + ".attention.o_proj", ".attention.proj" + ) + if trimmed in vision_encoder_dict: + param = vision_encoder_dict[trimmed] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + with torch.no_grad(): + weight_loader(param, w) + + elif name.startswith("vision_projection."): + remaining = name[len("vision_projection.") :] + + if remaining.startswith("patch_merger."): + trimmed = remaining[len("patch_merger.") :] + if trimmed in patch_merger_dict: + param = patch_merger_dict[trimmed] + with torch.no_grad(): + default_weight_loader(param, w) + + elif remaining.startswith("norm."): + trimmed = remaining[len("norm.") :] + if trimmed in norm_dict: + param = norm_dict[trimmed] + with torch.no_grad(): + default_weight_loader(param, w) + + else: + # linear_1 -> w_in, linear_2 -> w_out + trimmed = remaining.replace("linear_1.", "w_in.").replace( + "linear_2.", "w_out." + ) + if trimmed in adapter_dict: + param = adapter_dict[trimmed] + with torch.no_grad(): + default_weight_loader(param, w) + + else: + # Language model weights and any other weights + if name.startswith("language_model."): + # Qwen3ForCausalLM expects model.* prefix + name = "model." + name[len("language_model.") :] + yield (name, w) + + self.language_model.load_weights(llm_weights_generator()) + + +EntryClass = LightOnOCRForConditionalGeneration diff --git a/python/sglang/srt/multimodal/processors/lightonocr.py b/python/sglang/srt/multimodal/processors/lightonocr.py new file mode 100644 index 000000000..59c6c1429 --- /dev/null +++ b/python/sglang/srt/multimodal/processors/lightonocr.py @@ -0,0 +1,110 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +""" +Multimodal processor for lightonai/LightOnOCR-2-1B. + +Key difference from Pixtral: LightOnOCR does NOT use image break/end tokens. +The parent PixtralProcessor inserts row-break and image-end tokens between +image patch rows. This processor removes them after the parent processing +to produce a single contiguous range of image tokens per image. +""" + +from typing import List, Union + +from sglang.srt.models.lightonocr import LightOnOCRForConditionalGeneration +from sglang.srt.multimodal.processors.pixtral import PixtralProcessor + + +class LightOnOCRProcessor(PixtralProcessor): + """Processor for LightOnOCR model.""" + + models = [LightOnOCRForConditionalGeneration] + + def __init__(self, hf_config, server_args, _processor, *args, **kwargs): + # LightOnOCR uses image_token_id instead of image_token_index + if not hasattr(hf_config, "image_token_index"): + hf_config.image_token_index = getattr(hf_config, "image_token_id", 151655) + + # Propagate spatial_merge_size from root config to vision_config + spatial_merge_size = getattr(hf_config, "spatial_merge_size", 2) + if hasattr(hf_config, "vision_config"): + vc = hf_config.vision_config + if not hasattr(vc, "spatial_merge_size") or vc.spatial_merge_size is None: + vc.spatial_merge_size = spatial_merge_size + + if hasattr(_processor, "patch_size"): + _processor.spatial_merge_size = spatial_merge_size + + super().__init__(hf_config, server_args, _processor, *args, **kwargs) + + # Identify break/end token IDs for removal + self._break_token_ids = set() + for attr in ("image_break_token_id", "image_break_id"): + tid = getattr(_processor, attr, None) + if tid is not None: + self._break_token_ids.add(tid) + for attr in ("image_end_token_id", "image_end_id"): + tid = getattr(_processor, attr, None) + if tid is not None: + self._break_token_ids.add(tid) + + async def process_mm_data_async( + self, + image_data: List[Union[str, bytes]], + input_text, + request_obj, + *args, + **kwargs, + ): + result = await super().process_mm_data_async( + image_data=image_data, + input_text=input_text, + request_obj=request_obj, + *args, + **kwargs, + ) + + if not result or not self._break_token_ids: + return result + + # Remove break/end tokens and fix multimodal item offsets + input_ids = result.get("input_ids", []) + mm_items = result.get("mm_items", []) + + new_input_ids = [] + old_to_new = {} + for old_idx, token_id in enumerate(input_ids): + if token_id not in self._break_token_ids: + old_to_new[old_idx] = len(new_input_ids) + new_input_ids.append(token_id) + + if len(new_input_ids) == len(input_ids): + return result + + # Remap multimodal item offsets to account for removed tokens + for mm_item in mm_items: + if not mm_item.offsets: + continue + new_indices = sorted( + old_to_new[idx] + for start, end in mm_item.offsets + for idx in range(start, end + 1) + if idx in old_to_new + ) + if new_indices: + mm_item.offsets = [(new_indices[0], new_indices[-1])] + + result["input_ids"] = new_input_ids + return result