Initial support for Qwen3. Will udpate when the model is released
This commit is contained in:
datta0
parent
f6b8463b8d
commit
7e928545b1
1 changed files with 241 additions and 0 deletions
241
unsloth/models/qwen3.py
Normal file
241
unsloth/models/qwen3.py
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth 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 .llama import *
|
||||
import os
|
||||
from ._utils import __version__
|
||||
from .llama import (
|
||||
LlamaRotaryEmbedding,
|
||||
LlamaLinearScalingRotaryEmbedding,
|
||||
)
|
||||
from transformers.models.qwen3.modeling_qwen3 import (
|
||||
Qwen3Attention,
|
||||
Qwen3DecoderLayer,
|
||||
Qwen3Model,
|
||||
Qwen3ForCausalLM,
|
||||
)
|
||||
# For Pytorch 2.1.1
|
||||
try:
|
||||
from transformers.models.qwen3.modeling_qwen3 import (
|
||||
Qwen3SdpaAttention,
|
||||
Qwen3FlashAttention2,
|
||||
)
|
||||
except:
|
||||
Qwen3SdpaAttention = Qwen3Attention
|
||||
Qwen3FlashAttention2 = Qwen3Attention
|
||||
pass
|
||||
from unsloth_zoo.utils import Version, _get_dtype
|
||||
|
||||
|
||||
def Qwen3Attention_fast_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
causal_mask: Optional[BlockDiagonalCausalMask] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
padding_mask: Optional[torch.LongTensor] = None,
|
||||
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
*args, **kwargs,
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||||
|
||||
# Clear inference
|
||||
if hasattr(self, "paged_attention"):
|
||||
del self.paged_attention_K
|
||||
del self.paged_attention_V
|
||||
del self.paged_attention
|
||||
del self.temp_QA
|
||||
del self.temp_KV
|
||||
del self.RH_Q
|
||||
del self.attention
|
||||
pass
|
||||
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
|
||||
n_heads = self.config.num_attention_heads
|
||||
n_groups = self.num_key_value_groups
|
||||
n_kv_heads = self.config.num_key_value_heads
|
||||
head_dim = self.head_dim
|
||||
assert(n_kv_heads * n_groups == n_heads)
|
||||
|
||||
Q, K, V = self.apply_qkv(self, hidden_states)
|
||||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
|
||||
#Qwen3 has QKNorm. This seems to be the only difference from Qwen2.
|
||||
Q = fast_layernorm_compiled(self.q_norm, Q)
|
||||
K = fast_layernorm_compiled(self.k_norm, K)
|
||||
|
||||
kv_seq_len = K.shape[-2]
|
||||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
if position_ids is None:
|
||||
cos = self.rotary_emb.cos_cached
|
||||
sin = self.rotary_emb.sin_cached
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
else:
|
||||
cos, sin = self.rotary_emb(V, seq_len = kv_seq_len)
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
pass
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
pass
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if (not HAS_FLASH_ATTENTION and attention_mask is None):
|
||||
# Xformers memory efficient attention
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
K_M = V_M = bsz * kv_seq_len
|
||||
Q_M = bsz * q_len
|
||||
|
||||
has_swa = isinstance(causal_mask, xformers.attn_bias.BlockDiagonalCausalMask)
|
||||
|
||||
# Group query attention
|
||||
K = K .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, Q_M, n_heads, head_dim)
|
||||
K = K.view(1, K_M, n_heads, head_dim)
|
||||
V = V.view(1, V_M, n_heads, head_dim)
|
||||
pass
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, Q_M, n_kv_heads, n_groups, head_dim)
|
||||
K = K.view(1, K_M, n_kv_heads, n_groups, head_dim)
|
||||
V = V.view(1, V_M, n_kv_heads, n_groups, head_dim)
|
||||
pass
|
||||
pass
|
||||
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
sw = getattr(self.config, "sliding_window", None)
|
||||
sw = kv_seq_len if (sw is None or sw == "null") else sw
|
||||
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
|
||||
else:
|
||||
# Grouped query attention
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
# pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
pass
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
return attn_output, attn_weights, past_key_value
|
||||
pass
|
||||
|
||||
|
||||
class FastQwen3Model(FastLlamaModel):
|
||||
|
||||
@staticmethod
|
||||
def pre_patch():
|
||||
init_name, function = patch_linear_scaling(
|
||||
model_name = "Qwen3",
|
||||
rope_module = LlamaRotaryEmbedding,
|
||||
scaled_rope_module = LlamaLinearScalingRotaryEmbedding,
|
||||
attention_module = Qwen3Attention,
|
||||
)
|
||||
if init_name is not None:
|
||||
exec(function, globals())
|
||||
Qwen3Attention.__init__ = eval(init_name)
|
||||
pass
|
||||
Qwen3Attention .forward = Qwen3Attention_fast_forward
|
||||
Qwen3SdpaAttention .forward = Qwen3Attention_fast_forward
|
||||
Qwen3FlashAttention2.forward = Qwen3Attention_fast_forward
|
||||
Qwen3DecoderLayer .forward = LlamaDecoderLayer_fast_forward
|
||||
Qwen3Model .forward = LlamaModel_fast_forward
|
||||
Qwen3ForCausalLM .forward = CausalLM_fast_forward(LlamaModel_fast_forward_inference)
|
||||
PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward
|
||||
fix_prepare_inputs_for_generation(Qwen3ForCausalLM)
|
||||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.qwen3.modeling_qwen3
|
||||
transformers.models.Qwen3.modeling_qwen3.Qwen3RotaryEmbedding = LlamaRotaryEmbedding
|
||||
return
|
||||
pass
|
||||
|
||||
|
||||
@staticmethod
|
||||
def from_pretrained( #TODO: Change after release
|
||||
model_name = "Qwen/Qwen3-7B",
|
||||
max_seq_length = 4096,
|
||||
dtype = None,
|
||||
load_in_4bit = True,
|
||||
token = None,
|
||||
device_map = "sequential",
|
||||
rope_scaling = None,
|
||||
fix_tokenizer = True,
|
||||
model_patcher = None,
|
||||
tokenizer_name = None,
|
||||
trust_remote_code = False,
|
||||
**kwargs,
|
||||
):
|
||||
return FastLlamaModel.from_pretrained(
|
||||
model_name = model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = load_in_4bit,
|
||||
token = token,
|
||||
device_map = device_map,
|
||||
rope_scaling = rope_scaling,
|
||||
fix_tokenizer = fix_tokenizer,
|
||||
model_patcher = FastQwen3Model,
|
||||
tokenizer_name = tokenizer_name,
|
||||
trust_remote_code = trust_remote_code,
|
||||
**kwargs,
|
||||
)
|
||||
pass
|
||||
pass
|
||||
Loading…
Add table
Add a link
Reference in a new issue