File size: 2,704 Bytes
30902bd |
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 |
import importlib.metadata
from packaging import version
from torch import nn
from transformers import Qwen3Config, Qwen3Model, Qwen3PreTrainedModel
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.models.qwen3.modeling_qwen3 import (
Qwen3Attention,
Qwen3DecoderLayer,
Qwen3MLP,
Qwen3RMSNorm,
Qwen3RotaryEmbedding,
)
from transformers.utils import logging
from transformers.utils.import_utils import _is_package_available
logger = logging.get_logger(__name__)
def is_transformers_attn_greater_or_equal_4_56_2():
if not _is_package_available("transformers"):
return False
return version.parse(importlib.metadata.version("transformers")) >= version.parse(
"4.56.2"
)
class ModifiedQwen3Attention(Qwen3Attention):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.is_causal = False
class ModifiedQwen3DecoderLayer(Qwen3DecoderLayer):
def __init__(self, config: Qwen3Config, layer_idx: int):
GradientCheckpointingLayer.__init__(self)
self.hidden_size = config.hidden_size
self.self_attn = ModifiedQwen3Attention(config=config, layer_idx=layer_idx)
self.mlp = Qwen3MLP(config)
self.input_layernorm = 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]
class Qwen3EncoderModel(Qwen3Model):
_no_split_modules = ["ModifiedQwen3DecoderLayer"]
def __init__(self, config: Qwen3Config):
if not is_transformers_attn_greater_or_equal_4_56_2():
raise ValueError(
"The current implementation of Qwen2EncoderModel follows modeling_qwen2.py of transformers version >= 4.56.2"
)
Qwen3PreTrainedModel.__init__(self, config)
self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size
self.embed_tokens = nn.Embedding(
config.vocab_size, config.hidden_size, self.padding_idx
)
self.layers = nn.ModuleList(
[
ModifiedQwen3DecoderLayer(config, layer_idx)
for layer_idx in range(config.num_hidden_layers)
]
)
self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.rotary_emb = Qwen3RotaryEmbedding(config=config)
self.gradient_checkpointing = False
self.has_sliding_layers = "sliding_attention" in self.config.layer_types
# Initialize weights and apply final processing
self.post_init()
|