ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into nightly

Browse source
This commit is contained in:
Daniel Han 2025-04-29 23:48:43 -07:00
commit efad91f704
11 changed files with 1177 additions and 103 deletions
Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -10,7 +10,7 @@
<a href="https://discord.com/invite/unsloth"><img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/Discord button.png" height="48"></a>
<a href="https://docs.unsloth.ai"><img src="https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/images/Documentation%20Button.png" height="48"></a>
### 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)

4
pyproject.toml
View file

@ -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",

15
unsloth/dataprep/__init__.py Normal file
View file

@ -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 *

261
unsloth/dataprep/synthetic.py Normal file
View file

@ -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

18
unsloth/models/__init__.py
View file

@ -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

6
unsloth/models/_utils.py
View file

@ -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:

164
unsloth/models/llama.py
View file

@ -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

13
unsloth/models/loader.py
View file

@ -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

40
unsloth/models/mapper.py
View file

@ -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",
),
}

527
unsloth/models/qwen3.py Normal file
View file

@ -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

224
unsloth/models/qwen3_moe.py Normal file
View file

@ -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