From f1bbd26ff79be9f0b89fddd2a3bbe8ed5cb31928 Mon Sep 17 00:00:00 2001 From: Stefan He Date: Sat, 13 Dec 2025 02:56:54 -0800 Subject: [PATCH] Clean up GDN Init (#14855) --- python/sglang/srt/models/qwen3_next.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 9cde50468..b38a83d57 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -7,7 +7,7 @@ from torch import nn from sglang.srt.compilation.piecewise_context_manager import get_forward_context from sglang.srt.configs.qwen3_next import Qwen3NextConfig -from sglang.srt.distributed import divide, get_pp_group +from sglang.srt.distributed import get_pp_group from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated @@ -221,7 +221,6 @@ class Qwen3GatedDeltaNet(nn.Module): self.activation = config.hidden_act self.layer_norm_epsilon = config.rms_norm_eps - # QKV self.conv_dim = self.key_dim * 2 + self.value_dim self.conv1d = ColumnParallelLinear( input_size=self.conv_kernel_size, @@ -232,7 +231,6 @@ class Qwen3GatedDeltaNet(nn.Module): tp_size=self.attn_tp_size, ) self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) - # projection of the input hidden states projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2 projection_size_ba = self.num_v_heads * 2 @@ -272,17 +270,11 @@ class Qwen3GatedDeltaNet(nn.Module): }, ) - # selective projection used to make dt, B and C input dependent + self.dt_bias = nn.Parameter(torch.zeros(self.num_v_heads // self.attn_tp_size)) - # time step projection (discretization) - # instantiate once and copy inv_dt in init_weights of PretrainedModel - self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads // self.attn_tp_size)) - - A = torch.empty( - divide(self.num_v_heads, self.attn_tp_size), dtype=torch.float32 - ).uniform_(0, 16) - self.A_log = nn.Parameter(torch.log(A)) - self.A_log._no_weight_decay = True + self.A_log = nn.Parameter( + torch.zeros(self.num_v_heads // self.attn_tp_size, dtype=torch.float32) + ) set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)}) set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})