Add output_gate_type to Qwen3NextConfig and update models to utilize it (#25401)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user