From 73c42e5ded7ae216fee1a7738aad47aac501f217 Mon Sep 17 00:00:00 2001 From: datta0 Date: Thu, 27 Mar 2025 14:35:57 +0000 Subject: [PATCH 1/4] Initial support for Qwen3. Will udpate when the model is released --- unsloth/models/qwen3.py | 241 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 241 insertions(+) create mode 100644 unsloth/models/qwen3.py diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py new file mode 100644 index 0000000000..6e2ebe9e8b --- /dev/null +++ b/unsloth/models/qwen3.py @@ -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 From 71424b35ad3910f6d1c4a543d032a580a7674e77 Mon Sep 17 00:00:00 2001 From: datta0 Date: Fri, 28 Mar 2025 05:50:43 +0000 Subject: [PATCH 2/4] Add Qwen3Moe and necessitate transformers version --- unsloth/models/qwen3.py | 25 +++- unsloth/models/qwen3_moe.py | 223 ++++++++++++++++++++++++++++++++++++ 2 files changed, 242 insertions(+), 6 deletions(-) create mode 100644 unsloth/models/qwen3_moe.py diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index 6e2ebe9e8b..aa99aa88db 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -19,12 +19,25 @@ from .llama import ( LlamaRotaryEmbedding, LlamaLinearScalingRotaryEmbedding, ) -from transformers.models.qwen3.modeling_qwen3 import ( - Qwen3Attention, - Qwen3DecoderLayer, - Qwen3Model, - Qwen3ForCausalLM, -) +try: + from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3Attention, + Qwen3DecoderLayer, + Qwen3Model, + Qwen3ForCausalLM, + ) +except: + from packaging.version import Version + transformers_version = Version(transformers_version) + if not transformers_version >= Version("4.50.3"): #TODO: Update when transformers is updated + raise ImportError( + f"Unsloth: Your transformers version of {transformers_version} does not support Qwen3 and Qwen3Moe.\n"\ + f"The minimum required version is 4.50.3.\n"\ + f'Try `pip install --upgrade "transformers>=4.50.3"`\n'\ + f"to obtain the latest transformers build, then restart this session."\ + ) + pass + # For Pytorch 2.1.1 try: from transformers.models.qwen3.modeling_qwen3 import ( diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py new file mode 100644 index 0000000000..055c2dc907 --- /dev/null +++ b/unsloth/models/qwen3_moe.py @@ -0,0 +1,223 @@ +# 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 .qwen3 import ( + Qwen3Attention_fast_forward, + FastQwen3Model, +) +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeAttention, + Qwen3MoeSparseMoeBlock, + Qwen3MoeMLP, + Qwen3MoeDecoderLayer, + Qwen3MoeModel, + Qwen3MoeForCausalLM, +) +# For Pytorch 2.1.1 +# TODO: Transformers moved to `attention_interface`. So we might not need these anymore +# try: +# from transformers.models.qwen3_moe.modeling_qwen3_moe import ( +# Qwen3SdpaAttention, +# Qwen3FlashAttention2, +# ) +# except: +# Qwen3SdpaAttention = Qwen3Attention +# Qwen3FlashAttention2 = Qwen3Attention +# pass +from unsloth_zoo.utils import Version, _get_dtype + + +torch_nn_functional_softmax = torch.nn.functional.softmax +def Qwen3MoeSparseMoeBlock_fast_forward(self, X, temp_gate = None, temp_up = None): + # adapted from https://github.com/huggingface/transformers/pull/36878/files#diff-0855b77fc27ad9449158a1c74953f909b011c00de7125f7c8e68d0ff209c092aR356-R370 + + bsz, seq_len, hd = X.shape + X = X.view(-1, hd) + + router_logits = fast_linear_forward(self.gate_proj, X, out = temp_gate) #pretty much the only change from transformers implementation. + + routing_weights = torch_nn_functional_softmax(router_logits, dim = -1) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(X.dtype) + final_X = torch.zeros( + (bsz * seq_len, hd), dtype=X.dtype, device=X.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = X[None, top_x].reshape(-1, hd) + current_X = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_X.index_add_(0, top_x, current_X.to(X.dtype)) + final_X = final_X.reshape(bsz, seq_len, hd) + return final_X, router_logits +pass + + +def Qwen3MoeDecoderLayer_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: Optional[bool] = False, + output_router_logits: Optional[bool] = False, + use_cache: Optional[bool] = False, + padding_mask: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + *args, **kwargs, +): + residual = hidden_states + + if use_cache and hasattr(self, "_flag_for_generation"): #past_key_value is not None: + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states) + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + padding_mask=padding_mask, + position_embeddings = position_embeddings, + _flag_for_generation=self._flag_for_generation, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) + hidden_states = Qwen3MoeSparseMoeBlock_fast_forward(self.mlp, hidden_states) + hidden_states = residual + hidden_states + else: + residual = hidden_states + hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states) + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + padding_mask=padding_mask, + position_embeddings = position_embeddings, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + pass + + outputs = (hidden_states,) + if output_attentions: outputs += (self_attn_weights,) + if use_cache: outputs += (present_key_value,) + return outputs + + + +class FastQwen3MoeModel(FastQwen3Model): + + @staticmethod + def pre_patch(): + init_name, function = patch_linear_scaling( + model_name = "Qwen3Moe", + rope_module = LlamaRotaryEmbedding, + scaled_rope_module = LlamaLinearScalingRotaryEmbedding, + attention_module = Qwen3MoeAttention, + ) + if init_name is not None: + exec(function, globals()) + Qwen3MoeAttention.__init__ = eval(init_name) + pass + Qwen3MoeAttention .forward = Qwen3Attention_fast_forward + # Qwen3SdpaAttention .forward = Qwen3Attention_fast_forward + # Qwen3FlashAttention2 .forward = Qwen3Attention_fast_forward + Qwen3MoeSparseMoeBlock .forward = Qwen3MoeSparseMoeBlock_fast_forward + Qwen3MoeMLP .forward = fast_swiglu_inference # This is analogous to Dense models' MLP + Qwen3MoeDecoderLayer .forward = LlamaDecoderLayer_fast_forward + Qwen3MoeModel .forward = LlamaModel_fast_forward + Qwen3MoeForCausalLM .forward = CausalLM_fast_forward(LlamaModel_fast_forward_inference) + PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward + fix_prepare_inputs_for_generation(Qwen3MoeForCausalLM) + + # 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_moe.modeling_qwen3_moe + transformers.models.Qwen3Moe.modeling_qwen3_moe.Qwen3MoeRotaryEmbedding = 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 \ No newline at end of file From 406eb8cc715c0ed627886adddca74148072534f4 Mon Sep 17 00:00:00 2001 From: datta0 Date: Fri, 28 Mar 2025 05:58:01 +0000 Subject: [PATCH 3/4] Enable qwen3 and qwen3moe --- unsloth/models/__init__.py | 18 ++++++++++-------- unsloth/models/_utils.py | 2 +- unsloth/models/loader.py | 13 +++++++++++++ 3 files changed, 24 insertions(+), 9 deletions(-) diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index 317525c793..99db55c086 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .llama import FastLlamaModel -from .loader import FastLanguageModel, FastVisionModel, FastTextModel, FastModel -from .mistral import FastMistralModel -from .qwen2 import FastQwen2Model -from .granite import FastGraniteModel -from .dpo import PatchDPOTrainer, PatchKTOTrainer -from ._utils import is_bfloat16_supported, __version__ -from .rl import PatchFastRL, vLLMSamplingParams +from .llama import FastLlamaModel +from .loader import FastLanguageModel, FastVisionModel, FastTextModel, FastModel +from .mistral import FastMistralModel +from .qwen2 import FastQwen2Model +from .qwen3 import FastQwen3Model +from .qwen3_moe import FastQwen3MoeModel +from .granite import FastGraniteModel +from .dpo import PatchDPOTrainer, PatchKTOTrainer +from ._utils import is_bfloat16_supported, __version__ +from .rl import PatchFastRL, vLLMSamplingParams diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 840c15c003..e60881fa9c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -243,7 +243,7 @@ pass from transformers import __version__ as transformers_version from transformers import PretrainedConfig -model_architectures = ["llama", "mistral", "gemma", "gemma2", "qwen2", "granite"] +model_architectures = ["llama", "mistral", "gemma", "gemma2", "qwen2", "granite", "qwen3", "qwen3_moe"] for model_name in model_architectures: config_filepath = f"transformers.models.{model_name}.configuration_{model_name}" diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index cac5acd838..caa49d0bfb 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -23,6 +23,8 @@ from .granite import FastGraniteModel from .llama import FastLlamaModel, logger from .mistral import FastMistralModel from .qwen2 import FastQwen2Model +from .qwen3 import FastQwen3Model +from .qwen3_moe import FastQwen3MoeModel from .cohere import FastCohereModel from transformers import AutoConfig from transformers import __version__ as transformers_version @@ -51,6 +53,8 @@ SUPPORTS_GEMMA2 = transformers_version >= Version("4.42") SUPPORTS_LLAMA31 = transformers_version >= Version("4.43.2") SUPPORTS_LLAMA32 = transformers_version > Version("4.45.0") SUPPORTS_GRANITE = transformers_version >= Version("4.46.0") +SUPPORTS_QWEN3 = transformers_version >= Version("4.50.3") +SUPPORTS_QWEN3_MOE = transformers_version >= Version("4.50.3") if SUPPORTS_GEMMA: from .gemma import FastGemmaModel if SUPPORTS_GEMMA2: @@ -298,6 +302,15 @@ class FastLanguageModel(FastLlamaModel): dispatch_model = FastGemma2Model elif model_type == "qwen2": dispatch_model = FastQwen2Model + elif model_type == "qwen3" or model_type == "qwen3_moe": + if not SUPPORTS_QWEN3 or not SUPPORTS_QWEN3_MOE: + raise ImportError( + f"Unsloth: Your transformers version of {transformers_version} does not support Qwen3.\n"\ + f"The minimum required version is 4.50.3.\n"\ + f'Try `pip install --upgrade "transformers>=4.50.3"`\n'\ + f"to obtain the latest transformers build, then restart this session."\ + ) + dispatch_model = FastQwen3Model if model_type == "qwen3" else FastQwen3MoeModel # Temporary disable optimized Cohere until errors match # elif model_type == "cohere": # dispatch_model = FastCohereModel From 00589310df8d80e0277a5ae79d2149c7c08ef6b9 Mon Sep 17 00:00:00 2001 From: datta0 Date: Wed, 2 Apr 2025 06:49:06 +0000 Subject: [PATCH 4/4] add comments and use modified function --- unsloth/models/qwen3_moe.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py index 055c2dc907..319e3bffe3 100644 --- a/unsloth/models/qwen3_moe.py +++ b/unsloth/models/qwen3_moe.py @@ -76,7 +76,7 @@ def Qwen3MoeSparseMoeBlock_fast_forward(self, X, temp_gate = None, temp_up = Non # the current expert. We need to make sure to multiply the output hidden # states by `routing_weights` on the corresponding tokens (top-1 and top-2) current_state = X[None, top_x].reshape(-1, hd) - current_X = expert_layer(current_state) * routing_weights[top_x, idx, None] + current_X = expert_layer(current_state) * routing_weights[top_x, idx, None] # Qwen3MoeMLP.forward = fast_swiglu_inference takes care of making this faster. Analogous to Dense models' MLP # However `index_add_` only support torch tensors for indexing so we'll use # the `top_x` tensor here. @@ -119,10 +119,10 @@ def Qwen3MoeDecoderLayer_fast_forward( ) hidden_states = residual + hidden_states - # Fully Connected + # MoE Router MLP residual = hidden_states hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) - hidden_states = Qwen3MoeSparseMoeBlock_fast_forward(self.mlp, hidden_states) + hidden_states, router_logits = Qwen3MoeSparseMoeBlock_fast_forward(self.mlp, hidden_states) hidden_states = residual + hidden_states else: residual = hidden_states @@ -140,15 +140,16 @@ def Qwen3MoeDecoderLayer_fast_forward( ) hidden_states = residual + hidden_states - # Fully Connected + # MoE Router MLP residual = hidden_states hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states, router_logits = self.mlp(hidden_states) hidden_states = residual + hidden_states pass outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) + if output_router_logits: outputs += (router_logits,) if use_cache: outputs += (present_key_value,) return outputs @@ -173,7 +174,7 @@ class FastQwen3MoeModel(FastQwen3Model): # Qwen3FlashAttention2 .forward = Qwen3Attention_fast_forward Qwen3MoeSparseMoeBlock .forward = Qwen3MoeSparseMoeBlock_fast_forward Qwen3MoeMLP .forward = fast_swiglu_inference # This is analogous to Dense models' MLP - Qwen3MoeDecoderLayer .forward = LlamaDecoderLayer_fast_forward + Qwen3MoeDecoderLayer .forward = Qwen3MoeDecoderLayer_fast_forward Qwen3MoeModel .forward = LlamaModel_fast_forward Qwen3MoeForCausalLM .forward = CausalLM_fast_forward(LlamaModel_fast_forward_inference) PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward