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

Faster inference?

Browse source
This commit is contained in:
Daniel Han 2025-02-02 13:45:25 -08:00
commit 5a609b9d34
2 changed files with 22 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/kernels/utils.py
View file

@ -15,6 +15,7 @@
import triton
MAX_FUSED_SIZE : int = 65536
next_power_of_2 = triton.next_power_of_2
import functools
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
import torch
@ -96,18 +97,20 @@ def get_lora_parameters(proj):
pass
@functools.cache
def get_lora_parameters_bias(proj):
# For DPO or disabled adapters
base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj)
base_layer = getattr(proj, "base_layer", proj) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
bias = base_layer.bias
if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
return W, QUANT_STATE(W), None, None, None, bias
pass
active_adapter = proj.active_adapters[0] if \
hasattr(proj, "active_adapters") else proj.active_adapter
getattr(proj, "active_adapters", ) else proj.active_adapter
A = proj.lora_A [active_adapter].weight
B = proj.lora_B [active_adapter].weight
s = proj.scaling[active_adapter]

24
unsloth/models/llama.py
View file

@ -917,10 +917,23 @@ def LlamaModel_fast_forward_inference(
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
# 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(self.config.torch_dtype)
bsz, q_len, hd = X.shape
mlp_size = self.config.intermediate_size
assert(q_len == 1)
seq_len = past_key_values[0][0].shape[-2]
if bsz != 1:
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
@ -933,15 +946,10 @@ def LlamaModel_fast_forward_inference(
else:
attention_mask = None
pass
print(attention_mask)
next_decoder_cache = []
residual = torch.empty_like(X)
_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]
for idx, decoder_layer in enumerate(self.model.layers):
residual.copy_(X) # residual = X
X = fast_rms_layernorm_inference(