diff --git a/python/sglang/srt/models/gemma3n_causal.py b/python/sglang/srt/models/gemma3n_causal.py index 0f710b0f8..c92f70971 100644 --- a/python/sglang/srt/models/gemma3n_causal.py +++ b/python/sglang/srt/models/gemma3n_causal.py @@ -12,6 +12,7 @@ from sglang.srt.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor @@ -183,21 +184,21 @@ class Gemma3nAltUp(nn.Module): self.correct_output_scale = nn.Parameter( torch.zeros(config.hidden_size, dtype=torch.float32) ) - self.correction_coefs = ColumnParallelLinear( + self.correction_coefs = ReplicatedLinear( config.altup_num_inputs, config.altup_num_inputs, bias=False, quant_config=quant_config, prefix=add_prefix("correction_coefs", prefix), ) - self.prediction_coefs = ColumnParallelLinear( + self.prediction_coefs = ReplicatedLinear( config.altup_num_inputs, config.altup_num_inputs**2, bias=False, quant_config=quant_config, prefix=add_prefix("prediction_coefs", prefix), ) - self.modality_router = ColumnParallelLinear( + self.modality_router = ReplicatedLinear( config.hidden_size, config.altup_num_inputs, bias=False, @@ -545,14 +546,14 @@ class Gemma3nDecoderLayer(nn.Module): config, quant_config, prefix=add_prefix("laurel", prefix) ) - self.per_layer_input_gate = ColumnParallelLinear( + self.per_layer_input_gate = ReplicatedLinear( self.hidden_size, self.hidden_size_per_layer_input, bias=False, quant_config=quant_config, prefix=add_prefix("per_layer_input_gate", prefix), ) - self.per_layer_projection = RowParallelLinear( + self.per_layer_projection = ReplicatedLinear( self.hidden_size_per_layer_input, self.hidden_size, bias=False, @@ -677,6 +678,7 @@ class Gemma3nTextModel(PreTrainedModel): self.hidden_size, config.num_hidden_layers * config.hidden_size_per_layer_input, bias=False, + gather_output=True, quant_config=quant_config, prefix=add_prefix("per_layer_model_projection", prefix), ) @@ -692,6 +694,7 @@ class Gemma3nTextModel(PreTrainedModel): self.hidden_size, self.hidden_size, bias=False, + gather_output=True, quant_config=quant_config, prefix=prefix, ), @@ -704,6 +707,7 @@ class Gemma3nTextModel(PreTrainedModel): self.hidden_size, self.hidden_size, bias=False, + gather_output=True, quant_config=quant_config, prefix=prefix, ), @@ -782,9 +786,6 @@ class Gemma3nTextModel(PreTrainedModel): per_layer_inputs = self.project_per_layer_inputs(input_embeds, per_layer_inputs) - if positions.dim() == 1: - positions = positions.unsqueeze(0) - # Expand hidden_states to support per-layer inputs target_magnitude = torch.mean(input_embeds**2, dim=-1, keepdim=True) ** 0.5 epsilon_tensor = torch.tensor(torch.finfo(input_embeds.dtype).min) diff --git a/python/sglang/srt/models/gemma3n_mm.py b/python/sglang/srt/models/gemma3n_mm.py index 86f7fd516..e2dfe99cc 100644 --- a/python/sglang/srt/models/gemma3n_mm.py +++ b/python/sglang/srt/models/gemma3n_mm.py @@ -14,7 +14,7 @@ from transformers import ( ) from transformers.models.auto.modeling_auto import AutoModel -from sglang.srt.layers.linear import RowParallelLinear +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.vocab_parallel_embedding import VocabParallelEmbedding @@ -90,7 +90,7 @@ class Gemma3nMultimodalEmbedder(nn.Module): eps=self.eps, ) - self.embedding_projection = RowParallelLinear( + self.embedding_projection = ReplicatedLinear( self.multimodal_hidden_size, self.text_hidden_size, bias=False, diff --git a/test/registered/eval/test_vlms_mmmu_eval.py b/test/registered/eval/test_vlms_mmmu_eval.py index 099694c09..40156a019 100644 --- a/test/registered/eval/test_vlms_mmmu_eval.py +++ b/test/registered/eval/test_vlms_mmmu_eval.py @@ -34,7 +34,9 @@ MODEL_THRESHOLDS = { 0.270, 23.8 ), ModelLaunchSettings("google/gemma-3-4b-it"): ModelEvalMetrics(0.360, 10.9), - ModelLaunchSettings("google/gemma-3n-E4B-it"): ModelEvalMetrics(0.270, 17.7), + ModelLaunchSettings( + "google/gemma-3n-E4B-it", extra_args=["--tp=2"] + ): ModelEvalMetrics(0.270, 17.7), ModelLaunchSettings("mistral-community/pixtral-12b"): ModelEvalMetrics(0.360, 16.6), ModelLaunchSettings("moonshotai/Kimi-VL-A3B-Instruct"): ModelEvalMetrics( 0.330, 23.5 @@ -50,7 +52,7 @@ MODEL_THRESHOLDS = { ModelLaunchSettings( "unsloth/Mistral-Small-3.1-24B-Instruct-2503" ): ModelEvalMetrics(0.30, 16.7), - ModelLaunchSettings("XiaomiMiMo/MiMo-VL-7B-RL"): ModelEvalMetrics(0.28, 32.0), + ModelLaunchSettings("XiaomiMiMo/MiMo-VL-7B-RL"): ModelEvalMetrics(0.28, 40.0), ModelLaunchSettings("zai-org/GLM-4.1V-9B-Thinking"): ModelEvalMetrics(0.280, 30.4), ModelLaunchSettings( "zai-org/GLM-4.5V-FP8", extra_args=["--tp=2"]