diff --git a/README.md b/README.md index 2353708297..f3af0ee7dc 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ -### Finetune Llama 4, Gemma 3, Phi-4, Qwen 2.5 & Mistral 2x faster with 80% less VRAM! +### Finetune Qwen3, Llama 4, Gemma 3, Phi-4 & Mistral 2x faster with 80% less VRAM! ![](https://i.ibb.co/sJ7RhGG/image-41.png) diff --git a/pyproject.toml b/pyproject.toml index 44bd74ea35..5bfe4fcf75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ triton = [ ] huggingface = [ - "unsloth_zoo>=2025.4.1", + "unsloth_zoo>=2025.4.2", "packaging", "tyro", "transformers>=4.46.1,!=4.47.0", @@ -355,7 +355,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo>=2025.4.1", + "unsloth_zoo>=2025.4.2", "packaging", "tyro", "transformers>=4.46.1,!=4.47.0", diff --git a/unsloth/dataprep/__init__.py b/unsloth/dataprep/__init__.py new file mode 100644 index 0000000000..b36122eb74 --- /dev/null +++ b/unsloth/dataprep/__init__.py @@ -0,0 +1,15 @@ +# 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 .synthetic import * diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py new file mode 100644 index 0000000000..8fcbc1bef1 --- /dev/null +++ b/unsloth/dataprep/synthetic.py @@ -0,0 +1,261 @@ +# 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. + +__all__ = [ + "check_vllm_status", + "async_load_vllm", + "destroy_vllm", + "configure_synthetic_data_kit", +] +import subprocess +import time +import os +import requests +import torch +import gc +import time +from unsloth_zoo.vllm_utils import load_vllm +from transformers import AutoConfig + +def check_vllm_status(): + try: + response = requests.get("http://localhost:8000/metrics") + if response.status_code == 200: + return True + except requests.exceptions.ConnectionError: + return False + pass +pass + + +def async_load_vllm( + model_name = "unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit", + max_seq_length = 2048, + gpu_memory_utilization = 0.85, + float8_kv_cache = False, + conservativeness = 1.0, + token = None, +): + config = AutoConfig.from_pretrained( + model_name, + token = token, + ) + engine_args = load_vllm( + model_name = model_name, + config = config, + gpu_memory_utilization = gpu_memory_utilization, + max_seq_length = max_seq_length, + disable_log_stats = True, + float8_kv_cache = float8_kv_cache, + conservativeness = conservativeness, + return_args = True, + enable_lora = False, + ) + if "device" in engine_args: del engine_args["device"] + if "model" in engine_args: del engine_args["model"] + + subprocess_commands = [ + "vllm", "serve", str(model_name), + ] + for key, value in engine_args.items(): + flag = "--" + key.replace("_", "-") + which = str(value).lower().replace("torch.", "") + subprocess_commands += [flag, which,] + pass + print(subprocess_commands) + vllm_process = subprocess.Popen( + subprocess_commands, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + start_new_session = True, + ) + ready_message_part = b"Starting vLLM API server on" + ready = False + while vllm_process.poll() is None: + output = vllm_process.stdout.readline() + if not output: + print("Stdout stream ended before readiness message detected.") + break + output_str = output.decode('utf-8', errors='ignore').strip() + print(f"vLLM STDOUT: {output_str}") + if ready_message_part in output: + print(f"\n--- vLLM Server Ready (Detected: '{ready_message_part.decode()}') ---") + ready = True + break + pass + pass + if vllm_process is None: + raise RuntimeError("Unsloth: vllm_process failed to load!") + trial = 0 + while not check_vllm_status(): + if trial >= 100: + raise RuntimeError("Unsloth: vllm_process failed to load!") + trial += 1 + time.sleep(1) + return vllm_process +pass + + +def destroy_vllm(vllm_process): + print("Attempting to terminate the VLLM server gracefully...") + try: + vllm_process.terminate() + vllm_process.wait(timeout=10) + print("Server terminated gracefully.") + except subprocess.TimeoutExpired: + print("Server did not terminate gracefully after 10 seconds. Forcing kill...") + vllm_process.kill() + vllm_process.wait() + print("Server killed forcefully.") + except Exception as e: + print(f"An error occurred while trying to stop the process: {e}") + try: + if vllm_process.poll() is None: + print("Attempting forceful kill due to error...") + vllm_process.kill() + vllm_process.wait() + print("Server killed forcefully after error.") + except Exception as kill_e: + print(f"Error during forceful kill: {kill_e}") + for _ in range(10): + torch.cuda.empty_cache() + gc.collect() +pass + + +synthetic_config_string = """\ +# Master configuration file for Synthetic Data Kit + +# Global paths configuration +paths: + # Input data locations + input: + pdf: "data/pdf" + html: "data/html" + youtube: "data/youtube" + docx: "data/docx" + ppt: "data/ppt" + txt: "data/txt" + + # Output locations + output: + parsed: "data/output" # Where parsed text files are saved + generated: "data/generated" # Where generated content is saved + cleaned: "data/cleaned" # Where cleaned content is saved + final: "data/final" # Where final formatted content is saved + +# VLLM server configuration +vllm: + api_base: "http://localhost:8000/v1" # Base URL for VLLM API + port: 8000 # Port for VLLM server + model: "{model_name}" # Default model to use + max_retries: 3 # Number of retries for API calls + retry_delay: 1.0 # Initial delay between retries (seconds) + +# Ingest configuration +ingest: + default_format: "txt" # Default output format for parsed files + youtube_captions: "auto" # Options: "auto", "manual" - caption preference + +# LLM generation parameters +generation: + temperature: {temperature} # Higher = more creative, lower = more deterministic + top_p: {top_p} # Nucleus sampling parameter + chunk_size: {chunk_size} # Size of text chunks for processing + overlap: {overlap} # Overlap between chunks to maintain context + max_tokens: {max_tokens} # Maximum tokens in LLM responses + num_pairs: {default_num_pairs} # Default number of QA pairs to generate + +# Content cleanup parameters +cleanup: + threshold: {cleanup_threshold} # Default quality threshold (1-10) + batch_size: {cleanup_batch_size} # Number of items per batch for rating + temperature: {cleanup_temperature} # Temperature for rating (lower = more consistent) + +# Format conversion parameters +format: + default: "jsonl" # Default output format + include_metadata: true # Include metadata in output files + pretty_json: true # Use indentation in JSON output + +# Prompts for different tasks +prompts: + # Summary generation prompt + summary: | + Summarize this document in 3-5 sentences, focusing on the main topic and key concepts. + + # QA pair generation prompt + qa_generation: | + Create {num_pairs} question-answer pairs from this text for LLM training. + + Rules: + 1. Questions must be about important facts in the text + 2. Answers must be directly supported by the text + 3. Return JSON format only: + + [ + {{ + "question": "Question 1?", + "answer": "Answer 1." + }}, + {{ + "question": "Question 2?", + "answer": "Answer 2." + }} + ] + + Text: + {text} + + # QA pair rating prompt + qa_rating: | + Rate each of these question-answer pairs for quality and return exactly this JSON format: + + [ + {{"question": "same question text", "answer": "same answer text", "rating": n}} + ] + + Where n is a number from 1-10. + + DO NOT include any text outside of the JSON array, just return valid JSON: + + {pairs}""" + + +def configure_synthetic_data_kit( + model_name = "unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit", + temperature = 0.7, + top_p = 0.95, + chunk_size = 4000, + overlap = 200, + max_tokens = 512, + default_num_pairs = 25, + cleanup_threshold = 1.0, + cleanup_batch_size = 4, + cleanup_temperature = 0.3, +): + config = synthetic_config_string\ + .replace("{model_name}", str(model_name))\ + .replace("{temperature}", str(temperature))\ + .replace("{top_p}", str(top_p))\ + .replace("{chunk_size}", str(chunk_size))\ + .replace("{overlap}", str(overlap))\ + .replace("{max_tokens}", str(max_tokens))\ + .replace("{default_num_pairs}", str(default_num_pairs))\ + .replace("{cleanup_threshold}", str(cleanup_threshold))\ + .replace("{cleanup_batch_size}", str(cleanup_batch_size))\ + .replace("{cleanup_temperature}", str(cleanup_temperature)) + + return config +pass 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 e18d024522..ed8a2ade64 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.4.1" +__version__ = "2025.4.3" __all__ = [ "SUPPORTS_BFLOAT16", @@ -243,12 +243,12 @@ 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}" model_filepath = f"transformers.models.{model_name}.modeling_{model_name}" - config_filename = f"{model_name.title()}Config" + config_filename = f"{model_name.title().replace('_','')}Config" # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now exec(f"from {config_filepath} import {config_filename}", globals()) try: diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 722b50d27a..2b07e5f1fe 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -911,98 +911,104 @@ pass # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825 -def LlamaModel_fast_forward_inference( - self, - input_ids, - past_key_values, - position_ids, - attention_mask = None, -): - input_ids = input_ids[:,:self.max_seq_length] - bsz, q_len = input_ids.shape - hd = self.config.hidden_size - mlp_size = self.config.intermediate_size +def _LlamaModel_fast_forward_inference(attention_fast_forward_inference=LlamaAttention_fast_forward_inference, mlp_fast_forward_inference=fast_swiglu_inference): + # This makes the attention and MLP customisable. + # Now for models like qwen3 or cohere which use custom attention operations, we can use this function + def LlamaModel_fast_forward_inference_custom( + self, + input_ids, + past_key_values, + position_ids, + attention_mask = None, + ): + input_ids = input_ids[:,:self.max_seq_length] + bsz, q_len = input_ids.shape + hd = self.config.hidden_size + mlp_size = self.config.intermediate_size - X = self.model.embed_tokens(input_ids) - X = X.to(_get_dtype(self.config.torch_dtype)) - bsz, q_len, hd = X.shape - assert(q_len == 1) - # Get saved buffers to reduce memory movement - residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") - _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") - XX, XX2 = _XX[0], _XX[1] - variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0") - temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0") - temp_gate, temp_up = temp_mlp[0], temp_mlp[1] + X = self.model.embed_tokens(input_ids) + X = X.to(_get_dtype(self.config.torch_dtype)) + bsz, q_len, hd = X.shape + assert(q_len == 1) + # Get saved buffers to reduce memory movement + residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") + _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") + XX, XX2 = _XX[0], _XX[1] + variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0") + temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0") + temp_gate, temp_up = temp_mlp[0], temp_mlp[1] - seq_len = past_key_values[0][0].shape[-2] - if bsz != 1: - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - (bsz, q_len), - X, - seq_len, - sliding_window = getattr(self.config, "sliding_window", None), - ) - else: - attention_mask = None - pass + seq_len = past_key_values[0][0].shape[-2] + if bsz != 1: + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (bsz, q_len), + X, + seq_len, + sliding_window = getattr(self.config, "sliding_window", None), + ) + else: + attention_mask = None + pass - next_decoder_cache = [] + next_decoder_cache = [] - for idx, decoder_layer in enumerate(self.model.layers): - residual.copy_(X) # residual = X + for idx, decoder_layer in enumerate(self.model.layers): + residual.copy_(X) # residual = X + X = fast_rms_layernorm_inference( + decoder_layer.input_layernorm, + X, + XX = XX, + XX2 = XX2, + variance = variance, + ) + X, present_key_value = attention_fast_forward_inference( + decoder_layer.self_attn, + hidden_states = X, + past_key_value = past_key_values[idx], + position_ids = position_ids, + attention_mask = attention_mask, + do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"), + ) + X += residual + + residual.copy_(X) # residual = X + X = fast_rms_layernorm_inference( + decoder_layer.post_attention_layernorm, + X, + XX = XX, + XX2 = XX2, + variance = variance, + ) + X = mlp_fast_forward_inference( + decoder_layer.mlp, + X, + temp_gate = temp_gate, + temp_up = temp_up, + ) + X += residual + + next_decoder_cache.append(present_key_value) + pass X = fast_rms_layernorm_inference( - decoder_layer.input_layernorm, + self.model.norm, X, XX = XX, XX2 = XX2, variance = variance, ) - X, present_key_value = LlamaAttention_fast_forward_inference( - decoder_layer.self_attn, - hidden_states = X, - past_key_value = past_key_values[idx], - position_ids = position_ids, - attention_mask = attention_mask, - do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"), - ) - X += residual - residual.copy_(X) # residual = X - X = fast_rms_layernorm_inference( - decoder_layer.post_attention_layernorm, - X, - XX = XX, - XX2 = XX2, - variance = variance, + return BaseModelOutputWithPast( + last_hidden_state = X, + past_key_values = next_decoder_cache, + hidden_states = [], + attentions = [], ) - X = fast_swiglu_inference( - decoder_layer.mlp, - X, - temp_gate = temp_gate, - temp_up = temp_up, - ) - X += residual - - next_decoder_cache.append(present_key_value) pass - X = fast_rms_layernorm_inference( - self.model.norm, - X, - XX = XX, - XX2 = XX2, - variance = variance, - ) - - return BaseModelOutputWithPast( - last_hidden_state = X, - past_key_values = next_decoder_cache, - hidden_states = [], - attentions = [], - ) -pass + return LlamaModel_fast_forward_inference_custom +# For ensuring backwards compatibility, we create LlamaModel_fast_forward_inference that is consumed by other models +LlamaModel_fast_forward_inference = _LlamaModel_fast_forward_inference() def CausalLM_fast_forward(fast_forward_inference): def _CausalLM_fast_forward( @@ -2487,6 +2493,8 @@ class FastLlamaModel: elif model_type == "gemma2": apply_lora_mlp = apply_lora_mlp_geglu_approx elif model_type == "cohere": apply_lora_mlp = apply_lora_mlp_swiglu elif model_type == "granite": apply_lora_mlp = apply_lora_mlp_swiglu + elif model_type == "qwen3": apply_lora_mlp = apply_lora_mlp_swiglu + elif model_type == "qwen3moe": apply_lora_mlp = apply_lora_mlp_swiglu else: raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!") pass diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 3cd8508ffa..3d75c35117 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 diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index b8128968c9..82bec922e0 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -738,15 +738,39 @@ __INT_TO_FLOAT_MAPPER = \ "canopylabs/orpheus-3b-0.1-ft", "unsloth/orpheus-3b-0.1-ft-bnb-4bit", ), - "unsloth/Llama-4-Scout-17B-16E-Instruct-unsloth-dynamic-bnb-4bit" : ( - "unsloth/Llama-4-Scout-17B-16E-Instruct-unsloth", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "unsloth/Llama-4-Scout-17B-16E-Instruct-unsloth-bnb-4bit", + "unsloth/Qwen3-0.6B-unsloth-bnb-4bit" : ( + "unsloth/Qwen3-0.6B", + "Qwen/Qwen3-0.6B", + "unsloth/Qwen3-0.6B-bnb-4bit", ), - "unsloth/Llama-4-Scout-17B-16E-unsloth-dynamic-bnb-4bit" : ( - "unsloth/Llama-4-Scout-17B-16E-unsloth", - "meta-llama/Llama-4-Scout-17B-16E", - "unsloth/Llama-4-Scout-17B-16E-unsloth-bnb-4bit", + "unsloth/Qwen3-1.7B-unsloth-bnb-4bit" : ( + "unsloth/Qwen3-1.7B", + "Qwen/Qwen3-1.7B", + "unsloth/Qwen3-1.7B-bnb-4bit", + ), + "unsloth/Qwen3-4B-unsloth-bnb-4bit" : ( + "unsloth/Qwen3-4B", + "Qwen/Qwen3-4B", + "unsloth/Qwen3-4B-bnb-4bit", + ), + "unsloth/Qwen3-8B-unsloth-bnb-4bit" : ( + "unsloth/Qwen3-8B", + "Qwen/Qwen3-8B", + "unsloth/Qwen3-8B-bnb-4bit", + ), + "unsloth/Qwen3-14B-unsloth-bnb-4bit" : ( + "unsloth/Qwen3-14B", + "Qwen/Qwen3-14B", + "unsloth/Qwen3-14B-bnb-4bit", + ), + "unsloth/Qwen3-32B-unsloth-bnb-4bit" : ( + "unsloth/Qwen3-32B", + "Qwen/Qwen3-32B", + "unsloth/Qwen3-32B-bnb-4bit", + ), + "unsloth/Qwen3-30B-A3B-bnb-4bit" : ( + "unsloth/Qwen3-30B-A3B", + "Qwen/Qwen3-30B-A3B", ), } diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py new file mode 100644 index 0000000000..c0ceefd10c --- /dev/null +++ b/unsloth/models/qwen3.py @@ -0,0 +1,527 @@ +# 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, + _LlamaModel_fast_forward_inference, +) +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 +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask_for_sdpa, +) +# 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) # we will transpose after normalisation + K = K.view(bsz, q_len, n_kv_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation + 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. + # Note that using fast_layernorm_compiled causes issues as the dimensions don't match up. + # I tried to add a compiled version of the new norm but the numbers don't match up with Transformers + # TODO: Check on the differences here. + Q = fast_rms_layernorm(self.q_norm, Q) + K = fast_rms_layernorm(self.k_norm, K) + + Q = Q.transpose(1, 2) + K = K.transpose(1, 2) + + kv_seq_len = K.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[-2] + + if position_embeddings: + cos, sin = position_embeddings + else: + # Extend RoPE dynamically to fit in VRA + rotary_emb = self.rotary_emb + rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len) + + if position_ids is None: + # Useful for LongRoPE + cos, sin = rotary_emb.get_cached(kv_seq_len) + else: + cos, sin = rotary_emb(V, seq_len = kv_seq_len) + Q, K = fast_rope_embedding(Q, K, cos, sin) + + 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 = kv_seq_len + 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 + +torch_matmul = torch.matmul +def Qwen3Attention_fast_forward_inference( + self, + hidden_states: torch.Tensor, + past_key_value: Optional[Tuple[torch.Tensor]], + position_ids, + do_prefill = False, + attention_mask = None, +): + """ + https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406 + Fast inference using KV cache. + QK^T can be computed in 4 chunks + + [Q, q] @ [K, k].T where q, k are the new tokens. + [QK^T, Qk^T] + [qK^T, qk^T] + + Since the attention mask wipes Qk^T, we just get + [QK^T, 0] + [qK^T, qk^T] + + Since softmax is row-wise, we get + softmax([QK^T, 0]) + softmax([qK^T, qk^T]) + + We then multiply by [V] + [v] + softmax([QK^T, 0]) [softmax(QK^T)V] * + softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]] + + But notice * [softmax(QK^T)V] is just the last attention. + We just need to compute the last final row. + + This means we can pass in a row of Q, but we need to + remember K and V, which are called the KV cache. + """ + Xn = hidden_states + bsz, _, hd = hidden_states.size() + K1, V1 = past_key_value + dtype = Xn.dtype + + 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) + + hidden_size = self.config.hidden_size + attention_size = n_heads*head_dim + seq_len = K1.shape[-2] + kv_seq_len = seq_len + 1 + + # Prefill phase + # if not hasattr(self, "paged_attention"): + device = hidden_states.device + if do_prefill: + self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device) + self.paged_attention_K = self.paged_attention[:,0] + self.paged_attention_V = self.paged_attention[:,1] + self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) + self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) + self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device) + self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device) + self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device) + + # Mistral Nemo 12b has weird dimensions + if attention_size != hidden_size: + self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device) + else: + self.temp_O = self.temp_QA[1][:,:,:hidden_size] + pass + + self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device) + self.scalar = 1.0 / math_sqrt(self.head_dim) + self.half_head_dim = head_dim // 2 + elif kv_seq_len >= self.paged_attention.shape[0]: + self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim)) + self.paged_attention_K = self.paged_attention[:,0] + self.paged_attention_V = self.paged_attention[:,1] + self.attention.resize_((bsz, n_heads, 1, self.attention.shape[-1]+KV_CACHE_INCREMENT)) + pass + + Qn = fast_linear_forward(self.q_proj, Xn, out = self.temp_QA[0]) + Kn = fast_linear_forward(self.k_proj, Xn, out = self.temp_KV[0]) + Vn = fast_linear_forward(self.v_proj, Xn, out = self.temp_KV[1]) + Qn = Qn.view(bsz, 1, n_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation + Kn = Kn.view(bsz, 1, n_kv_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation + Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) + + Qn = fast_rms_layernorm(self.q_norm, Qn) + Kn = fast_rms_layernorm(self.k_norm, Kn) + + Qn = Qn.transpose(1, 2) + Kn = Kn.transpose(1, 2) + + # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) + # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) + + # Need to do it prior 2 steps before hitting full on short KV cache + # or else error + self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2) + cos, sin = self.rotary_emb.get_cached(kv_seq_len) + cos = cos[position_ids].unsqueeze(1) + sin = sin[position_ids].unsqueeze(1) + h = self.half_head_dim + + RH_Q = self.RH_Q + RH_Q[:,:,:,:h] = Qn[:,:,:,h:] + RH_Q[:,:,:,h:] = Qn[:,:,:,:h] + RH_Q[:,:,:,:h].neg_() # torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]) + Qn *= cos + Qn.addcmul_(RH_Q, sin) + + RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0") + RH_K[:,:,:,:h] = Kn[:,:,:,h:] + RH_K[:,:,:,h:] = Kn[:,:,:,:h] + RH_K[:,:,:,:h].neg_() #torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]) + Kn *= cos + Kn.addcmul_(RH_K, sin) + + # New KV cache + # Kn = torch.cat([K1, Kn], dim = 2) + # Vn = torch.cat([V1, Vn], dim = 2) + self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3) + self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3) + Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3) + Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3) + + # Handle sliding windows + sliding_window = getattr(self.config, "sliding_window", None) + if sliding_window is not None and kv_seq_len > sliding_window: + # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193 + slicing_tokens = 1 - sliding_window + Knn = Kn[:, :, slicing_tokens:, :]#.contiguous() + Vnn = Vn[:, :, slicing_tokens:, :]#.contiguous() + else: + Knn, Vnn = Kn, Vn + pass + + # Grouped query attention + _, _, cached_len, _ = Knn.shape + if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1: + Knn = Knn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim) + Vnn = Vnn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim) + Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim) + Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim) + pass + # else: + # Knn, Vnn = Knn, Vnn + # pass + + # Attention + if bsz == 1: + Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963 + # It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows + A = torch_matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len]) + # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched + A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype) + A = torch_matmul(A, Vnn, out = Qn) + else: + if SDPA_HAS_GQA: + A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False, enable_gqa = True) + else: + A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False) + pass + A = A.transpose(1, 2) + A = A.reshape(bsz, 1, attention_size) + A = fast_linear_forward(self.o_proj, A, out = self.temp_O) + return A, (Kn, Vn) +pass + +# def Qwen3Model_fast_forward_inference( +# self, +# input_ids, +# past_key_values, +# position_ids, +# attention_mask = None, +# ): +# input_ids = input_ids[:,:self.max_seq_length] +# bsz, q_len = input_ids.shape +# hd = self.config.hidden_size +# mlp_size = self.config.intermediate_size + +# X = self.model.embed_tokens(input_ids) +# X = X.to(_get_dtype(self.config.torch_dtype)) +# bsz, q_len, hd = X.shape +# assert(q_len == 1) +# # Get saved buffers to reduce memory movement +# residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") +# _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") +# XX, XX2 = _XX[0], _XX[1] +# variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0") +# temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0") +# temp_gate, temp_up = temp_mlp[0], temp_mlp[1] + +# seq_len = past_key_values[0][0].shape[-2] +# if bsz != 1: +# attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( +# attention_mask, +# (bsz, q_len), +# X, +# seq_len, +# sliding_window = getattr(self.config, "sliding_window", None), +# ) +# else: +# attention_mask = None +# pass + +# next_decoder_cache = [] + +# for idx, decoder_layer in enumerate(self.model.layers): +# residual.copy_(X) # residual = X +# X = fast_rms_layernorm_inference( +# decoder_layer.input_layernorm, +# X, +# XX = XX, +# XX2 = XX2, +# variance = variance, +# ) +# X, present_key_value = Qwen3Attention_fast_forward_inference( +# decoder_layer.self_attn, +# hidden_states = X, +# past_key_value = past_key_values[idx], +# position_ids = position_ids, +# attention_mask = attention_mask, +# do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"), +# ) +# X += residual + +# residual.copy_(X) # residual = X +# X = fast_rms_layernorm_inference( +# decoder_layer.post_attention_layernorm, +# X, +# XX = XX, +# XX2 = XX2, +# variance = variance, +# ) +# X = fast_swiglu_inference( +# decoder_layer.mlp, +# X, +# temp_gate = temp_gate, +# temp_up = temp_up, +# ) +# X += residual + +# next_decoder_cache.append(present_key_value) +# pass +# X = fast_rms_layernorm_inference( +# self.model.norm, +# X, +# XX = XX, +# XX2 = XX2, +# variance = variance, +# ) + +# return BaseModelOutputWithPast( +# last_hidden_state = X, +# past_key_values = next_decoder_cache, +# hidden_states = [], +# attentions = [], +# ) +# 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(Qwen3Attention_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 diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py new file mode 100644 index 0000000000..319e3bffe3 --- /dev/null +++ b/unsloth/models/qwen3_moe.py @@ -0,0 +1,224 @@ +# 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] # 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. + 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 + + # MoE Router MLP + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) + hidden_states, router_logits = 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 + + # MoE Router MLP + residual = hidden_states + hidden_states = fast_rms_layernorm(self.post_attention_layernorm, 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 + + + +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 = Qwen3MoeDecoderLayer_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