Fix RoPE Scaling issues (#52)

* Fix RoPE Scaling

* Update llama.py

* Update llama.py
This commit is contained in:
Daniel Han 2023-12-26 04:32:04 +11:00 committed by GitHub
parent 5dda3f4058
commit 04093395ed
2 changed files with 11 additions and 5 deletions

View file

@ -369,6 +369,7 @@ 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
assert(seq_length <= self.max_seq_length)
past_key_values_length = 0
if past_key_values is not None:
@ -661,6 +662,9 @@ class FastLlamaModel:
bnb_4bit_compute_dtype = dtype,
)
# https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/12
# RoPE Scaling's max_position_embeddings must be updated
max_position_embeddings = max(max_seq_length, model_max_seq_length)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map = device_map,
@ -668,6 +672,7 @@ class FastLlamaModel:
quantization_config = bnb_config,
token = token,
rope_scaling = rope_scaling,
max_position_embeddings = max_position_embeddings,
)
tokenizer = AutoTokenizer.from_pretrained(
model_name,
@ -685,7 +690,7 @@ class FastLlamaModel:
layer.self_attn.apply_o = original_apply_o
pass
model.max_seq_length = max_seq_length
model.max_seq_length = max_position_embeddings
return model, tokenizer
pass
@ -746,7 +751,7 @@ class FastLlamaModel:
layers_to_transform = None,
use_gradient_checkpointing = True,
random_state = 3407,
max_seq_length = 2048,
max_seq_length = 2048, # not used anymore
**kwargs,
):
assert(max_seq_length <= model.max_seq_length)
@ -824,6 +829,7 @@ class FastLlamaModel:
# Patch cross entropy loss labels
# Fixes https://github.com/unslothai/unsloth/issues/10
max_seq_length = model.max_seq_length
extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda")
model.model.extra_ignored_labels = extra_ignored_labels
internal_model = model

View file

@ -125,7 +125,7 @@ def MistralAttention_fast_forward(
V = V.transpose(1, 2)
# Flash Attention v2 auto supports grouped query attention
sliding_window = self.config.sliding_window
sliding_window = getattr(self.config, "sliding_window")
sliding_window = q_len if sliding_window is None else sliding_window
window = (-1, -1) if (q_len <= sliding_window) else (sliding_window, sliding_window)
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
@ -169,7 +169,7 @@ def MistralForCausalLM_fast_forward(
if causal_mask is None:
bsz, q_len = input_ids.shape
sliding_window = self.config.sliding_window
sliding_window = getattr(self.config, "sliding_window")
if sliding_window is None or sliding_window <= 0:
causal_mask = xformers.attn_bias.LowerTriangularMask()
elif q_len <= sliding_window:
@ -312,7 +312,7 @@ class FastMistralModel(FastLlamaModel):
layer.self_attn.apply_o = original_apply_o
pass
model.max_seq_length = max_seq_length
model.max_seq_length = max(max_seq_length, model.config.max_position_embeddings)
return model, tokenizer
pass
pass