Add output_gate_type to Qwen3NextConfig and update models to utilize it (#25401)

This commit is contained in:
Gaoji Liu
2026-05-19 00:18:08 +08:00
committed by GitHub
parent d96e593fd0
commit 3e2a109636
3 changed files with 21 additions and 1 deletions
+4
View File
@@ -68,6 +68,8 @@ class Qwen3NextConfig(PretrainedConfig):
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
hidden_act (`str`, *optional*, defaults to `"silu"`):
The non-linear activation function in the decoder.
output_gate_type (`str`, *optional*, defaults to `None`):
The gate activation function used by the linear attention output norm.
max_position_embeddings (`int`, *optional*, defaults to 32768):
The maximum sequence length that this model might ever be used with.
initializer_range (`float`, *optional*, defaults to 0.02):
@@ -186,6 +188,7 @@ class Qwen3NextConfig(PretrainedConfig):
num_attention_heads=16,
num_key_value_heads=2,
hidden_act="silu",
output_gate_type=None,
max_position_embeddings=32768,
initializer_range=0.02,
rms_norm_eps=1e-6,
@@ -223,6 +226,7 @@ class Qwen3NextConfig(PretrainedConfig):
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.output_gate_type = output_gate_type
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
+6
View File
@@ -143,6 +143,7 @@ class Qwen3_5GatedDeltaNet(nn.Module):
self.conv_kernel_size = config.linear_conv_kernel_dim
self.layer_id = layer_id
self.activation = config.hidden_act
self.output_gate_type = config.output_gate_type
self.layer_norm_epsilon = config.rms_norm_eps
# Conv1d layer
@@ -237,6 +238,11 @@ class Qwen3_5GatedDeltaNet(nn.Module):
norm_before_gate=True,
device=torch.get_device_module().current_device(),
dtype=config.torch_dtype,
**(
{"activation": self.output_gate_type}
if self.output_gate_type is not None
else {}
),
)
self.out_proj = RowParallelLinear(
+11 -1
View File
@@ -106,6 +106,7 @@ class Qwen3GatedDeltaNet(nn.Module):
self.conv_kernel_size = config.linear_conv_kernel_dim
self.layer_id = layer_id
self.activation = config.hidden_act
self.output_gate_type = config.output_gate_type
self.layer_norm_epsilon = config.rms_norm_eps
self.conv_dim = self.key_dim * 2 + self.value_dim
@@ -186,12 +187,21 @@ class Qwen3GatedDeltaNet(nn.Module):
norm_before_gate=True,
device=torch.get_device_module().current_device(),
dtype=config.torch_dtype,
**(
{"activation": self.output_gate_type}
if self.output_gate_type is not None
else {}
),
)
if not get_global_server_args().disable_piecewise_cuda_graph
else FusedRMSNormGated(
self.head_v_dim,
eps=self.layer_norm_epsilon,
activation=self.activation,
activation=(
self.output_gate_type
if self.output_gate_type is not None
else self.activation
),
device=torch.get_device_module().current_device(),
dtype=config.torch_dtype,
)