From 45fa4bf3561018402e5dfaf2c3a4e3ad22cae5de Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 6 Jan 2024 19:13:39 +1100 Subject: [PATCH] Fix tokenizer, bias, dropout supported for LoRA (#69) * Fix tokenizer, dropout, bias for LoRA * Update loader.py --- README.md | 13 ++--- unsloth/models/_utils.py | 37 +++++++++++++ unsloth/models/llama.py | 114 +++++++++++++++++++++------------------ unsloth/models/loader.py | 8 +++ 4 files changed, 113 insertions(+), 59 deletions(-) diff --git a/README.md b/README.md index b1c08d47de..d93bd0090c 100644 --- a/README.md +++ b/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** diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index d9e5686be5..2f882cc7c3 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -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, diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 88c095265c..e2361d389c 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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 diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index eb8d4960b9..61585d77a5 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -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)