DPO, SWA fixes (#57)
* Pytorch 2.1.1 install path, 4bit loading * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update loader.py * Update loader.py * Update loader.py * Update loader.py * Spelling errors * Update __init__.py * DPO loss fix * Update fast_lora.py * Update fast_lora.py * Out of bounds tokenization * Fix Mistral SWA
This commit is contained in:
Daniel Han
GitHub
parent
0369e7aa7a
commit
24133feda6
3 changed files with 63 additions and 24 deletions
|
|
@ -17,10 +17,16 @@ from .utils import fast_dequantize, QUANT_STATE
|
|||
from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
|
||||
|
||||
def get_lora_parameters(proj):
|
||||
active_adapter = proj.active_adapters[0] if \
|
||||
hasattr(proj, "active_adapters") else proj.active_adapter
|
||||
# For DPO or disabled adapters
|
||||
base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj)
|
||||
W = base_layer.weight
|
||||
|
||||
if proj.disable_adapters or proj.merged:
|
||||
return W, QUANT_STATE(W), None, None, None
|
||||
pass
|
||||
|
||||
active_adapter = proj.active_adapters[0] if \
|
||||
hasattr(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]
|
||||
|
|
@ -31,7 +37,6 @@ pass
|
|||
def matmul_lora(X, W, W_quant, A, B, s, out = None):
|
||||
dtype = X.dtype
|
||||
W = fast_dequantize(W.t(), W_quant)
|
||||
A, B = A.t(), B.t()
|
||||
|
||||
if X.dim() == 3:
|
||||
batch, seq_len, d = X.shape
|
||||
|
|
@ -43,7 +48,13 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
|
|||
|
||||
out = torch.matmul(X, W, out = out)
|
||||
if W_quant is not None: del W
|
||||
out += (X @ A.to(dtype)) @ (s * B.to(dtype))
|
||||
|
||||
if A is not None:
|
||||
# LoRA is enabled
|
||||
A, B = A.t(), B.t()
|
||||
out += (X @ A.to(dtype)) @ (s * B.to(dtype))
|
||||
pass
|
||||
|
||||
return out.view(batch, seq_len, -1) if reshape else out
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -369,8 +369,21 @@ def LlamaModel_fast_forward(
|
|||
raise ValueError("Unsloth: You have to specify either decoder_input_ids or decoder_inputs_embeds")
|
||||
|
||||
seq_length_with_past = seq_length
|
||||
|
||||
# Fix out of bounds tokenization
|
||||
if hasattr(self, "max_seq_length"):
|
||||
assert(seq_length <= self.max_seq_length)
|
||||
if seq_length > self.max_seq_length:
|
||||
logger.warning_once(
|
||||
f"Unsloth: Input IDs of length {seq_length} > the model's max sequence length of {self.max_seq_length}.\n"\
|
||||
"We shall truncate it ourselves. It's imperative if you correct this issue first."
|
||||
)
|
||||
if input_ids is not None:
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
elif inputs_embeds is not None:
|
||||
inputs_embeds = inputs_embeds[:,:self.max_seq_length,:]
|
||||
pass
|
||||
pass
|
||||
|
||||
past_key_values_length = 0
|
||||
|
||||
if past_key_values is not None:
|
||||
|
|
|
|||
|
|
@ -97,21 +97,36 @@ def MistralAttention_fast_forward(
|
|||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
M = bsz * q_len
|
||||
|
||||
has_sliding_window = isinstance(causal_mask, xformers.attn_bias.BlockDiagonalCausalMask)
|
||||
|
||||
# Group query attention
|
||||
if n_groups != 1:
|
||||
K = K .view(bsz, q_len, n_kv_heads, 1, head_dim)
|
||||
V = V .view(bsz, q_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
# Xformers does not support backward, so we have to convert
|
||||
# GQA to MQA by cloning K and V
|
||||
K = K.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made
|
||||
V = V.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
# if n_groups != 1:
|
||||
K = K .view(bsz, q_len, n_kv_heads, 1, head_dim)
|
||||
V = V .view(bsz, q_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
# Xformers does not support backward, so we have to convert
|
||||
# GQA to MQA by cloning K and V
|
||||
K = K.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made
|
||||
V = V.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made
|
||||
|
||||
if has_sliding_window:
|
||||
Q = Q.view(1, M, n_heads, head_dim)
|
||||
K = K.view(1, M, n_heads, head_dim)
|
||||
V = V.view(1, 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_sliding_window:
|
||||
Q = Q.view(1, M, n_kv_heads, n_groups, head_dim)
|
||||
K = K.view(1, M, n_kv_heads, n_groups, head_dim)
|
||||
V = V.view(1, M, n_kv_heads, n_groups, head_dim)
|
||||
pass
|
||||
pass
|
||||
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
|
|
@ -131,12 +146,12 @@ def MistralAttention_fast_forward(
|
|||
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, q_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, q_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, q_len, head_dim)
|
||||
pass
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, q_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, q_len, head_dim)
|
||||
# pass
|
||||
# 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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue