From 7106f6c8e1509cd57abeafd5d50cb1beaffbc63c Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang <2448370773@qq.com> Date: Tue, 27 Jan 2026 14:24:00 +0800 Subject: [PATCH] [GLM-OCR] Support GLM-OCR Model (#17582) Signed-off-by: Xinyuan Tong Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: Xinyuan Tong --- .../multimodal_language_models.md | 1 + python/sglang/srt/configs/model_config.py | 8 +- python/sglang/srt/layers/attention/vision.py | 68 ++- python/sglang/srt/models/glm4.py | 24 +- python/sglang/srt/models/glm4v.py | 2 - python/sglang/srt/models/glm_ocr.py | 435 ++++++++++++++++++ python/sglang/srt/models/glm_ocr_nextn.py | 162 +++++++ .../sglang/srt/multimodal/processors/glm4v.py | 7 +- python/sglang/srt/utils/common.py | 1 + 9 files changed, 679 insertions(+), 29 deletions(-) create mode 100644 python/sglang/srt/models/glm_ocr.py create mode 100644 python/sglang/srt/models/glm_ocr_nextn.py diff --git a/docs/supported_models/multimodal_language_models.md b/docs/supported_models/multimodal_language_models.md index 1677bb574..9f645ab95 100644 --- a/docs/supported_models/multimodal_language_models.md +++ b/docs/supported_models/multimodal_language_models.md @@ -42,6 +42,7 @@ in the GitHub search bar. | **Phi-4-multimodal-instruct** | `microsoft/Phi-4-multimodal-instruct` | Phi-4-multimodal-instruct is the multimodal variant of the Phi-4-mini model, enhanced with LoRA for improved multimodal capabilities. It supports text, vision and audio modalities in SGLang. | | | **MiMo-VL** (7B) | `XiaomiMiMo/MiMo-VL-7B-RL` | Xiaomi's compact yet powerful vision-language model featuring a native resolution ViT encoder for fine-grained visual details, an MLP projector for cross-modal alignment, and the MiMo-7B language model optimized for complex reasoning tasks. | | | **GLM-4.5V** (106B) / **GLM-4.1V**(9B) | `zai-org/GLM-4.5V` | GLM-4.5V and GLM-4.1V-Thinking: Towards Versatile Multimodal Reasoning with Scalable Reinforcement Learning | Use `--chat-template glm-4v` | +| **GLM-OCR** | `zai-org/GLM-OCR` | GLM-OCR: A fast and accurate general OCR model | | | **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. | diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index fbcb9754c..359aa811e 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -278,6 +278,11 @@ class ModelConfig: ]: self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" + if is_draft_model and self.hf_config.architectures[0] in [ + "GlmOcrForConditionalGeneration", + ]: + self.hf_config.architectures[0] = "GlmOcrForConditionalGenerationNextN" + if ( is_draft_model and self.hf_config.architectures[0] == "LongcatFlashForCausalLM" @@ -935,7 +940,7 @@ class ModelConfig: needs_tf_v5 = is_glm_46vmoe tf_version = version.parse(tf_version_str) - required_version = version.parse("5.0.0") + required_version = version.parse("5.0.0dev0") if tf_version < required_version: if needs_tf_v5: @@ -1132,6 +1137,7 @@ multimodal_model_archs = [ "Gemma3nForConditionalGeneration", "Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", + "GlmOcrForConditionalGeneration", "GlmAsrForConditionalGeneration", "Grok1VForCausalLM", "Grok1AForCausalLM", diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 68593ab08..648ff091a 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -590,6 +590,7 @@ class VisionAttention(nn.Module): num_dummy_heads: int = 0, qkv_bias: bool = True, qk_normalization: bool = False, + qk_normalization_by_head_size: bool = False, layer_norm_eps: float = 1e-06, customized_position_embedding_applier: Callable[ [torch.Tensor, torch.Tensor, Any, Any], Tuple[torch.Tensor, torch.Tensor] @@ -617,30 +618,19 @@ class VisionAttention(nn.Module): self.kv_size = self.num_attention_kv_heads_per_partition * self.head_size self.qk_normalization = qk_normalization + self.qk_normalization_by_head_size = qk_normalization_by_head_size # Additional dummy heads are used to enable TP for common GPU counts. self.dummy_dim = (num_dummy_heads + num_heads) * self.head_size if self.qk_normalization: - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - ) - if get_global_server_args().rl_on_policy_target is not None - else {} + self.q_norm, self.k_norm = self._init_qk_norm( + self.dummy_dim, layer_norm_eps, embed_dim ) - self.q_norm = RMSNorm( - self.dummy_dim, - eps=layer_norm_eps, - var_hidden_size=embed_dim, - **norm_kwargs, - ) - self.k_norm = RMSNorm( - self.dummy_dim, - eps=layer_norm_eps, - var_hidden_size=embed_dim, - **norm_kwargs, + + elif self.qk_normalization_by_head_size: + self.q_norm, self.k_norm = self._init_qk_norm( + self.head_size, layer_norm_eps ) # Select attention backend via a unified method @@ -702,6 +692,31 @@ class VisionAttention(nn.Module): self.aux_stream = aux_stream self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] if aux_stream else [] + def _init_qk_norm( + self, norm_dim: int, eps: float, var_hidden_size: Optional[int] = None + ): + norm_kwargs = ( + dict( + weight_dtype=torch.float32, + cast_x_before_out_mul=True, + ) + if get_global_server_args().rl_on_policy_target is not None + else {} + ) + q_norm = RMSNorm( + norm_dim, + eps=eps, + var_hidden_size=var_hidden_size, + **norm_kwargs, + ) + k_norm = RMSNorm( + norm_dim, + eps=eps, + var_hidden_size=var_hidden_size, + **norm_kwargs, + ) + return q_norm, k_norm + def _determine_attention_backend(self, passed_backend: Optional[str]) -> str: """Decide the multimodal attention backend string. @@ -734,6 +749,16 @@ class VisionAttention(nn.Module): return backend + def _apply_qk_norm_head_size(self, q: torch.Tensor, k: torch.Tensor): + """apply qk norm for GLM-OCR vit attn""" + q_by_head = q.reshape(-1, self.head_size) + q_by_head = self.q_norm(q_by_head) + k_by_head = k.reshape(-1, self.head_size) + k_by_head = self.k_norm(k_by_head) + q = q_by_head.view(q.shape) + k = k_by_head.view(k.shape) + return q, k + def _apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor): """apply qk norm for internvl vit attn""" @@ -816,6 +841,8 @@ class VisionAttention(nn.Module): q = q.reshape(bsz * s, head, -1).contiguous() k = k.reshape(bsz * s, kv_head, -1).contiguous() v = v.reshape(bsz * s, kv_head, -1).contiguous() + if self.qk_normalization_by_head_size: + q, k = self._apply_qk_norm_head_size(q, k) else: # [b, s, embed_dim] --> [s, b, embed_dim] x = rearrange(x, "b s ... -> s b ...") @@ -837,6 +864,9 @@ class VisionAttention(nn.Module): rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v) ] + if self.qk_normalization_by_head_size: + q, k = self._apply_qk_norm_head_size(q, k) + cos = None sin = None @@ -881,7 +911,7 @@ class VisionAttention(nn.Module): assert v.dim() == 3, v.dim() # internvl - if self.qk_normalization: + if self.qk_normalization and not self.qk_normalization_by_head_size: # jit kernel if can_use_jit_qk_norm(self.head_size, q.dtype): diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py index f5d1c17f8..ba40a1f74 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -119,6 +119,7 @@ class Glm4Attention(nn.Module): quant_config: Optional[QuantizationConfig] = None, dual_chunk_attention_config: Optional[dict[str, Any]] = None, partial_rotary_factor: float = 0.5, + bias: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -153,7 +154,7 @@ class Glm4Attention(nn.Module): self.head_dim, self.total_num_heads, self.total_num_kv_heads, - bias=True, + bias=bias, quant_config=quant_config, prefix=add_prefix("qkv_proj", prefix), ) @@ -216,13 +217,23 @@ class Glm4DecoderLayer(nn.Module): ) -> None: super().__init__() self.hidden_size = config.hidden_size - rope_theta = getattr(config, "rope_theta", 1000000) - rope_scaling = getattr(config, "rope_scaling", None) + + rp = getattr(config, "rope_parameters", None) + if isinstance(rp, dict): + rope_theta = rp.get("rope_theta", getattr(config, "rope_theta", 1000000)) + partial_rotary_factor = rp.get( + "partial_rotary_factor", + getattr(config, "partial_rotary_factor", 0.5), + ) + rope_scaling = getattr(config, "rope_scaling", None) + else: + rope_theta = getattr(config, "rope_theta", 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + partial_rotary_factor = getattr(config, "partial_rotary_factor", 0.5) + + bias = getattr(config, "attention_bias", True) max_position_embeddings = getattr(config, "max_position_embeddings", 32768) head_dim = getattr(config, "head_dim", None) - partial_rotary_factor = getattr( - getattr(config, "rope_parameters", None), "partial_rotary_factor", None - ) or getattr(config, "partial_rotary_factor", 0.5) dual_chunk_attention_config = getattr( config, "dual_chunk_attention_config", None ) @@ -238,6 +249,7 @@ class Glm4DecoderLayer(nn.Module): quant_config=quant_config, dual_chunk_attention_config=dual_chunk_attention_config, partial_rotary_factor=partial_rotary_factor, + bias=bias, prefix=add_prefix("self_attn", prefix), ) diff --git a/python/sglang/srt/models/glm4v.py b/python/sglang/srt/models/glm4v.py index 243677f82..44f9837a9 100644 --- a/python/sglang/srt/models/glm4v.py +++ b/python/sglang/srt/models/glm4v.py @@ -758,8 +758,6 @@ class Glm4vForConditionalGeneration(nn.Module): name = name.replace(r"model.language_model.", r"model.") if "model.visual." in name: name = name.replace("model.visual.", "visual.") - if name.startswith("lm_head.") and not self.pp_group.is_last_rank: - continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: diff --git a/python/sglang/srt/models/glm_ocr.py b/python/sglang/srt/models/glm_ocr.py new file mode 100644 index 000000000..c2f4adc4c --- /dev/null +++ b/python/sglang/srt/models/glm_ocr.py @@ -0,0 +1,435 @@ +# Copyright 2023-2024 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. +# ============================================================================== + +# Modeling from: +# ./llama.py and +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/GlmOcr/modular_GlmOcr.py +"""Inference-only GLM-OCR model compatible with HuggingFace weights.""" + +import logging +from functools import lru_cache +from typing import Iterable, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.models.glm_ocr.configuration_glm_ocr import ( + GlmOcrConfig, + GlmOcrVisionConfig, +) + +from sglang.srt.distributed.parallel_state import get_pp_group +from sglang.srt.layers.attention import vision_utils +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.glm4 import Glm4Model +from sglang.srt.models.glm4v import ( + Glm4vForConditionalGeneration, + Glm4vPatchMerger, + Glm4vRMSNorm, + Glm4vVisionMLP, + Glm4vVisionModel, + Glm4vVisionPatchEmbed, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix +from sglang.srt.utils.hf_transformers_utils import get_processor + +logger = logging.getLogger(__name__) + +cached_get_processor = lru_cache(get_processor) + + +class GlmOcrRMSNorm(Glm4vRMSNorm): + pass + + +class GlmOcrVisionMLP(Glm4vVisionMLP): + pass + + +class GlmOcrVisionBlock(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + num_heads: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + attn_qkv_bias: bool = True, + num_dummy_heads: int = 0, + rms_norm_eps: float = 1e-5, + use_data_parallel: bool = False, + ) -> None: + super().__init__() + self.norm1 = RMSNorm(dim, eps=rms_norm_eps) + self.norm2 = RMSNorm(dim, eps=rms_norm_eps) + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + use_qkv_parallel=True, + qkv_bias=attn_qkv_bias, + proj_bias=True, + qk_normalization_by_head_size=True, + flatten_batch=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + num_dummy_heads=num_dummy_heads, + use_data_parallel=use_data_parallel, + ) + self.mlp = GlmOcrVisionMLP( + dim, + intermediate_dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + use_data_parallel=use_data_parallel, + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + ) -> torch.Tensor: + S, B, H = x.shape + # norm1: flatten to 2D -> [S*B, H], then reshape back + x2d = x.reshape(-1, H) + hidden_states = self.norm1(x2d).reshape(S, B, H) + + # Attention expects [B, S, H] + hidden_states = rearrange(hidden_states, "s b h -> b s h") + attn = self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + ) + attn = rearrange(attn, "b s h -> s b h") + + # norm2 with fused residual-add: also 2D + attn2d = attn.reshape(-1, H) + x_norm_2d, x_after_add_2d = self.norm2(x2d, residual=attn2d) + x_norm = x_norm_2d.reshape(S, B, H) + x_after_add = x_after_add_2d.reshape(S, B, H) + + # MLP and final residual + mlp_out = self.mlp(x_norm) + x = x_after_add + mlp_out + return x + + +class GlmOcrVisionPatchEmbed(Glm4vVisionPatchEmbed): + pass + + +class GlmOcrVisionPatchMerger(Glm4vPatchMerger): + pass + + +class GlmOcrVisionModel(Glm4vVisionModel): + def __init__( + self, + vision_config: GlmOcrVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + super().__init__(vision_config, quant_config, prefix, use_data_parallel) + + patch_size = vision_config.patch_size + temporal_patch_size = vision_config.temporal_patch_size + in_channels = vision_config.in_channels + depth = vision_config.depth + self.hidden_size = vision_config.hidden_size + self.num_heads = vision_config.num_heads + + self.patch_size = vision_config.patch_size + self.spatial_merge_size = vision_config.spatial_merge_size + self.out_hidden_size = vision_config.out_hidden_size + self.intermediate_size = vision_config.intermediate_size + self.use_data_parallel = use_data_parallel + + self.patch_embed = GlmOcrVisionPatchEmbed( + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + in_channels=in_channels, + hidden_size=self.hidden_size, + ) + + head_dim = self.hidden_size // self.num_heads + self.rotary_pos_emb = get_rope( + head_size=head_dim, + rotary_dim=head_dim // 2, + max_position=8192, + base=10000.0, + is_neox_style=True, + ) + + self.blocks = nn.ModuleList( + [ + GlmOcrVisionBlock( + dim=self.hidden_size, + intermediate_dim=self.intermediate_size, + num_heads=self.num_heads, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{layer_idx}", prefix), + rms_norm_eps=vision_config.rms_norm_eps, + attn_qkv_bias=vision_config.attention_bias, + use_data_parallel=use_data_parallel, + ) + for layer_idx in range(depth) + ] + ) + self.merger = GlmOcrVisionPatchMerger( + d_model=vision_config.out_hidden_size, + context_dim=vision_config.out_hidden_size * vision_config.in_channels, + quant_config=quant_config, + bias=False, + prefix=add_prefix("merger", prefix), + use_data_parallel=use_data_parallel, + ) + + self.downsample = nn.Conv2d( + in_channels=vision_config.hidden_size, + out_channels=vision_config.out_hidden_size, + kernel_size=vision_config.spatial_merge_size, + stride=vision_config.spatial_merge_size, + ) + self.post_layernorm = GlmOcrRMSNorm( + vision_config.hidden_size, eps=vision_config.rms_norm_eps + ) + + def forward(self, x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + # patchify + x = x.to(device=self.device, dtype=self.dtype) + x = self.patch_embed(x) + + # compute position embedding + rotary_pos_emb_cos, rotary_pos_emb_sin, image_type_ids = self.rot_pos_emb( + grid_thw + ) + # compute cu_seqlens + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum(dim=0, dtype=torch.int32) + cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens]) + + rotary_pos_emb_cos = torch.cat([rotary_pos_emb_cos, rotary_pos_emb_cos], dim=-1) + rotary_pos_emb_sin = torch.cat([rotary_pos_emb_sin, rotary_pos_emb_sin], dim=-1) + + # x.shape: (s, b, d) where b=1 for vision processing + # transformers + x = x.unsqueeze(1) + for blk in self.blocks: + x = blk( + x, + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + ) + + # adapter + x = self.post_layernorm(x) + x = x.view(-1, self.spatial_merge_size, self.spatial_merge_size, x.shape[-1]) + x = x.permute(0, 3, 1, 2) + x = self.downsample(x).view(-1, self.out_hidden_size) + x = self.merger(x) + + return x + + +class GlmOcrForConditionalGeneration(Glm4vForConditionalGeneration): + def __init__( + self, + config: GlmOcrConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, quant_config, prefix) + + self.pp_group = get_pp_group() + self.config = config + self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder + self.visual = GlmOcrVisionModel( + vision_config=config.vision_config, + quant_config=quant_config, + prefix=add_prefix("visual", prefix), + use_data_parallel=self.use_data_parallel, + ) + + vision_utils.update_vit_attn_dummy_heads_config(self.config) + + self.model = Glm4Model( + config, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + if self.pp_group.is_last_rank: + if self.pp_group.world_size == 1 and self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + else: + # ranks other than the last rank will have a placeholder layer + self.lm_head = PPMissingLayer() + + self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling + + self.logits_processor = LogitsProcessor(config) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + + # For EAGLE3 support + self.capture_aux_hidden_states = False + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): + if is_nextn: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + assert num_nextn_layers == 1, "Only 1 nextn layer is supported" + # compatible with old design + nextn_layer_id = ( + 0 + if self.config.num_hidden_layers == 1 + else self.config.num_hidden_layers + ) + else: + raise ValueError("num_nextn_predict_layers is not in the config") + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".up_proj", 1), + (".gate_up_proj", ".gate_proj", 0), + ] + + if is_nextn: + nextn_layer_prefix = f"model.layers.{nextn_layer_id}" + nextn_spec_weight_names = [ + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", + ] + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + + # For the PP case, we add special handling for lm_head.weight, + # - On non–last ranks: we continue, because this stage is supposed to + # be just an empty PPMissingLayer shell. + # - On the last rank: params_dict is expected to contain lm_head.weight, + # so it will never hit the branch "if name not in params_dict". + # + # For all other parameters, such like + # "model.visual.blocks.20.mlp.gate_proj.weight", the unified rule is: + # If this name does not exist in the current rank’s params_dict, + # it does not belong to this pipeline stage, thus we simply continue. + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "language_model" in name: + name = name.replace(r"model.language_model.", r"model.") + if "model.visual." in name: + name = name.replace("model.visual.", "visual.") + + if not is_nextn: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + continue + else: + if not name.startswith(nextn_layer_prefix): + continue + + # Use shared head and embed weights from target model + if "shared_head.head" in name or "embed_tokens" in name: + continue + + is_decoder = True + # For nextn specific weights + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace(nextn_layer_prefix, "model") + is_decoder = False + break + # For decoder layer weights + if is_decoder: + name = name.replace(nextn_layer_prefix, "model.decoder") + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if "visual" in name: + # adapt to VisionAttention + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + + try: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + except KeyError: + print(params_dict.keys()) + raise + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if "visual" in name: + loaded_weight = vision_utils.pad_vit_attn_dummy_heads( + self.config, name, loaded_weight + ) + weight_loader(param, loaded_weight) + + +EntryClass = [GlmOcrForConditionalGeneration] diff --git a/python/sglang/srt/models/glm_ocr_nextn.py b/python/sglang/srt/models/glm_ocr_nextn.py new file mode 100644 index 000000000..ae771af53 --- /dev/null +++ b/python/sglang/srt/models/glm_ocr_nextn.py @@ -0,0 +1,162 @@ +# Copyright 2023-2024 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. +# ============================================================================== + +"""Inference-only GLM-OCR Speculative Decoding.""" + +import logging +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.glm4 import Glm4DecoderLayer +from sglang.srt.models.glm_ocr import GlmOcrForConditionalGeneration +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class GlmOcrModelNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + if quant_config is not None and quant_config.get_name() == "modelopt_fp4": + logger.warning( + "Overriding GlmOcrModelNextN quant config for modelopt_fp4 GLM-OCR model." + ) + quant_config = None + + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + self.decoder = Glm4DecoderLayer( + config, + 0, + quant_config=quant_config, + prefix=add_prefix("decoder", prefix), + ) + + self.shared_head = nn.Module() + self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + if hidden_states.shape[0] > 0: + hidden_states = self.eh_proj( + torch.cat( + ( + self.enorm(hidden_states), + self.hnorm(forward_batch.spec_info.hidden_states), + ), + dim=-1, + ) + ) + + residual = None + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states, residual = self.decoder( + positions, hidden_states, forward_batch, residual + ) + + if not forward_batch.forward_mode.is_idle(): + if residual is not None: + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + else: + hidden_states = self.shared_head.norm(hidden_states) + + return hidden_states + + +class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.tp_size = get_tensor_model_parallel_world_size() + self.quant_config = quant_config + self.model = GlmOcrModelNextN( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("model.shared_head.head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + + self.num_fused_shared_experts = ( + 0 if get_global_server_args().disable_shared_experts_fusion else 1 + ) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + super().load_weights(weights, is_nextn=True) + + +EntryClass = [GlmOcrForConditionalGenerationNextN] diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py index 80d717a7a..45dc785fb 100644 --- a/python/sglang/srt/multimodal/processors/glm4v.py +++ b/python/sglang/srt/multimodal/processors/glm4v.py @@ -3,6 +3,7 @@ from typing import List, Union from sglang.srt.layers.rotary_embedding import MRotaryEmbedding from sglang.srt.models.glm4v import Glm4vForConditionalGeneration from sglang.srt.models.glm4v_moe import Glm4vMoeForConditionalGeneration +from sglang.srt.models.glm_ocr import GlmOcrForConditionalGeneration from sglang.srt.multimodal.processors.base_processor import ( BaseMultimodalProcessor as SGLangBaseProcessor, ) @@ -10,7 +11,11 @@ from sglang.srt.multimodal.processors.base_processor import MultimodalSpecialTok class Glm4vImageProcessor(SGLangBaseProcessor): - models = [Glm4vForConditionalGeneration, Glm4vMoeForConditionalGeneration] + models = [ + Glm4vForConditionalGeneration, + Glm4vMoeForConditionalGeneration, + GlmOcrForConditionalGeneration, + ] def __init__(self, hf_config, server_args, _processor, *args, **kwargs): super().__init__(hf_config, server_args, _processor, *args, **kwargs) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 8e39ee4ca..2b560fcc3 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -2906,6 +2906,7 @@ def is_fa3_default_architecture(hf_config): "Glm4MoeForCausalLM", "Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", + "GlmOcrForConditionalGeneration", "Step3VLForConditionalGeneration", "StepVLForConditionalGeneration", "MiMoV2FlashForCausalLM",