chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support LightOnOCR-2-1B (#17806)

Browse Source
This commit is contained in:
Shivam jindal
2026-01-29 20:33:41 +05:30
committed by GitHub
parent 3c9cc44ff5
commit 0769de9b0f
3 changed files with 409 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/configs/model_config.py
View File

@@ -1147,6 +1147,7 @@ multimodal_model_archs = [
"LlavaQwenForCausalLM",
"LlavaForConditionalGeneration",
"LlavaVidForCausalLM",
"LightOnOCRForConditionalGeneration",
"MiniCPMO",
"MiniCPMV",
"Mistral3ForConditionalGeneration",

298
python/sglang/srt/models/lightonocr.py Normal file
View File

@@ -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

110
python/sglang/srt/multimodal/processors/lightonocr.py Normal file
View File

@@ -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