diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 9c04746d2..ce6be6ef8 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -14,7 +14,7 @@ import logging import re -from typing import Iterable, List, Optional, Set, Tuple +from typing import Iterable, List, Optional, Set, Tuple, Union import torch from torch import nn @@ -25,6 +25,7 @@ from transformers import ( ) from sglang.srt.distributed import ( + get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) @@ -45,8 +46,9 @@ from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead -from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -68,6 +70,59 @@ Gemma4MLP = Gemma3MLP Gemma4TextScaledWordEmbedding = Gemma3TextScaledWordEmbedding +def pp_filter_load_weight( + name, + loaded_weight, + *, + pp_group, + start_layer, + end_layer, + params_dict, + loaded_params, + tie_word_embeddings, + embed_weight_name, + first_rank_only_patterns=(), + last_rank_only_prefixes=(), + head_param_name="lm_head.weight", +): + """Shared PP filter for Gemma4 load_weights paths. + + Returns True if the caller should ``continue`` (handled or skipped), + False otherwise. No-op when ``pp_group.world_size == 1``. + + Handles three concerns in order: + 1. Drop transformer-layer weights outside [start_layer, end_layer). + 2. Route the tied ``embed_tokens.weight`` to ``lm_head`` on the last + rank (under PP, embed and lm_head live on different ranks so they + can't be tied via module aliasing). + 3. Skip rank-local module weights on the wrong rank. + """ + if pp_group.world_size <= 1: + return False + + layer_id = get_layer_id(name) + if layer_id is not None and (layer_id < start_layer or layer_id >= end_layer): + return True + + if tie_word_embeddings and pp_group.is_last_rank and name == embed_weight_name: + head_param = params_dict.get(head_param_name) + if head_param is not None: + wl = getattr(head_param, "weight_loader", default_weight_loader) + wl(head_param, loaded_weight) + loaded_params.add(head_param_name) + return True + + if not pp_group.is_first_rank and any(p in name for p in first_rank_only_patterns): + return True + + if not pp_group.is_last_rank and any( + name.startswith(p) for p in last_rank_only_prefixes + ): + return True + + return False + + class Gemma4Router(nn.Module): """Router for Gemma4 MoE that preprocesses input before projection. @@ -673,15 +728,12 @@ class Gemma4TextModel(PreTrainedModel): self.quant_config = quant_config self.vocab_size = config.vocab_size self.padding_idx = getattr(config, "pad_token_id", None) + self.pp_group = get_pp_group() - self.embed_tokens = Gemma4TextScaledWordEmbedding( - config.vocab_size, - config.hidden_size, - self.padding_idx, - embed_scale=self.config.hidden_size**0.5, # embedded normalizer - ) - - # Per-layer input embeddings + # Token / per-layer embedding tables and the per-layer projection only + # produce activations consumed at the model entry, so they live on the + # first PP rank only. Other ranks substitute PPMissingLayer so that + # parameter iteration still works (load_weights skips them explicitly). self.hidden_size = config.hidden_size self.hidden_size_per_layer_input = ( getattr(config, "hidden_size_per_layer_input", None) or 0 @@ -690,7 +742,43 @@ class Gemma4TextModel(PreTrainedModel): getattr(config, "vocab_size_per_layer_input", None) or config.vocab_size ) - if self.hidden_size_per_layer_input and self.hidden_size_per_layer_input > 0: + # PLE-enabled variants (E2B/E4B) forward `per_layer_inputs` through + # the PP proxy, but cuda_graph_runner hardcodes the proxy schema to + # {hidden_states, residual} and silently drops any extra keys at + # replay time. Empirically this corrupts E4B output to garbage on + # non-first PP ranks (eager path produces correct output and + # GSM8K ~0.92, cuda-graph path emits token soup). Refuse the + # combination until the runner becomes schema-aware; users can run + # PP + PLE eagerly with --disable-cuda-graph. + if self.pp_group.world_size > 1 and self.hidden_size_per_layer_input > 0: + sa = get_global_server_args() + if sa is not None and not sa.disable_cuda_graph: + raise ValueError( + "Pipeline parallelism is currently incompatible with " + "per-layer-input (PLE) embeddings under CUDA graph: " + "the runner's PP proxy schema is hardcoded to " + "{hidden_states, residual} and silently drops " + "per_layer_inputs, corrupting per-layer contributions on " + "non-first PP ranks. Workarounds: (a) pass " + "--disable-cuda-graph to fall back to eager replay, or " + "(b) use tensor parallelism (--tp-size) instead of PP." + ) + + if self.pp_group.is_first_rank: + self.embed_tokens = Gemma4TextScaledWordEmbedding( + config.vocab_size, + config.hidden_size, + self.padding_idx, + embed_scale=self.config.hidden_size**0.5, # embedded normalizer + ) + else: + self.embed_tokens = PPMissingLayer() + + if ( + self.pp_group.is_first_rank + and self.hidden_size_per_layer_input + and self.hidden_size_per_layer_input > 0 + ): self.embed_tokens_per_layer = Gemma4TextScaledWordEmbedding( self.vocab_size_per_layer_input, config.num_hidden_layers * self.hidden_size_per_layer_input, @@ -721,7 +809,7 @@ class Gemma4TextModel(PreTrainedModel): self.per_layer_input_scale = None self.per_layer_projection_scale = None - self.layers = make_layers( + self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: Gemma4DecoderLayer( layer_id=idx, @@ -729,10 +817,15 @@ class Gemma4TextModel(PreTrainedModel): quant_config=quant_config, prefix=prefix, ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() self.layers_to_capture = [] self.post_init() @@ -817,24 +910,35 @@ class Gemma4TextModel(PreTrainedModel): forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, per_layer_inputs: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, **kwargs, - ) -> torch.Tensor: - if (input_ids is None) ^ (input_embeds is not None): - raise ValueError( - "You must specify exactly one of input_ids or inputs_embeds" + ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]], PPProxyTensors]: + if self.pp_group.is_first_rank: + if (input_ids is None) ^ (input_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + if input_ids is not None: + input_embeds = self.embed_tokens(input_ids) + per_layer_inputs = self.get_per_layer_inputs(input_ids) + per_layer_inputs = self.project_per_layer_inputs( + input_embeds, per_layer_inputs ) - - if input_ids is not None: - input_embeds = self.embed_tokens(input_ids) - per_layer_inputs = self.get_per_layer_inputs(input_ids) - per_layer_inputs = self.project_per_layer_inputs(input_embeds, per_layer_inputs) - - hidden_states = input_embeds + hidden_states = input_embeds + else: + assert ( + pp_proxy_tensors is not None + ), "pp_proxy_tensors is required on non-first PP ranks" + hidden_states = pp_proxy_tensors["hidden_states"] + # PLE inputs were computed on rank 0 and forwarded along the + # pipeline; non-PLE models simply omit the key. + per_layer_inputs = pp_proxy_tensors.tensors.get("per_layer_inputs", None) aux_hidden_states = [] - num_layers = len(self.layers) + num_layers = self.config.num_hidden_layers - for layer_idx, layer in enumerate(self.layers): + for layer_idx in range(self.start_layer, self.end_layer): if layer_idx in self.layers_to_capture: aux_hidden_states.append(hidden_states) @@ -842,6 +946,7 @@ class Gemma4TextModel(PreTrainedModel): per_layer_input = per_layer_inputs[:, layer_idx, :] else: per_layer_input = None + layer = self.layers[layer_idx] layer_outputs = layer( positions=positions, hidden_states=hidden_states, @@ -850,7 +955,23 @@ class Gemma4TextModel(PreTrainedModel): **kwargs, ) hidden_states = layer_outputs[0] - residual = layer_outputs[1] if len(layer_outputs) > 1 else None + # Gemma4DecoderLayer.forward always returns (hidden_states, None); + # the residual is fused inside the layer, so nothing to thread. + + if not self.pp_group.is_last_rank: + # cuda_graph_runner allocates a fixed PP-proxy schema of + # {hidden_states, residual} and KeyErrors if a model omits a key. + # Gemma4 fuses the residual inside each layer so we don't have a + # standalone tensor to forward; emit a zero placeholder instead so + # graph replay can still copy it. The receiving stage never reads + # this key. + proxy = { + "hidden_states": hidden_states, + "residual": torch.zeros_like(hidden_states), + } + if per_layer_inputs is not None: + proxy["per_layer_inputs"] = per_layer_inputs + return PPProxyTensors(proxy) # Capture the output of the last layer if requested. # layers_to_capture uses +1 offset, so num_layers means @@ -858,10 +979,7 @@ class Gemma4TextModel(PreTrainedModel): if num_layers in self.layers_to_capture: aux_hidden_states.append(hidden_states) - if residual is None: - hidden_states = self.norm(hidden_states) - else: - hidden_states, _ = self.norm(hidden_states, residual) + hidden_states = self.norm(hidden_states) if len(aux_hidden_states) == 0: return hidden_states @@ -918,25 +1036,46 @@ class Gemma4ForCausalLM(PreTrainedModel): prefix: str = "", ) -> None: super().__init__(config=config) + self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config + self.model = Gemma4TextModel( config=config, quant_config=quant_config, prefix=add_prefix("model", prefix) ) self.logits_processor = LogitsProcessor(config) - if self.config.tie_word_embeddings: + # tie_word_embeddings ties lm_head to embed_tokens, but with PP those + # tensors live on opposite ranks (first vs last). In the PP > 1 case + # we materialize a real ParallelLMHead on the last rank and route the + # checkpoint's embed_tokens.weight into it during load_weights. + if self.pp_group.world_size == 1 and self.config.tie_word_embeddings: self.lm_head = self.model.embed_tokens - else: + elif self.pp_group.is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), ) + else: + self.lm_head = PPMissingLayer() + self.capture_aux_hidden_states = False self.post_init() + def tie_weights(self, *args, **kwargs): + # HF's PreTrainedModel.tie_weights uses ``_tied_weights_keys`` to bind + # ``lm_head.weight`` to ``model.embed_tokens.weight``. Under PP those + # tensors live on different ranks (embed on first, head on last) and + # the missing side is a PPMissingLayer with no ``weight`` attribute, + # which makes the default tie_weights crash. load_weights routes the + # checkpoint embedding into lm_head explicitly, so the tie is a no-op + # here when PP is active. + if self.pp_group.world_size > 1: + return + super().tie_weights(*args, **kwargs) + def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens @@ -957,17 +1096,24 @@ class Gemma4ForCausalLM(PreTrainedModel): forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, per_layer_inputs: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, **kwargs, - ) -> LogitsProcessor: + ) -> Union[LogitsProcessor, PPProxyTensors]: hidden_states = self.model( input_ids, positions, forward_batch, input_embeds, per_layer_inputs, + pp_proxy_tensors=pp_proxy_tensors, **kwargs, ) + if not self.pp_group.is_last_rank: + # `hidden_states` here is actually a PPProxyTensors handed off to + # the next stage; logits processing only happens on the last rank. + return hidden_states + aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states @@ -1022,6 +1168,25 @@ class Gemma4ForCausalLM(PreTrainedModel): if ".experts." in name and ".moe.experts." not in name: name = name.replace(".experts.", ".moe.experts.") + if pp_filter_load_weight( + name, + loaded_weight, + pp_group=self.pp_group, + start_layer=self.model.start_layer, + end_layer=self.model.end_layer, + params_dict=params_dict, + loaded_params=loaded_params, + tie_word_embeddings=self.config.tie_word_embeddings, + embed_weight_name="model.embed_tokens.weight", + first_rank_only_patterns=( + "embed_tokens", + "per_layer_model_projection", + "per_layer_projection_norm", + ), + last_rank_only_prefixes=("model.norm.", "lm_head."), + ): + continue + # 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 @@ -1120,6 +1285,19 @@ class Gemma4ForCausalLM(PreTrainedModel): return self._shard_weight(self.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 Gemma4ForCausalLM " + "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._shard_weight(self.model.embed_tokens.weight) head = self._shard_weight(self.lm_head.weight) return embed, head diff --git a/python/sglang/srt/models/gemma4_mm.py b/python/sglang/srt/models/gemma4_mm.py index 4ce8a5909..6d82e31cf 100644 --- a/python/sglang/srt/models/gemma4_mm.py +++ b/python/sglang/srt/models/gemma4_mm.py @@ -28,11 +28,14 @@ from transformers import ( PreTrainedModel, ) +from sglang.srt.distributed import get_pp_group 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.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, @@ -43,13 +46,17 @@ from sglang.srt.managers.schedule_batch import ( MultimodalInputs, flatten_nested_list, ) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + ForwardMode, + PPProxyTensors, +) 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 +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 @@ -170,38 +177,46 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): 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) - 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), - ) - - # Audio components - if getattr(config, "audio_config", None) is not None: - self.audio_tower = Gemma4AudioEncoder( - config=config.audio_config, + # 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("audio_tower", prefix), + prefix=add_prefix("vision_tower", prefix), ) - self.embed_audio = Gemma4MultimodalEmbedder( - config.audio_config, + self.embed_vision = Gemma4MultimodalEmbedder( + config.vision_config, config.text_config, quant_config=quant_config, - prefix=add_prefix("embed_audio", prefix), + 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 @@ -212,13 +227,31 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): config.text_config.vocab_size, ) - # Text model + # 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 @@ -548,17 +581,26 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, **kwargs: object, - ) -> LogitsProcessor: + ) -> Union[LogitsProcessor, PPProxyTensors]: """Forward pass for multimodal Gemma4.""" - if (input_ids is None) ^ (input_embeds is not None): + 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" ) positions += 1 per_layer_inputs = None - if input_ids is not 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 ple_ids[input_ids == self.config.image_token_id] = pad_id @@ -567,9 +609,8 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): per_layer_inputs = self.get_per_layer_inputs(ple_ids) # Prepare bidirectional attention masks for image tokens during prefill. - # Gemma 4 uses bidirectional attention for image soft tokens. - # Only TritonAttnBackend supports this; incompatible with CUDA Graph and - # chunked 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() @@ -580,7 +621,8 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): mask_dtype=torch.bool, ) - # Use general_mm_embed_routine for handling multimodal data + # 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, @@ -592,24 +634,43 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): }, 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 - # Process hidden states through logits processor + # 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, - self.language_model.embed_tokens, + 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 @@ -764,6 +825,10 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): 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: @@ -776,6 +841,29 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): 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") @@ -948,6 +1036,18 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): 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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9c5e2822f..af4f2a401 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2194,7 +2194,10 @@ class ServerArgs: f"Disable hybrid SWA memory for {model_arch} as it is not yet supported." ) self.disable_hybrid_swa_memory = True - elif model_arch == "Gemma4ForConditionalGeneration": + elif model_arch in ( + "Gemma4ForConditionalGeneration", + "Gemma4ForCausalLM", + ): default_attention_backend = ( "trtllm_mha" if is_sm100_supported() else "triton" ) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 6ef153bd8..f34c73a4d 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -69,6 +69,8 @@ DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST = "Qwen/Qwen3-Next-80B-A3B-Instruct" # VL test models DEFAULT_MODEL_NAME_FOR_TEST_VL_PP = "Qwen/Qwen3-VL-2B-Thinking" DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP = "zai-org/GLM-4.1V-9B-Thinking" +DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP = "google/gemma-4-26B-A4B-it" +DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP = "google/gemma-4-E4B-it" # NVFP4 models DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST = "nvidia/DeepSeek-V3-0324-FP4" diff --git a/test/registered/distributed/test_pp_single_node.py b/test/registered/distributed/test_pp_single_node.py index 85dafad5c..47942db38 100644 --- a/test/registered/distributed/test_pp_single_node.py +++ b/test/registered/distributed/test_pp_single_node.py @@ -4,6 +4,9 @@ python3 -m unittest test_pp_single_node.TestPPAccuracy.test_gsm8k python3 -m unittest test_pp_single_node.TestQwenPPAccuracy.test_pp_consistency python3 -m unittest test_pp_single_node.TestFixedBugs.test_chunked_prefill_with_small_bs python3 -m unittest test_pp_single_node.TestQwenVLPPAccuracy.test_mmmu +python3 -m unittest test_pp_single_node.TestGemma4PPAccuracy.test_gsm8k +python3 -m unittest test_pp_single_node.TestGemma4PPAccuracy.test_mmmu +python3 -m unittest test_pp_single_node.TestGemma4PLEPPAccuracy.test_gsm8k python3 -m unittest test_pp_single_node.TestPPMixedChunk.test_gsm8k """ @@ -21,6 +24,8 @@ from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( DEFAULT_MLA_MODEL_NAME_FOR_TEST, DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP, + DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP, DEFAULT_MODEL_NAME_FOR_TEST_GLM_41V_PP, DEFAULT_MODEL_NAME_FOR_TEST_VL_PP, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -202,6 +207,137 @@ class TestQwenVLPPAccuracy(unittest.TestCase): self.assertGreater(metrics["score"], 0.26) +@unittest.skipIf( + is_in_amd_ci(), + "Gemma4 PP not yet validated on AMD", +) +class TestGemma4PPAccuracy(unittest.TestCase): + """End-to-end PP=2 accuracy gate for Gemma4 multimodal. + + Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so + sglang auto-selects the triton attention backend; no manual flag needed. + The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an + H100's 80 GB. + """ + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP + cls.base_url = "http://127.0.0.1:23333" + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + 1, + "--pp-size", + 2, + "--trust-remote-code", + "--enable-multimodal", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + # Gemma4 is instruction-tuned and doesn't follow few-shot completion + # prompts well — use the chat API (default in run_eval), which scores + # ~0.98 on this model vs ~0.44 with api="completion". + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=200, + num_threads=32, + ) + metrics = run_eval(args) + print(f"{metrics=}") + + # Chat-API baseline ~0.98; gate well below to absorb sample-noise + # without missing a real PP-routing regression (pre-PP-fix the model + # produced garbage outputs scoring ≈ 0). + self.assertGreaterEqual(metrics["score"], 0.90) + # Wait a little bit so that the memory check happens. + time.sleep(4) + + @unittest.skipIf(is_in_ci(), "To reduce the CI execution time.") + def test_mmmu(self): + # Multimodal accuracy gate covering the vision_tower → embed_vision + # (first rank) → PP-proxy handoff → LM tail (last rank) chain. + # Measured 0.71 on 200 examples; full eval (~900 questions) takes + # ~5-7 min on H100 so this is manual-only. + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mmmu", + num_examples=None, + num_threads=32, + ) + metrics = run_eval(args) + print(f"{metrics=}") + # Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in + # 0.69-0.73. Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary + # samples ≈ 0.015) while still catching mid-grade vision/PP + # regressions, not just complete breakage. + self.assertGreater(metrics["score"], 0.65) + + +@unittest.skipIf( + is_in_amd_ci(), + "Gemma4 PP not yet validated on AMD", +) +class TestGemma4PLEPPAccuracy(unittest.TestCase): + """PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path). + + 26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP + test never crosses the PLE branch. Cuda graph + PLE corrupts outputs + (the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema + drops ``per_layer_inputs``), so this test pins the eager configuration. + """ + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP + cls.base_url = "http://127.0.0.1:23339" + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + 1, + "--pp-size", + 2, + "--trust-remote-code", + "--enable-multimodal", + # Required for PLE under PP — see Gemma4TextModel guard. + "--disable-cuda-graph", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + # Eager-path baseline ~0.92; gate 0.80 catches PLE breakage + # (corruption collapses score to ~0). + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=100, + num_threads=32, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], 0.80) + time.sleep(4) + + class TestQwenPPAccuracy(unittest.TestCase): @classmethod def setUpClass(cls):