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

Qwen3 inference fixes

Browse source
This commit is contained in:
Dattu Sharma 2025-04-30 03:03:38 +00:00
commit 44834c1c12
2 changed files with 368 additions and 95 deletions
Download patch file Download diff file Expand all files Collapse all files

162
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(

297
unsloth/models/qwen3.py
View file

@ -18,6 +18,7 @@ from ._utils import __version__
from .llama import (
LlamaRotaryEmbedding,
LlamaLinearScalingRotaryEmbedding,
_LlamaModel_fast_forward_inference,
)
try:
from transformers.models.qwen3.modeling_qwen3 import (
@ -37,7 +38,9 @@ except:
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 (
@ -103,17 +106,19 @@ def Qwen3Attention_fast_forward(
if past_key_value is not None:
kv_seq_len += past_key_value[0].shape[-2]
# Extend RoPE dynamically to fit in VRAM
self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
if position_ids is None:
cos = self.rotary_emb.cos_cached
sin = self.rotary_emb.sin_cached
Q, K = fast_rope_embedding(Q, K, cos, sin)
if position_embeddings:
cos, sin = position_embeddings
else:
cos, sin = self.rotary_emb(V, seq_len = kv_seq_len)
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
pass
# 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)
@ -164,8 +169,7 @@ def Qwen3Attention_fast_forward(
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)
V = V.transpose(1, 2)
sw = getattr(self.config, "sliding_window", None)
sw = kv_seq_len if (sw is None or sw == "null") else sw
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:
@ -185,13 +189,276 @@ def Qwen3Attention_fast_forward(
# 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):
@ -212,7 +479,7 @@ class FastQwen3Model(FastLlamaModel):
Qwen3FlashAttention2.forward = Qwen3Attention_fast_forward
Qwen3DecoderLayer .forward = LlamaDecoderLayer_fast_forward
Qwen3Model .forward = LlamaModel_fast_forward
Qwen3ForCausalLM .forward = CausalLM_fast_forward(LlamaModel_fast_forward_inference)
Qwen3ForCausalLM .forward = CausalLM_fast_forward(_LlamaModel_fast_forward_inference(Qwen3Attention_fast_forward_inference))
PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward
fix_prepare_inputs_for_generation(Qwen3ForCausalLM)