# 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. # ============================================================================== import logging import re from functools import lru_cache from typing import Iterable, List, Optional, Set, Tuple, TypedDict, Union import torch from torch import nn from transformers import ( Gemma4AudioConfig, Gemma4Config, Gemma4TextConfig, Gemma4VisionConfig, PreTrainedModel, ) from sglang.srt.distributed import get_pp_group from sglang.srt.environ import envs from sglang.srt.layers.attention.triton_backend import TritonAttnBackend from sglang.srt.layers.layernorm import Gemma4RMSNorm from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.managers.mm_utils import ( MultiModalityDataPaddingPatternMultimodalTokens, general_mm_embed_routine, ) from sglang.srt.managers.schedule_batch import ( Modality, MultimodalDataItem, MultimodalInputs, flatten_nested_list, ) from sglang.srt.model_executor.forward_batch_info import ( ForwardBatch, ForwardMode, PPProxyTensors, ) from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from sglang.srt.models.gemma4_audio import Gemma4AudioEncoder from sglang.srt.models.gemma4_causal import Gemma4TextModel, pp_filter_load_weight from sglang.srt.models.gemma4_vision import Gemma4VisionEncoder 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 Gemma4ImagePixelInputs(TypedDict): pixel_values: torch.Tensor """Shape: `(batch_size * num_images, num_channels, height, width)`""" class Gemma4AudioInputs(TypedDict): input_features_padded: torch.Tensor """Shape: `(batch_size * num_audio, seq_length, num_features)`""" input_features_mask: torch.Tensor """Shape: `(batch_size * num_audio, seq_length)`""" class Gemma4MultimodalEmbedder(nn.Module): """Projects vision/audio soft tokens into LM embedding space.""" def __init__( self, multimodal_config: Union[Gemma4AudioConfig, Gemma4VisionConfig], text_config: Gemma4TextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.eps = multimodal_config.rms_norm_eps self.text_hidden_size = text_config.hidden_size # Audio tower uses output_proj_dims (1536) rather than hidden_size # (1024); vision uses hidden_size (768) directly. embedding_dim = ( getattr(multimodal_config, "output_proj_dims", None) or multimodal_config.hidden_size ) self.embedding_projection = ReplicatedLinear( embedding_dim, self.text_hidden_size, bias=False, quant_config=quant_config, prefix=add_prefix("embedding_projection", prefix), ) self.embedding_pre_projection_norm = Gemma4RMSNorm( embedding_dim, eps=self.eps, with_scale=False, ) def forward( self, inputs_embeds: torch.Tensor, ) -> torch.Tensor: """Project soft tokens from a multimodal tower into LM space.""" embs_normed = self.embedding_pre_projection_norm(inputs_embeds) embs_proj, _ = self.embedding_projection(embs_normed) return embs_proj class Gemma4ForConditionalGeneration(PreTrainedModel): config_class = Gemma4Config """Gemma4 multimodal model for conditional generation.""" # BitandBytes specific attributes default_bitsandbytes_target_modules = [ ".gate_proj.", ".down_proj.", ".up_proj.", ".q_proj.", ".k_proj.", ".v_proj.", ".o_proj.", ] bitsandbytes_stacked_params_mapping = { "q_proj": ("qkv_proj", 0), "k_proj": ("qkv_proj", 1), "v_proj": ("qkv_proj", 2), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } packed_modules_mapping = { "qkv_proj": [ "q_proj", "k_proj", "v_proj", ], "gate_up_proj": [ "gate_proj", "up_proj", ], } # LoRA specific attributes supported_lora_modules = [ "qkv_proj", "o_proj", "gate_up_proj", "down_proj", ] # Gemma does not apply LoRA to the embedding layer embedding_modules = {} embedding_padding_modules = [] supports_lora = True def __init__( self, config: Gemma4Config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__(config=config) self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config text_config = config.text_config prefix = add_prefix("model", prefix) # Vision/audio encoders + their projection embedders are only consumed # at the input-embedding stage, so they live on the first PP rank only. if self.pp_group.is_first_rank: self.vision_tower = Gemma4VisionEncoder( config=config.vision_config, quant_config=quant_config, prefix=add_prefix("vision_tower", prefix), ) self.embed_vision = Gemma4MultimodalEmbedder( config.vision_config, config.text_config, quant_config=quant_config, prefix=add_prefix("embed_vision", prefix), ) if getattr(config, "audio_config", None) is not None: self.audio_tower = Gemma4AudioEncoder( config=config.audio_config, quant_config=quant_config, prefix=add_prefix("audio_tower", prefix), ) self.embed_audio = Gemma4MultimodalEmbedder( config.audio_config, config.text_config, quant_config=quant_config, prefix=add_prefix("embed_audio", prefix), ) else: self.audio_tower = None self.embed_audio = None else: self.vision_tower = PPMissingLayer() self.embed_vision = PPMissingLayer() self.audio_tower = None self.embed_audio = None self.vocab_size = config.text_config.vocab_size self.vocab_size_per_layer_input = getattr( config.text_config, "vocab_size_per_layer_input", config.text_config.vocab_size, ) # Text model — internal Gemma4TextModel is already PP-aware. self.language_model = Gemma4TextModel( config.text_config, quant_config, prefix=add_prefix("language_model", prefix), ) # Tied embeddings: under PP the embed_tokens lives on the first rank # while logits run on the last rank, so we can't reuse the embedding # module directly. For PP=1 keep the original tying; for PP>1 # materialize a real ParallelLMHead on the last rank and route the # checkpoint embedding into it during load_weights. text_tie = getattr(text_config, "tie_word_embeddings", True) if self.pp_group.world_size == 1 and text_tie: self.lm_head = self.language_model.embed_tokens elif self.pp_group.is_last_rank: self.lm_head = ParallelLMHead( text_config.vocab_size, text_config.hidden_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), ) else: self.lm_head = PPMissingLayer() # Create logits processor for the multimodal model self.logits_processor = LogitsProcessor(config.text_config) self.capture_aux_hidden_states = False self.post_init() @property def model(self): # Alias .model to .language_model so this class satisfies the piecewise # CUDA graph gate (which checks `hasattr(model, "model")`). Implemented # as a property to avoid registering a duplicate submodule in # `_modules`, which would double state_dict keys and disturb # ShardedStateLoader / CPU-offload / dummy-init paths. return self.language_model def __setattr__(self, name, value): # Block writes to "model" so the runner's # `self.model.model = resolve_language_model(self.model)` (which for # this class returns language_model itself) is a no-op rather than a # nn.Module submodule registration. Without this, nn.Module.__setattr__ # would bypass the @property's setter for Module values and pollute # `_modules` with a duplicate alias, doubling state_dict keys. if name == "model": return super().__setattr__(name, value) def pad_input_ids( self, input_ids: List[int], mm_inputs: MultimodalInputs, ) -> List[int]: """Pad input IDs with image and audio tokens.""" pattern = MultiModalityDataPaddingPatternMultimodalTokens() return pattern.pad_input_tokens(input_ids, mm_inputs) def get_input_embeddings(self) -> nn.Embedding: return self.language_model.get_input_embeddings() def get_embed_and_head(self) -> Tuple[torch.Tensor, torch.Tensor]: # Gemma 4 multimodal ties its LM head to the text embed_tokens embed = self.language_model.embed_tokens.weight return embed, embed def get_attention_sliding_window_size(self): return getattr(self.config.text_config, "sliding_window", -1) - 1 def set_dflash_layers_to_capture(self, layer_ids: List[int]): if layer_ids is None: raise ValueError( "DFLASH requires explicit layer_ids for aux hidden capture." ) self.capture_aux_hidden_states = True self.language_model.layers_to_capture = [val + 1 for val in layer_ids] def prepare_attn_masks( self, forward_batch: ForwardBatch, input_ids: torch.Tensor, mask_dtype: torch.dtype, ): """Prepare bidirectional attention masks for image tokens. Gemma 4 uses bidirectional attention for image soft tokens during prefill. Following the HF implementation, bidirectional attention is only enabled within each individual image group (same-item tokens), not across items. Currently only the TritonAttnBackend supports this. TODO(kpham-sgl): Guard appropriately for gemma3_mm.py:prepare_attn_masks() """ if not isinstance(get_attn_backend(), TritonAttnBackend): logger.warning_once( "Bidirectional attention for image tokens requires TritonAttnBackend. " "Falling back to causal attention, which may degrade image quality." ) return assert forward_batch.forward_mode == ForwardMode.EXTEND bidirectional_attn_masks_list = [] bidirectional_attn_mask_indptr = torch.zeros( forward_batch.batch_size + 1, dtype=torch.int32, device=input_ids.device ) split_images = [] for i in range(forward_batch.batch_size): extend_seq_len = forward_batch.extend_seq_lens[i] prefix_len = forward_batch.extend_prefix_lens[i] bidirectional_attn_mask = torch.zeros( extend_seq_len, extend_seq_len + prefix_len, dtype=mask_dtype, device=input_ids.device, ) # Start with causal mask bidirectional_attn_mask.fill_(1) bidirectional_attn_mask = bidirectional_attn_mask.tril(diagonal=prefix_len) # HF only enables bidirectional attention for image tokens, # not video or audio (see create_causal_mask_mapping). mm_inputs = forward_batch.mm_inputs[i] if mm_inputs is not None: for mm_item in mm_inputs.mm_items: if mm_item.is_image(): for im_begin, im_end in mm_item.offsets: # Note(kpham-sgl): We only apply bidirectional attention when the image token span # is fully contained in the extend window. Otherwise, we silently fall back to # causal attention. # FIXME(kpham-sgl): This is a hack to work around the fact that the image token span # might not be fully contained in the extend window during chunked prefill. # We should fix this by properly making chunked prefill mask aware. if ( im_begin >= prefix_len and im_end < prefix_len + extend_seq_len ): bidirectional_attn_mask[ im_begin - prefix_len : im_end + 1 - prefix_len, im_begin : im_end + 1, ] = 1 elif ( im_end >= prefix_len and im_begin < prefix_len + extend_seq_len ): split_images.append((i, im_begin, im_end)) bidirectional_attn_masks_list.append(bidirectional_attn_mask.flatten()) bidirectional_attn_mask_indptr[i + 1] = ( bidirectional_attn_mask_indptr[i] + bidirectional_attn_mask.nelement() ) if split_images: num_split_images = len(split_images) logger.warning_once( f"{num_split_images} images are split across chunk boundaries. " "Below are the first 5 images that are split across chunk boundaries: " ) for i, im_begin, im_end in split_images[:5]: logger.warning_once( f"Image {i}:{im_begin}-{im_end} is split across chunk boundaries.\n", ) logger.warning_once( "Those images will receive causal attention. Disable chunked prefill (--chunked-prefill-size=-1) for full bidirectional attention.", ) if bidirectional_attn_masks_list: bidirectional_attn_masks = torch.cat(bidirectional_attn_masks_list, dim=0) get_attn_backend().forward_metadata.mask_indptr = ( bidirectional_attn_mask_indptr ) get_attn_backend().forward_metadata.custom_mask = bidirectional_attn_masks def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: vt = self.vision_tower all_embeds = [] for item in items: all_pixel_values = flatten_nested_list([item.feature]) all_position_ids = flatten_nested_list( [getattr(item, "image_position_ids", None)] ) for pv_idx, pv in enumerate(all_pixel_values): if ( pv.dim() in (2, 3) and pv.shape[-1] == self.config.text_config.hidden_size ): all_embeds.append(pv.to(self.language_model.device)) continue if pv_idx >= len(all_position_ids) or all_position_ids[pv_idx] is None: raise ValueError( f"pixel_values[{pv_idx}] has no matching image_position_ids. " "The HF image processor likely renamed this output — " "update ATTR_NAME_TO_MODALITY in the Gemma4 processor." ) pp = all_position_ids[pv_idx] # Vision tower expects 3-D (batch, num_patches, ...). # A single image may arrive as 2-D; add the batch dim if needed. if pv.dim() == 2: pv = pv.unsqueeze(0) if pp.dim() == 2: pp = pp.unsqueeze(0) pv = pv.to(device=vt.device, dtype=self.language_model.dtype()) pp = pp.to(device=vt.device) pooled, pooler_mask = vt(pv, pp) for hs, mask in zip(pooled, pooler_mask): real_tokens = hs[mask] all_embeds.append( self.embed_vision( inputs_embeds=real_tokens.unsqueeze(0) ).squeeze(0) ) if all_embeds: return torch.cat(all_embeds, dim=0) else: return torch.empty( 0, self.language_model.config.hidden_size, device=next(self.parameters()).device, dtype=self.language_model.dtype(), ) def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: """Encode video frames through the vision tower with video-specific pooling. Each video is (num_frames, num_patches, patch_pixels) with matching position_ids (num_frames, num_patches, 2). Frames are flattened into the batch dimension so each frame is encoded independently, then pooled dynamically based on the input patch count and pooling_kernel_size. """ vt = self.vision_tower all_embeds = [] for item in items: all_pixel_values = flatten_nested_list([item.feature]) all_position_ids = flatten_nested_list( [getattr(item, "video_position_ids", None)] ) for pv_idx, pv in enumerate(all_pixel_values): if ( pv.dim() in (2, 3) and pv.shape[-1] == self.config.text_config.hidden_size ): all_embeds.append(pv.to(self.language_model.device)) continue if pv_idx >= len(all_position_ids) or all_position_ids[pv_idx] is None: raise ValueError( f"pixel_values_videos[{pv_idx}] has no matching video_position_ids." ) pp = all_position_ids[pv_idx] # HF processor returns 4-D tensors # (num_videos, num_frames, num_patches, ...) — collapse to # 3-D (num_frames, num_patches, ...) so each frame is a # batch element for the vision tower. if pv.dim() == 4: pv = pv.reshape(-1, pv.shape[-2], pv.shape[-1]) if pp.dim() == 4: pp = pp.reshape(-1, pp.shape[-2], pp.shape[-1]) pv = pv.to(device=vt.device, dtype=self.language_model.dtype()) pp = pp.to(device=vt.device) pooled, pooler_mask = vt(pv, pp) for hs, mask in zip(pooled, pooler_mask): real_tokens = hs[mask] all_embeds.append( self.embed_vision( inputs_embeds=real_tokens.unsqueeze(0) ).squeeze(0) ) if all_embeds: return torch.cat(all_embeds, dim=0) else: return torch.empty( 0, self.language_model.config.hidden_size, device=next(self.parameters()).device, dtype=self.language_model.dtype(), ) def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: if self.audio_tower is None: raise ValueError( "Audio inputs provided but the model does not have an audio tower." ) all_input_features = flatten_nested_list([item.feature for item in items]) all_input_features_mask = flatten_nested_list( [~item.input_features_mask for item in items] ) all_embeds = [] for input_features, input_features_mask in zip( all_input_features, all_input_features_mask ): if input_features.dim() == 2: input_features = input_features.unsqueeze(0) if input_features_mask.dim() == 1: input_features_mask = input_features_mask.unsqueeze(0) input_features = input_features.to( device=self.audio_tower.device, dtype=self.language_model.dtype(), ) input_features_mask = input_features_mask.to(device=input_features.device) # audio_mel_mask convention: True = padding audio_encodings, audio_mask = self.audio_tower( input_features, input_features_mask ) audio_features = self.embed_audio(inputs_embeds=audio_encodings) for enc, mask in zip(audio_features, audio_mask): all_embeds.append(enc[~mask]) if all_embeds: return torch.cat(all_embeds, dim=0) else: return torch.empty( 0, self.language_model.config.hidden_size, device=next(self.parameters()).device, dtype=self.language_model.dtype(), ) def get_per_layer_inputs( self, input_ids: torch.LongTensor ) -> Optional[torch.Tensor]: return self.language_model.get_per_layer_inputs(input_ids) def project_per_layer_inputs( self, inputs_embeds: torch.Tensor, per_layer_inputs: Optional[torch.Tensor] = None, ) -> torch.Tensor: return self.language_model.project_per_layer_inputs( inputs_embeds, per_layer_inputs ) @torch.no_grad() def forward( self, input_ids: torch.LongTensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, **kwargs: object, ) -> Union[LogitsProcessor, PPProxyTensors]: """Forward pass for multimodal Gemma4.""" is_first_rank = self.pp_group.is_first_rank is_last_rank = self.pp_group.is_last_rank # Only the first PP rank consumes input_ids/input_embeds; later stages # receive activations through pp_proxy_tensors. if is_first_rank and (input_ids is None) ^ (input_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) if envs.SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION.get(): positions = positions + 1 else: positions += 1 per_layer_inputs = None # PLE table and the per-layer projection live on the first rank only, # so non-first ranks must skip this and pull per_layer_inputs from the # PP proxy (forwarded by Gemma4TextModel). if is_first_rank and input_ids is not None: ple_ids = input_ids.clone() pad_id = self.config.text_config.pad_token_id # Use torch.where instead of boolean indexing for NPU graph compatibility ple_ids = torch.where( input_ids == self.config.image_token_id, pad_id, ple_ids ) ple_ids = torch.where( input_ids == self.config.video_token_id, pad_id, ple_ids ) ple_ids = torch.where( input_ids == self.config.audio_token_id, pad_id, ple_ids ) per_layer_inputs = self.get_per_layer_inputs(ple_ids) # Prepare bidirectional attention masks for image tokens during prefill. # mm_inputs is preserved on every PP rank up to the first-rank embed # routine, so each rank's attn_backend can install the mask locally. if ( forward_batch.forward_mode == ForwardMode.EXTEND and forward_batch.contains_image_inputs() ): self.prepare_attn_masks( forward_batch, input_ids, mask_dtype=torch.bool, ) # general_mm_embed_routine already handles PP: it skips the embedding # work on non-first ranks and forwards pp_proxy_tensors via **kwargs. hidden_states = general_mm_embed_routine( input_ids=input_ids, forward_batch=forward_batch, language_model=self.language_model, data_embedding_funcs={ Modality.IMAGE: self.get_image_feature, Modality.VIDEO: self.get_video_feature, Modality.AUDIO: self.get_audio_feature, }, positions=positions, per_layer_inputs=per_layer_inputs, pp_proxy_tensors=pp_proxy_tensors, **kwargs, ) if not is_last_rank: # `hidden_states` is actually a PPProxyTensors flowing to the next # stage; logits processing happens on the last rank only. return hidden_states # Unpack aux_hidden_states if Eagle3 capture is active aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states # PP=1 keeps the original tied-weight behavior of using embed_tokens # directly; under PP we route through the dedicated lm_head module. head = ( self.language_model.embed_tokens if self.pp_group.world_size == 1 and getattr(self.config.text_config, "tie_word_embeddings", True) else self.lm_head ) return self.logits_processor( input_ids, hidden_states, head, forward_batch, aux_hidden_states, ) def tie_weights(self, recompute_mapping=False): # Under PP, embed_tokens (first rank) and lm_head (last rank) live on # different processes, so HF's automatic tying would crash on the # PPMissingLayer side. load_weights routes the embedding into lm_head # on the last rank explicitly, so the tie is a no-op under PP. if self.pp_group.world_size > 1: return return self.language_model.tie_weights() # Standard stacked-params mapping for fused QKV / GateUp linears # in the text decoder. Also consumed by the tower QKV remap (step 2). 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), ] # Regex for fused QKV in vision/audio towers. # Vision: *.self_attn.{q,k,v}_proj.* Audio: *.attn.{q,k,v}_proj.* _RE_TOWER_QKV = re.compile( r"(.+\.(?:self_attn|attn))\.(q_proj|k_proj|v_proj)\.(.*)" ) # Regex for fused GateUp in the vision tower MLP. _RE_TOWER_GATE_UP = re.compile(r"(.+\.mlp)\.(gate_proj|up_proj)\.(.*)") _RE_AUDIO_LAYER = re.compile(r"(audio_tower)\.layers\.(\d+)\.(.*)") @staticmethod def _remap_audio_tower_name(name: str) -> str: """Remap audio tower checkpoint names to our module tree. Checkpoint naming (``layers``, ``self_attn``, ``feed_forward1/2``, etc.) differs from our module tree (``conformer``, ``attention.attn``, ``ffw_layer_start/end``, etc.). Applied before ``_remap_tower_name``. """ if "audio_tower." not in name: return name # SSCP conv block: layer0/layer1 → conv_0/conv_1 name = name.replace( "subsample_conv_projection.layer0.", "subsample_conv_projection.conv_0.", ) name = name.replace( "subsample_conv_projection.layer1.", "subsample_conv_projection.conv_1.", ) # Conformer layers: audio_tower.layers.{i} → audio_tower.conformer.{i} m = Gemma4ForConditionalGeneration._RE_AUDIO_LAYER.match(name) if m: tower, layer_idx, suffix = m.groups() # Order matters: more specific patterns first. # relative_k_proj → relative_position_embedding.pos_proj suffix = suffix.replace( "self_attn.relative_k_proj.", "attention.attn.relative_position_embedding.pos_proj.", ) # self_attn.post → attention.post (the output projection) suffix = suffix.replace("self_attn.post.", "attention.post.") # general self_attn → attention.attn suffix = suffix.replace("self_attn.", "attention.attn.") # norms suffix = suffix.replace("norm_pre_attn.", "attention.pre_attn_norm.") suffix = suffix.replace("norm_post_attn.", "attention.post_norm.") suffix = suffix.replace("norm_out.", "norm.") # feed-forward blocks suffix = suffix.replace("feed_forward1.", "ffw_layer_start.") suffix = suffix.replace("feed_forward2.", "ffw_layer_end.") name = f"{tower}.conformer.{layer_idx}.{suffix}" return name @staticmethod def _remap_tower_name(name: str, params_dict: dict) -> str: """Remap a vision/audio tower checkpoint name to our module tree. Three transformations, applied in order: 1. **Fused QKV** — ``{q,k,v}_proj.*`` → ``qkv.*`` Weight/bias are redirected into the fused ``qkv.{proj}.{attr}`` namespace (stacked-params then merges them into ``qkv_proj``). Clip buffers are split: ``input_*`` → shared ``qkv.input_*``, ``output_*`` → per-projection ``qkv.{q,k,v}_output_*``. 2. **Fused GateUp** — ``{gate,up}_proj.*`` → ``gate_up.*`` Same pattern as QKV. 3. **Clippable wrapper** — ``*.weight``/``*.bias`` → ``*.linear.weight`` Catches the remaining (non-fused) clippable linears whose inner ``RowParallelLinear``/``ColumnParallelLinear`` lives at ``.linear``. Falls back to the original name when ``.linear.`` does not exist in ``params_dict`` (plain linears, norms, conv weights, etc.). """ # Step 1: fused QKV m = Gemma4ForConditionalGeneration._RE_TOWER_QKV.match(name) if m: pfx, proj, attr = m.groups() if attr in ("weight", "bias", "linear.weight", "linear.bias"): bare_attr = attr.rsplit(".", 1)[-1] return f"{pfx}.qkv.{proj}.{bare_attr}" if attr.startswith("output_"): return f"{pfx}.qkv.{proj[0]}_{attr}" if attr.startswith("input_"): return f"{pfx}.qkv.{attr}" # Step 2: fused GateUp m = Gemma4ForConditionalGeneration._RE_TOWER_GATE_UP.match(name) if m: pfx, proj, attr = m.groups() short = proj.split("_")[0] # "gate" or "up" if attr in ("weight", "bias", "linear.weight", "linear.bias"): bare_attr = attr.rsplit(".", 1)[-1] return f"{pfx}.gate_up.{proj}.{bare_attr}" if attr.startswith("output_"): return f"{pfx}.gate_up.{short}_{attr}" if attr.startswith("input_"): return f"{pfx}.gate_up.{attr}" # Step 3: clippable wrapper (.weight → .linear.weight) if name.endswith(".weight") or name.endswith(".bias"): base, attr = name.rsplit(".", 1) alt = f"{base}.linear.{attr}" if alt in params_dict: return alt return name def _get_k_eq_v_layers(self) -> set: """Return set of layer indices where attention_k_eq_v applies (full-attention layers).""" text_config = self.config.text_config if not getattr(text_config, "attention_k_eq_v", False): return set() return { i for i, lt in enumerate(text_config.layer_types) if lt == "full_attention" } def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): k_eq_v_layers = self._get_k_eq_v_layers() num_experts = getattr(self.config.text_config, "num_experts", 0) or 0 expert_params_mapping = [ # (param_name, ckpt_weight_name, shard_ids) # gate_up_proj is fused [E, 2*I, H] — chunk into w1 (gate) + w3 (up) ("experts.w13_weight", "experts.gate_up_proj", ("w1", "w3")), ("experts.w2_weight", "experts.down_proj", ("w2",)), ] # Per-expert checkpoint format used by compressed-tensors / FP8 # (e.g. RedHatAI/*-FP8-Dynamic) and by ModelOpt NVFP4 # (e.g. nvidia/Gemma-4-*-NVFP4). Each expert is stored as a # separate key with shape (out, in): # experts..{gate,up,down}_proj.{weight,weight_scale, # weight_scale_2,input_scale} # `make_expert_params_mapping` emits tuples whose `weight_name` ends # in a trailing dot, so the standard `name.replace(weight_name, # param_name)` collapses every suffix uniformly to the fused # FusedMoE params (experts.w13_*, experts.w2_*). per_expert_params_mapping = ( FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=num_experts, ) if num_experts else [] ) params_dict = dict(self.named_parameters()) params_dict.update(dict(self.named_buffers())) non_persistent_buffers: Set[str] = set() for mod_name, mod in self.named_modules(): for buf_name in getattr(mod, "_non_persistent_buffers_set", set()): full = f"{mod_name}.{buf_name}" if mod_name else buf_name non_persistent_buffers.add(full) text_tie = getattr(self.config.text_config, "tie_word_embeddings", True) start_layer = self.language_model.start_layer end_layer = self.language_model.end_layer loaded_params: Set[str] = set() for name, loaded_weight in weights: if "embed_vision.embedding." in name or "embed_audio.embedding." in name: continue if self.audio_tower is None and ( "audio_tower." in name or "embed_audio." in name ): continue name = re.sub(r"^model\.", "", name) if pp_filter_load_weight( name, loaded_weight, pp_group=self.pp_group, start_layer=start_layer, end_layer=end_layer, params_dict=params_dict, loaded_params=loaded_params, tie_word_embeddings=text_tie, embed_weight_name="language_model.embed_tokens.weight", first_rank_only_patterns=( "language_model.embed_tokens", "language_model.per_layer_model_projection", "language_model.per_layer_projection_norm", "vision_tower.", "embed_vision.", "audio_tower.", "embed_audio.", ), last_rank_only_prefixes=("language_model.norm.", "lm_head."), ): continue # HF has router.per_expert_scale and experts.* on the decoder layer; # remap into our moe.* subtree since Gemma4MoE owns both. name = name.replace(".router.per_expert_scale", ".moe.per_expert_scale") if ".experts." in name and ".moe.experts." not in name: name = name.replace(".experts.", ".moe.experts.") # Remap audio tower checkpoint names to our module tree if "audio_tower." in name: name = self._remap_audio_tower_name(name) # Remap vision / audio tower names (fused QKV/GateUp, clippable wrappers) if "vision_tower." in name or "audio_tower." in name: name = self._remap_tower_name(name, params_dict) # attention_k_eq_v: full-attention layers have no v_proj in the # checkpoint (K and V share weights). When we see a k_proj weight # for one of these layers, load it into both the "k" and "v" shards # of the fused QKV so the forward produces v_raw == k_raw. should_dup_k_to_v = ( ".k_proj." in name and k_eq_v_layers and "language_model." in name and (m := re.search(r"layers\.(\d+)\.", name)) is not None and int(m.group(1)) in k_eq_v_layers ) # MoE expert weights checked first (gate_up_proj contains "up_proj" # which would false-match the stacked dense MLP mapping). orig_name = name # 1) Per-expert checkpoint layout (compressed-tensors FP8 like # RedHatAI/*-FP8-Dynamic, ModelOpt NVFP4 like # nvidia/Gemma-4-*-NVFP4): experts..{gate,up,down}_proj.* # The trailing dot in `weight_name` lets a single mapping fold # weight, weight_scale, weight_scale_2, and input_scale into # their corresponding fused FusedMoE params (experts.w13_*, # experts.w2_*). for ( param_name, weight_name, expert_id, shard_id, ) in per_expert_params_mapping: if weight_name not in orig_name: continue name = orig_name.replace(weight_name, param_name) if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader( param, loaded_weight, name, shard_id=shard_id, expert_id=expert_id, ) loaded_params.add(name) break else: # 2) BF16 fused checkpoint layout: experts.gate_up_proj is a # [E, 2*I, H] tensor that needs per-expert chunking into # w1 (gate) and w3 (up). for param_name, weight_name, shard_ids in expert_params_mapping: name = orig_name if weight_name not in name: continue name = name.replace(weight_name, param_name) if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader for i in range(num_experts): chunks = loaded_weight[i].chunk(len(shard_ids), dim=0) for chunk, sid in zip(chunks, shard_ids): weight_loader(param, chunk, name, sid, i) loaded_params.add(name) break else: for ( param_name, weight_name, shard_id, ) in self.stacked_params_mapping: name = orig_name if weight_name not in name: continue name = name.replace(weight_name, param_name) if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) if should_dup_k_to_v: weight_loader(param, loaded_weight, "v") loaded_params.add(name) break else: name = orig_name if name.endswith(".bias") and name not in params_dict: continue name = maybe_remap_kv_scale_name(name, params_dict) if name is None: continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(name) unloaded_params = params_dict.keys() - loaded_params if unloaded_params: param_names = set(dict(self.named_parameters()).keys()) buckets = { logging.WARNING: ( "Some weights are not initialized from checkpoints", lambda p: p in param_names, ), logging.INFO: ( "Persistent buffers not in checkpoint (using default init)", lambda p: p not in param_names and p not in non_persistent_buffers, ), logging.DEBUG: ( "Non-persistent buffers not in checkpoint (expected)", lambda p: p in non_persistent_buffers, ), } for level, (msg, pred) in buckets.items(): names = sorted(p for p in unloaded_params if pred(p)) if names: logger.log(level, "%s: %s", msg, names) return loaded_params lora_pattern = re.compile( r"^language_model\.layers\.(\d+)\.(?:self_attn|mlp)\.(?:qkv_proj|o_proj|down_proj|gate_up_proj)" ) def should_apply_lora(self, module_name: str) -> bool: return bool(self.lora_pattern.match(module_name)) def get_hidden_dim(self, module_name, layer_idx): # return input_dim, output_dim if module_name == "qkv_proj": return ( self.config.hidden_size, self.config.head_dim * ( self.config.num_attention_heads + self.config.num_key_value_heads * 2 ), ) elif module_name == "o_proj": return ( self.config.head_dim * self.config.num_attention_heads, self.config.hidden_size, ) elif module_name == "gate_up_proj": assert len(set(self.config.intermediate_size)) == 1, ( "Currently SGLang requires uniform intermediate size for all layers. " "Please file an issue if you need support for non-uniform intermediate sizes." ) return self.config.hidden_size, self.config.intermediate_size[0] * 2 elif module_name == "down_proj": assert len(set(self.config.intermediate_size)) == 1, ( "Currently SGLang requires uniform intermediate size for all layers. " "Please file an issue if you need support for non-uniform intermediate sizes." ) return self.config.intermediate_size[0], self.config.hidden_size else: raise NotImplementedError() def get_embed(self): return self.language_model.embed_tokens.weight def get_embed_and_head(self): if self.pp_group.world_size > 1: # Under PP, embed_tokens lives on the first rank and lm_head on the # last; neither rank holds both tensors, so we can't return the # pair locally without a cross-stage gather. Callers (RL weight # sync, remote weight loader) currently assume a single-rank view — # fail loudly rather than dereference a PPMissingLayer. raise NotImplementedError( "get_embed_and_head() is not implemented for Gemma4 " "multimodal under pipeline parallelism. embed_tokens lives " "on the first PP rank and lm_head on the last; use " "--pp-size 1 if you need this API." ) embed = self.language_model.embed_tokens.weight # Gemma4 ties word embeddings, so embed_tokens serves as lm_head return embed, embed def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): self.capture_aux_hidden_states = True text_config = self.config.text_config if layer_ids is None: num_layers = text_config.num_hidden_layers self.language_model.layers_to_capture = [ 2, num_layers // 2, num_layers - 3, ] else: # we plus 1 here because in sglang, for the ith layer, it takes the output # of the (i-1)th layer as aux hidden state self.language_model.layers_to_capture = [val + 1 for val in layer_ids] EntryClass = Gemma4ForConditionalGeneration