# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # 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 copy from typing import Iterable, List, Optional, Set, Tuple import einops import torch from torch import nn from transformers import ( ROPE_INIT_FUNCTIONS, Gemma3TextConfig, PretrainedConfig, PreTrainedModel, ) from sglang.srt.layers.activation import GeluAndMul from sglang.srt.layers.layernorm import Gemma3RMSNorm from sglang.srt.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import AttentionType, RadixAttention from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb, get_rope from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu, make_layers _is_cpu = is_cpu() _is_cpu_amx_available = cpu_has_amx_support() # Aligned with HF's implementation, using sliding window inclusive with the last token # SGLang assumes exclusive def get_attention_sliding_window_size(config): return config.sliding_window - 1 # Adapted from: # https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma3.py def extract_layer_index(prefix: str) -> int: """Extract the layer index from a prefix string.""" parts = prefix.split(".") for part in parts: if part.startswith("layers."): layer_str = part.split(".")[-1] try: return int(layer_str) except ValueError: continue return -1 class Gemma3MLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_activation: str, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=add_prefix("gate_up_proj", prefix), ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, prefix=add_prefix("down_proj", prefix), ) if hidden_activation != "gelu_pytorch_tanh": raise ValueError( f"{self.__class__.__name__} uses `gelu_pytorch_tanh` as the hidden activation " "function. Please set `hidden_activation` to " "`gelu_pytorch_tanh`." ) self.act_fn = GeluAndMul() self.prefix = prefix def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x class Gemma3Attention(nn.Module): def __init__( self, layer_id: int, config: Gemma3TextConfig, max_position_embeddings: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.layer_id = layer_id self.config = config tp_size = get_parallel().tp_size self.total_num_heads = config.num_attention_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = config.num_key_value_heads self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) if self.total_num_kv_heads >= tp_size: # Number of KV heads is greater than TP size, so we partition # the KV heads across multiple tensor parallel GPUs. assert self.total_num_kv_heads % tp_size == 0 else: # Number of KV heads is less than TP size, so we replicate # the KV heads across multiple tensor parallel GPUs. assert tp_size % self.total_num_kv_heads == 0 hidden_size = config.hidden_size head_dim = getattr( config, "head_dim", hidden_size // config.num_attention_heads ) self.head_dim = head_dim partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) self.rotary_dim = int(partial_rotary_factor * self.head_dim) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = config.query_pre_attn_scalar**-0.5 self.qkv_proj = QKVParallelLinear( hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=config.attention_bias, quant_config=quant_config, prefix=add_prefix("qkv_proj", prefix), ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, hidden_size, bias=config.attention_bias, quant_config=quant_config, prefix=add_prefix("o_proj", prefix), ) self.is_sliding = config.layer_types[layer_id] == "sliding_attention" # In transformers v5, rope_parameters is nested per layer type: # {"sliding_attention": {"rope_theta": 10000}, "full_attention": {"rope_theta": 1000000}} # In v4 it was flat: {"rope_type": "default", "rope_theta": ...} rope_params = config.rope_parameters is_nested = isinstance(rope_params, dict) and "full_attention" in rope_params # Initialize the rotary embedding. if self.is_sliding: # Local attention. Override the values in config.json. if is_nested: self.rope_theta = rope_params["sliding_attention"].get( "rope_theta", 10000.0 ) else: self.rope_theta = getattr(config, "rope_local_base_freq", 10000.0) self.rope_scaling = {"rope_type": "default"} # FIXME(mick): idk why vllm does this # self.sliding_window = config.interleaved_sliding_window self.sliding_window = get_attention_sliding_window_size(config) else: # Global attention. Use the values in config.json. if is_nested: self.rope_theta = rope_params["full_attention"].get( "rope_theta", 1000000.0 ) else: self.rope_theta = ( rope_params.get("rope_theta", 10000.0) if rope_params else 10000.0 ) self.rope_scaling = {"rope_type": "default"} self.sliding_window = None self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.rotary_dim, max_position=max_position_embeddings, base=self.rope_theta, rope_scaling=self.rope_scaling, is_neox_style=getattr(config, "rope_is_neox_style", True), ) self.attn = RadixAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, layer_id=layer_id, logit_cap=0.0, # Module must also define `get_attention_sliding_window_size` to correctly initialize # attention backend in `ForwardBatch`. sliding_window_size=self.sliding_window, quant_config=quant_config, prefix=add_prefix("attn", prefix), attn_type=AttentionType.DECODER_BIDIRECTIONAL, ) # Gemma3 adds normalization for q and k self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) def forward_cpu( self, positions: torch.Tensor, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], forward_batch: ForwardBatch, **kwargs, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) # [s, h * head_dim] q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # [s, h, head_dim] q = q.unflatten(-1, (self.num_heads, self.head_dim)).unsqueeze(0) q = self.q_norm(q) k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0) k = self.k_norm(k) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v, forward_batch=forward_batch) # Compatible with triton backend which returns [1, s, h, head_dim] if attn_output.dim() == 4 and attn_output.shape[0] == 1: attn_output = attn_output.squeeze(0) attn_output = attn_output.flatten(-2, -1) # [s, h * head_dim] output, _ = self.o_proj(attn_output) return output def forward_native( self, positions: torch.Tensor, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], forward_batch: ForwardBatch, **kwargs, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) # [s, h * head_dim] q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # [s, h, head_dim] q = q.unflatten(-1, (self.num_heads, self.head_dim)) # -> [h, s, head_dim] q = q.transpose(0, 1).unsqueeze(0) q = self.q_norm(q) k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) # -> [h, s, head_dim] k = k.transpose(0, 1).unsqueeze(0) k = self.k_norm(k) # q, k = self.rotary_emb(positions, q, k) cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) # [b, h, s, head_dim] -> [b, s, h, head_dim] q = q.permute(0, 2, 1, 3) k = k.permute(0, 2, 1, 3) attn_output = self.attn(q, k, v, forward_batch=forward_batch) # Compatible with triton backend which returns [1, s, h, head_dim] if attn_output.dim() == 4 and attn_output.shape[0] == 1: attn_output = attn_output.squeeze(0) attn_output = attn_output.flatten(-2, -1) # [s, h * head_dim] output, _ = self.o_proj(attn_output) return output def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], forward_batch: ForwardBatch, **kwargs, ) -> torch.Tensor: if _is_cpu and _is_cpu_amx_available: return self.forward_cpu( positions, hidden_states, position_embeddings, forward_batch, **kwargs ) return self.forward_native( positions, hidden_states, position_embeddings, forward_batch, **kwargs ) class Gemma3DecoderLayer(nn.Module): def __init__( self, layer_id: int, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size self.self_attn = Gemma3Attention( layer_id=layer_id, config=config, max_position_embeddings=config.max_position_embeddings, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), ) self.hidden_size = config.hidden_size self.mlp = Gemma3MLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, hidden_activation=config.hidden_activation, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) self.input_layernorm = Gemma3RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.post_attention_layernorm = Gemma3RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.pre_feedforward_layernorm = Gemma3RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.post_feedforward_layernorm = Gemma3RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.is_sliding = self.self_attn.is_sliding self.layer_id = layer_id def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, position_embeddings_global: torch.Tensor, position_embeddings_local: torch.Tensor, forward_batch: ForwardBatch, **kwargs, ) -> tuple[ torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]] ]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # apply global RoPE to non-sliding layer only if self.self_attn.is_sliding: position_embeddings = position_embeddings_local else: position_embeddings = position_embeddings_global hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, position_embeddings=position_embeddings, forward_batch=forward_batch, **kwargs, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.pre_feedforward_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) return outputs class Gemma3RotaryEmbedding(nn.Module): def __init__(self, config: Gemma3TextConfig, device=None): super().__init__() # BC: "rope_type" was originally "type" rope_scaling = config.rope_parameters if rope_scaling is not None: self.rope_type = rope_scaling.get( "rope_type", rope_scaling.get("type", "default") ) else: self.rope_type = "default" if self.rope_type is None: self.rope_type = "default" self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config if self.rope_type == "default": self.rope_init_fn = self.compute_default_rope_parameters else: self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq def _dynamic_frequency_update(self, position_ids, device): """ dynamic RoPE layers should recompute `inv_freq` in the following situations: 1 - growing beyond the cached sequence length (allow scaling) 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) """ seq_len = torch.max(position_ids) + 1 if seq_len > self.max_seq_len_cached: # growth inv_freq, self.attention_scaling = self.rope_init_fn( self.config, device, seq_len=seq_len ) self.register_buffer( "inv_freq", inv_freq, persistent=False ) # TODO joao: may break with compilation self.max_seq_len_cached = seq_len if ( seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len ): # reset # This .to() is needed if the model has been moved to a device after being initialized (because # the buffer is automatically moved, but not the original copy) self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len @staticmethod def compute_default_rope_parameters(config, device=None, seq_len=None): """Standard RoPE: no scaling, just base frequency.""" rope_params = config.rope_parameters if isinstance(rope_params, dict) and "rope_theta" not in rope_params: # Nested per-layer-type format; pick the first available theta for v in rope_params.values(): if isinstance(v, dict) and "rope_theta" in v: base = v["rope_theta"] break else: base = 10000.0 else: base = rope_params.get("rope_theta", 10000.0) if rope_params else 10000.0 dim = ( getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads ) inv_freq = 1.0 / ( base ** ( torch.arange(0, dim, 2, dtype=torch.int64).to( device=device, dtype=torch.float ) / dim ) ) return inv_freq, 1.0 @torch.no_grad() def forward(self, x, position_ids): if "dynamic" in self.rope_type: self._dynamic_frequency_update(position_ids, device=x.device) # Core RoPE block inv_freq_expanded = ( self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) ) position_ids_expanded = position_ids[:, None, :].float() # Force float32 (see https://github.com/huggingface/transformers/pull/29285) device_type = x.device.type device_type = ( device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" ) with torch.autocast(device_type=device_type, enabled=False): freqs = ( inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float() ).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention cos = cos * self.attention_scaling sin = sin * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class Gemma3TextScaledWordEmbedding(nn.Embedding): """ This module overrides nn.Embeddings' forward by multiplying with embeddings scale. """ def __init__( self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0, ): super().__init__(num_embeddings, embedding_dim, padding_idx) self.embed_scale = embed_scale def forward(self, input_ids: torch.Tensor): return super().forward(input_ids) * self.embed_scale class Gemma3TextModel(PreTrainedModel): def __init__( self, config: Gemma3TextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__(config=config) self.config = config self.quant_config = quant_config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size # Gemma3 downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402 self.embed_tokens = Gemma3TextScaledWordEmbedding( config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=self.config.hidden_size**0.5, ) self.norm = Gemma3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # In transformers v5, rope_parameters is nested per layer type: # {"sliding_attention": {"rope_type": ..., "rope_theta": 10000}, # "full_attention": {"rope_type": ..., "rope_theta": 1000000}} # Flatten into the format Gemma3RotaryEmbedding expects. rope_params = config.rope_parameters if isinstance(rope_params, dict) and "full_attention" in rope_params: global_theta = rope_params["full_attention"].get("rope_theta", 1000000.0) local_theta = rope_params["sliding_attention"].get("rope_theta", 10000.0) else: # v4 flat format fallback global_theta = ( rope_params.get("rope_theta", 10000.0) if rope_params else 10000.0 ) local_theta = getattr(config, "rope_local_base_freq", 10000.0) global_config = copy.deepcopy(config) global_config.rope_parameters = { **rope_params["full_attention"], "rope_theta": global_theta, } self.rotary_emb = Gemma3RotaryEmbedding(config=global_config) self.gradient_checkpointing = False local_config = copy.deepcopy(config) local_config.rope_parameters = { "rope_type": "default", "rope_theta": local_theta, } self.rotary_emb_local = Gemma3RotaryEmbedding(config=local_config) self.layers = make_layers( config.num_hidden_layers, lambda idx, prefix: Gemma3DecoderLayer( layer_id=idx, config=config, quant_config=quant_config, prefix=prefix, ), prefix=add_prefix("layers", prefix), ) self.norm = Gemma3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.layers_to_capture = [] self.post_init() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, **kwargs, ) -> torch.Tensor: if input_embeds is None: hidden_states = self.embed_tokens(input_ids) else: hidden_states = input_embeds aux_hidden_states = [] num_layers = len(self.layers) if _is_cpu and _is_cpu_amx_available: for i, layer in enumerate(self.layers): if i in self.layers_to_capture: aux_hidden_states.append(hidden_states) layer_outputs = layer( positions=positions, position_embeddings_global=None, position_embeddings_local=None, hidden_states=hidden_states, forward_batch=forward_batch, **kwargs, ) hidden_states = layer_outputs[0] else: if positions.dim() == 1: positions = einops.rearrange(positions, "s -> 1 s") position_embeddings_global = self.rotary_emb(hidden_states, positions) position_embeddings_local = self.rotary_emb_local(hidden_states, positions) for i, layer in enumerate(self.layers): if i in self.layers_to_capture: aux_hidden_states.append(hidden_states) layer_outputs = layer( positions=positions, position_embeddings_global=position_embeddings_global, position_embeddings_local=position_embeddings_local, hidden_states=hidden_states, forward_batch=forward_batch, **kwargs, ) hidden_states = layer_outputs[0] # Capture the output of the last layer if requested. # layers_to_capture uses +1 offset (captures input of layer i = output of i-1), # so index num_layers means the output of the final layer. if num_layers in self.layers_to_capture: aux_hidden_states.append(hidden_states) hidden_states = self.norm(hidden_states) if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states class Gemma3ForCausalLM(PreTrainedModel): config_class = Gemma3TextConfig _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config_class = Gemma3TextConfig base_model_prefix = "language_model" # 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 = { # shard_name, weight_name, index "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: Gemma3TextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__(config=config) self.config = config self.quant_config = quant_config self.model = Gemma3TextModel( config, quant_config, prefix=add_prefix("model", prefix) ) self.logits_processor = LogitsProcessor(config) if self.config.tie_word_embeddings: self.lm_head = self.model.embed_tokens else: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), ) self.capture_aux_hidden_states = False self.post_init() def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens def get_attention_sliding_window_size(self): return get_attention_sliding_window_size(self.config) def dtype(self) -> torch.dtype: return next(self.parameters()).dtype @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, **kwargs, ) -> LogitsProcessor: hidden_states = self.model( input_ids, positions, forward_batch, input_embeds, **kwargs ) aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states return self.logits_processor( input_ids, hidden_states, self.model.embed_tokens, forward_batch, aux_hidden_states, ) @torch.no_grad() def forward_split_prefill( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, split_interval: Tuple[int, int], # [start, end) 0-based input_embeds: torch.Tensor = None, ): start, end = split_interval # embed if start == 0: if input_embeds is None: hidden_states = self.model.embed_tokens(input_ids) else: hidden_states = input_embeds if positions.dim() == 1: positions = einops.rearrange(positions, "s -> 1 s") position_embeddings_global = self.model.rotary_emb(hidden_states, positions) position_embeddings_local = self.model.rotary_emb_local( hidden_states, positions ) forward_batch.hidden_states = hidden_states forward_batch.model_specific_states = { "positions": positions, "position_embeddings_global": position_embeddings_global, "position_embeddings_local": position_embeddings_local, } # decoder layer for i in range(start, end): layer = self.model.layers[i] layer_output = layer( positions=forward_batch.model_specific_states["positions"], position_embeddings_global=forward_batch.model_specific_states[ "position_embeddings_global" ], position_embeddings_local=forward_batch.model_specific_states[ "position_embeddings_local" ], hidden_states=forward_batch.hidden_states, forward_batch=forward_batch, ) forward_batch.hidden_states = layer_output[0] if end == self.model.config.num_hidden_layers: # norm forward_batch.hidden_states = self.model.norm(forward_batch.hidden_states) # logits process result = self.logits_processor( input_ids, forward_batch.hidden_states, self.model.embed_tokens, forward_batch, ) else: result = None return result def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): 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", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), ] params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() for name, loaded_weight in weights: remapped_name = maybe_remap_kv_scale_name(name, params_dict) if remapped_name is None: continue if remapped_name != name: param = params_dict[remapped_name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(remapped_name) continue for param_name, shard_name, shard_id in stacked_params_mapping: # if param_name in name: # print(f"{param_name} is already in {name}") if shard_name not in name: continue name = name.replace(shard_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: # lm_head is not used in vllm as it is tied with embed_token. # To prevent errors, skip loading lm_head.weight. if "lm_head.weight" in name: continue # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue # Remapping the name of FP8 kv-scale. name = maybe_remap_kv_scale_name(name, params_dict) if name is None: 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: # logger.warning( # "Some weights are not initialized from checkpoints: %s", unloaded_params # ) return loaded_params def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): if layer_ids is None: self.capture_aux_hidden_states = True num_layers = self.config.num_hidden_layers self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] else: self.capture_aux_hidden_states = True # 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.model.layers_to_capture = [val + 1 for val in layer_ids] def _shard_weight(self, weight: torch.Tensor) -> torch.Tensor: """Shard a full embedding/lm_head weight along vocab dim for the current TP rank. Gemma3 uses nn.Embedding (unsharded) but the Eagle3 draft model uses VocabParallelEmbedding (sharded). This method extracts the correct shard so the weights can be shared. """ tp_size = get_parallel().tp_size if tp_size <= 1: return weight tp_rank = get_parallel().tp_rank shard_size = (weight.shape[0] + tp_size - 1) // tp_size return weight[tp_rank * shard_size : (tp_rank + 1) * shard_size] def get_embed(self): return self._shard_weight(self.model.embed_tokens.weight) def get_embed_and_head(self): embed = self._shard_weight(self.model.embed_tokens.weight) head = self._shard_weight(self.lm_head.weight) return embed, head EntryClass = Gemma3ForCausalLM