chore: import upstream snapshot with attribution

python/sglang/srt/models/lightonocr.py

@@ -0,0 +1,299 @@
+# 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
+        config_dict = vision_config.to_dict()
+        if config_dict.get("rope_parameters"):
+            config_dict["rope_theta"] = config_dict["rope_parameters"].get("rope_theta")
+        dataclass_fields = {field.name for field in fields(VisionEncoderArgs)}
+        vision_args = {
+            key: value for key, value in config_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