Fix tokenizer, bias, dropout supported for LoRA (#69)
* Fix tokenizer, dropout, bias for LoRA * Update loader.py
This commit is contained in:
Daniel Han
GitHub
parent
e4a080b68e
commit
45fa4bf356
4 changed files with 117 additions and 63 deletions
13
README.md
13
README.md
|
|
@ -51,7 +51,7 @@ Do **NOT** use this if you have Anaconda. You must use the Conda install method,
|
|||
```python
|
||||
import torch; torch.version.cuda
|
||||
```
|
||||
2. For Pytorch 2.1.0: You can update Pytorch via Pip (interchange `cu121` / `cu118`). Go to https://pytorch.org/ to learn more. Select either `cu118` for CUDA 11.8 or `cu121` for CUDA 12.1. If you have a RTX 3060 or higher (A100, H100 etc), use the `"ampere"` path.
|
||||
2. For Pytorch 2.1.0: You can update Pytorch via Pip (interchange `cu121` / `cu118`). Go to https://pytorch.org/ to learn more. Select either `cu118` for CUDA 11.8 or `cu121` for CUDA 12.1. If you have a RTX 3060 or higher (A100, H100 etc), use the `"ampere"` path. For Pytorch 2.1.1: got to step 3.
|
||||
```bash
|
||||
pip install --upgrade --force-reinstall --no-cache-dir torch==2.1.0 triton \
|
||||
--index-url https://download.pytorch.org/whl/cu121
|
||||
|
|
@ -118,8 +118,8 @@ model = FastLanguageModel.get_peft_model(
|
|||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",],
|
||||
lora_alpha = 16,
|
||||
lora_dropout = 0, # Currently only supports dropout = 0
|
||||
bias = "none", # Currently only supports bias = "none"
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
use_gradient_checkpointing = True,
|
||||
random_state = 3407,
|
||||
max_seq_length = max_seq_length,
|
||||
|
|
@ -174,8 +174,8 @@ model = FastLanguageModel.get_peft_model(
|
|||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",],
|
||||
lora_alpha = 64,
|
||||
lora_dropout = 0, # Currently only supports dropout = 0
|
||||
bias = "none", # Currently only supports bias = "none"
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
use_gradient_checkpointing = True,
|
||||
random_state = 3407,
|
||||
max_seq_length = max_seq_length,
|
||||
|
|
@ -209,7 +209,8 @@ dpo_trainer.train()
|
|||
|
||||
# Future Milestones and limitations
|
||||
1. Support Mixtral.
|
||||
2. Does not support non Llama models - we do so in the future.
|
||||
2. Supports all Mistral, Llama type models, but some are unoptimized (Qwen with biases)
|
||||
3. Dropout, bias in LoRA matrices are supported, just not optimized.
|
||||
|
||||
# Performance comparisons on 1 Tesla T4 GPU:
|
||||
**Time taken for 1 epoch**
|
||||
|
|
|
|||
|
|
@ -141,12 +141,49 @@ def check_tokenizer(
|
|||
if index >= max_embedding_size:
|
||||
bad_indices = list(added_tokens_fast.keys ())[j:]
|
||||
bad_tokens = list(added_tokens_fast.values())[j:]
|
||||
|
||||
if not _reload:
|
||||
# Try removing the token
|
||||
added_tokens = [str(x) for x in tokenizer.added_tokens_decoder.values()]
|
||||
special_tokens = tokenizer.special_tokens_map
|
||||
import itertools
|
||||
special_tokens = frozenset(
|
||||
itertools.chain.from_iterable(
|
||||
[x] if type(x) is str else x for x in special_tokens.values()
|
||||
)
|
||||
)
|
||||
can_be_removed1 = [x for x in bad_tokens if x not in special_tokens]
|
||||
can_be_removed2 = [x for x in can_be_removed1 if x in tokenizer._added_tokens_encoder.keys()]
|
||||
|
||||
# Check of extra tokens can in fact we removed!
|
||||
|
||||
if (len(can_be_removed1) == len(bad_tokens)) and \
|
||||
(len(can_be_removed2) == len(bad_tokens)):
|
||||
# Yes it can be fixed!
|
||||
for bad_token in can_be_removed1:
|
||||
remove_id = tokenizer._added_tokens_encoder[bad_token]
|
||||
del tokenizer._added_tokens_decoder[remove_id]
|
||||
del tokenizer._added_tokens_encoder[bad_token]
|
||||
pass
|
||||
# Confirm 1 more time!
|
||||
if max(tokenizer.added_tokens_decoder.keys()) < max_embedding_size:
|
||||
logger.warning_once(
|
||||
f"Unsloth loaded a broken tokenizer `{model_name}`, but managed to repair it!\n"\
|
||||
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"\
|
||||
"We removed these bad tokens. If you think this is incorrect, fix your tokenizer first."
|
||||
)
|
||||
return tokenizer
|
||||
pass
|
||||
pass
|
||||
|
||||
# :( Failure
|
||||
raise RuntimeError(
|
||||
f"Unsloth tried to load `{model_name}`, but cannot succeed.\n"\
|
||||
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"\
|
||||
f"Fix your tokenizer since it'll perform out of bounds memory accesses."
|
||||
)
|
||||
pass
|
||||
|
||||
# Try slow tokenizer which can fix things!
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name,
|
||||
|
|
|
|||
|
|
@ -777,9 +777,15 @@ class FastLlamaModel:
|
|||
assert(max_seq_length <= model.max_seq_length)
|
||||
|
||||
if lora_dropout != 0:
|
||||
raise TypeError("Unsloth: Fast model patching only works with dropout = 0.")
|
||||
logger.warning_once(
|
||||
f"Unsloth: Dropout = 0 is supported for fast patching. You are using dropout = {lora_dropout}.\n"\
|
||||
f"Unsloth will patch all other layers, except LoRA matrices, causing a performance hit."
|
||||
)
|
||||
if bias != "none":
|
||||
raise TypeError("Unsloth: Fast model patching only works with bias = 'none'.")
|
||||
logger.warning_once(
|
||||
f"Unsloth: bias = `none` is supported for fast patching. You are using bias = {bias}.\n"\
|
||||
f"Unsloth will patch all other layers, except LoRA matrices, causing a performance hit."
|
||||
)
|
||||
|
||||
transformers_set_seed(random_state)
|
||||
|
||||
|
|
@ -795,8 +801,8 @@ class FastLlamaModel:
|
|||
r = r,
|
||||
lora_alpha = lora_alpha,
|
||||
target_modules = target_modules,
|
||||
lora_dropout = 0,
|
||||
bias = "none",
|
||||
lora_dropout = lora_dropout,
|
||||
bias = bias,
|
||||
task_type = TaskType.CAUSAL_LM,
|
||||
layers_to_transform = layers_to_transform,
|
||||
**kwargs,
|
||||
|
|
@ -813,62 +819,64 @@ class FastLlamaModel:
|
|||
n_mlp = 0
|
||||
n_qkv = 0
|
||||
n_o = 0
|
||||
for idx, layer in enumerate(model.model.model.layers):
|
||||
if lora_dropout == 0 and bias == "none":
|
||||
for idx, layer in enumerate(model.model.model.layers):
|
||||
|
||||
# MLP patching
|
||||
gate_proj = layer.mlp.gate_proj
|
||||
up_proj = layer.mlp. up_proj
|
||||
down_proj = layer.mlp.down_proj
|
||||
# MLP patching
|
||||
gate_proj = layer.mlp.gate_proj
|
||||
up_proj = layer.mlp. up_proj
|
||||
down_proj = layer.mlp.down_proj
|
||||
|
||||
if hasattr(gate_proj, "lora_A") and \
|
||||
hasattr( up_proj, "lora_A") and \
|
||||
hasattr(down_proj, "lora_A") and \
|
||||
(gate_proj.base_layer if hasattr(gate_proj, "base_layer") else gate_proj).bias is None and \
|
||||
( up_proj.base_layer if hasattr( up_proj, "base_layer") else up_proj).bias is None and \
|
||||
(down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None:
|
||||
if hasattr(gate_proj, "lora_A") and \
|
||||
hasattr( up_proj, "lora_A") and \
|
||||
hasattr(down_proj, "lora_A") and \
|
||||
(gate_proj.base_layer if hasattr(gate_proj, "base_layer") else gate_proj).bias is None and \
|
||||
( up_proj.base_layer if hasattr( up_proj, "base_layer") else up_proj).bias is None and \
|
||||
(down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None:
|
||||
|
||||
# https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
|
||||
layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp)
|
||||
n_mlp += 1
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\n"\
|
||||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
pass
|
||||
# https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
|
||||
layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp)
|
||||
n_mlp += 1
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\n"\
|
||||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
pass
|
||||
|
||||
# QKV attention patching
|
||||
q_proj = layer.self_attn.q_proj
|
||||
k_proj = layer.self_attn.k_proj
|
||||
v_proj = layer.self_attn.v_proj
|
||||
if hasattr(q_proj, "lora_A") and \
|
||||
hasattr(k_proj, "lora_A") and \
|
||||
hasattr(v_proj, "lora_A") and \
|
||||
(q_proj.base_layer if hasattr(q_proj, "base_layer") else q_proj).bias is None and \
|
||||
(k_proj.base_layer if hasattr(k_proj, "base_layer") else k_proj).bias is None and \
|
||||
(v_proj.base_layer if hasattr(v_proj, "base_layer") else v_proj).bias is None:
|
||||
# QKV attention patching
|
||||
q_proj = layer.self_attn.q_proj
|
||||
k_proj = layer.self_attn.k_proj
|
||||
v_proj = layer.self_attn.v_proj
|
||||
if hasattr(q_proj, "lora_A") and \
|
||||
hasattr(k_proj, "lora_A") and \
|
||||
hasattr(v_proj, "lora_A") and \
|
||||
(q_proj.base_layer if hasattr(q_proj, "base_layer") else q_proj).bias is None and \
|
||||
(k_proj.base_layer if hasattr(k_proj, "base_layer") else k_proj).bias is None and \
|
||||
(v_proj.base_layer if hasattr(v_proj, "base_layer") else v_proj).bias is None:
|
||||
|
||||
layer.self_attn.apply_qkv = apply_lora_qkv
|
||||
n_qkv += 1
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Unsloth cannot patch Attention layers with our manual autograd engine since either LoRA adapters\n"\
|
||||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
pass
|
||||
layer.self_attn.apply_qkv = apply_lora_qkv
|
||||
n_qkv += 1
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Unsloth cannot patch Attention layers with our manual autograd engine since either LoRA adapters\n"\
|
||||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
pass
|
||||
|
||||
# O attention patching
|
||||
o_proj = layer.self_attn.o_proj
|
||||
if hasattr(o_proj, "lora_A") and \
|
||||
(o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None:
|
||||
# O attention patching
|
||||
o_proj = layer.self_attn.o_proj
|
||||
if hasattr(o_proj, "lora_A") and \
|
||||
(o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None:
|
||||
|
||||
layer.self_attn.apply_o = apply_lora_o
|
||||
n_o += 1
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Unsloth cannot patch O projection layer with our manual autograd engine since either LoRA adapters\n"\
|
||||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
layer.self_attn.apply_o = apply_lora_o
|
||||
n_o += 1
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Unsloth cannot patch O projection layer with our manual autograd engine since either LoRA adapters\n"\
|
||||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ FOURBIT_MAPPER = \
|
|||
"unsloth/llama-2-13b-bnb-4bit" : "unsloth/llama-13-7b",
|
||||
"unsloth/codellama-34b-bnb-4bit" : "codellama/CodeLlama-34b-hf",
|
||||
"unsloth/zephyr-sft-bnb-4bit" : "unsloth/zephyr-sft",
|
||||
"unsloth/tinyllama-bnb-4bit" : "unsloth/tinyllama",
|
||||
}
|
||||
|
||||
# https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading!
|
||||
|
|
@ -54,6 +55,13 @@ class FastLanguageModel(FastLlamaModel):
|
|||
f"to obtain the latest transformers build, then restart this session.\n"\
|
||||
f"For now, we shall load `{model_name}` instead (still 4bit, just slower downloading)."
|
||||
)
|
||||
elif not load_in_4bit and model_name in FOURBIT_MAPPER:
|
||||
new_model_name = FOURBIT_MAPPER[model_name]
|
||||
logger.warning_once(
|
||||
f"Unsloth: You passed in `{model_name}` which is a 4bit model, yet you set\n"\
|
||||
f"`load_in_4bit = False`. We shall load `{new_model_name}` instead."
|
||||
)
|
||||
model_name = new_model_name
|
||||
pass
|
||||
|
||||
model_config = AutoConfig.from_pretrained(model_name)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue