File size: 5,621 Bytes
a1111f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | # coding=utf-8
# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Callable, Optional, Union
import torch
from torch import nn
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
from transformers.utils.deprecation import deprecate_kwarg
from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
from .summary_context import SummaryBatchContext, build_summary_sliding_context
from .runtime_config import _get_shared_batch_mask
from .layer_norm_rope import Qwen3RMSNorm
from .layer_attention import Qwen3Attention, Qwen3SummaryAttention
class Qwen3MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x):
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
return down_proj
class Qwen3DecoderLayer(GradientCheckpointingLayer):
def __init__(self, config: Qwen3Config, layer_idx: int):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
# Use SummaryAttention if enabled in config
if getattr(config, "use_summary_attention", False) is True and config.summary_layer_freq[layer_idx] == 1:
self.self_attn = Qwen3SummaryAttention(config=config, layer_idx=layer_idx)
elif getattr(config, "use_summary_attention", False) is False and config.summary_layer_freq[layer_idx] == 0:
self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
else:
raise ValueError(f'Check config.summary_layer_freq {config.summary_layer_freq} and config.use_summary_attention {config.use_summary_attention}')
self.mlp = Qwen3MLP(config)
self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
if getattr(config, "summary_independent_attention_layernorm", False):
self.input_layernorm_summary = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.attention_type = config.layer_types[layer_idx]
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Cache] = None,
use_cache: Optional[bool] = False,
cache_position: Optional[torch.LongTensor] = None,
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
summary_ctx: Optional[SummaryBatchContext] = None,
**kwargs: Unpack[TransformersKwargs],
) -> torch.Tensor:
residual = hidden_states
if getattr(self.config, "summary_independent_attention_layernorm", False):
summary_mask = summary_ctx.summary_mask
summary_pos = _get_shared_batch_mask(summary_mask, "summary_mask").to(residual.device)
hidden_states = self.input_layernorm(hidden_states)
if summary_pos.any():
hidden_summary = residual[:, summary_pos, :]
hidden_summary = self.input_layernorm_summary(hidden_summary)
hidden_states[:, summary_pos, :] = hidden_summary
else:
hidden_states = self.input_layernorm(hidden_states)
# Self Attention - pass summary_ctx if using summary attention
attn_kwargs = {
"hidden_states": hidden_states,
"attention_mask": attention_mask,
"position_ids": position_ids,
"past_key_values": past_key_values,
"use_cache": use_cache,
"cache_position": cache_position,
"position_embeddings": position_embeddings,
**kwargs,
}
if isinstance(self.self_attn, Qwen3SummaryAttention):
attn_kwargs["summary_ctx"] = summary_ctx
hidden_states, _ = self.self_attn(**attn_kwargs)
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
return hidden_states
|