Fix: loading models with resized vocabulary

[oKatanaaa]

This PR is intended to address the issue of loading models with resized vocabulary in Unsloth.

At the moment loading models with resized vocab fails because of tensor shapes mismatch.

The fix is plain simple: resize the embedding layer before loading peft modules (which include resized embeddings).

This fix is battle tested in my own experiments, I constantly use it to train Mistral with ChatML template by adding a new embedding for the <|im_start|> token. It works as intended. It also does not induce any compatibility issues since resize is conditional (embeddings are resized only when new vocab size is provided).

Example usage:

# Loading peft (with adapters) Mistral with newly added token (original mistral vocab is 32000)
model_loaded, _ = FastLanguageModel.from_pretrained(
    "updated_model", 
    resize_model_vocab=32001
)

To prove the 'safety' of the fix, I did additional testing with weights values (embed_tokens) comparison:

# %%
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
from unsloth import FastLanguageModel
import gc
import torch

def free_mem():
    gc.collect()
    torch.cuda.empty_cache()

# %%
model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/mistral-7b-instruct-v0.1-bnb-4bit",
    load_in_4bit=True
)

# %%
w_original = model.model.embed_tokens.weight.clone()

# %%
# Base vocab is 32000
model.resize_token_embeddings(32001)
free_mem()

# %%
model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
    modules_to_save=["embed_tokens", "lm_head"]
)

# %%
model.save_pretrained('/workspace/test_save_model', save_embedding_layers=True)

# %%
# Vanilla unsloth loading
try:
    model_loaded, _ = FastLanguageModel.from_pretrained("/workspace/test_save_model")
except Exception as ex:
    print('Load failed.')
    print(ex)
free_mem()

# %%
model_loaded, _ = FastLanguageModel.from_pretrained(
    "/workspace/test_save_model", 
    resize_model_vocab=32001
)

# %%
w_resized = model.base_model.base_model.embed_tokens.modules_to_save.default.weight
w_uploaded = model_loaded.base_model.base_model.embed_tokens.modules_to_save.default.weight.float()

# %%
print('Resized and uploaded are the same:', torch.allclose(w_resized, w_uploaded))
print('Original and resized slice are the same:', torch.allclose(w_original.float(), w_resized[:-1]))
print('Original and uploaded slice are the same:', torch.allclose(w_original.float(), w_uploaded[:-1]))

Here is the output:

# When loading using vanilla unsloth
==((====))==  Unsloth: Fast Mistral patching release 2024.4
   \\   [/](https://vscode-remote+ssh-002dremote-002bcuda3-005fnlp-005fllm-005fdev.vscode-resource.vscode-cdn.net/)|    GPU: NVIDIA GeForce RTX 3090. Max memory: 23.691 GB. Platform = Linux.
O^O/ \_/ \    Pytorch: 2.1.2. CUDA = 8.6. CUDA Toolkit = 12.1.
\        [/](https://vscode-remote+ssh-002dremote-002bcuda3-005fnlp-005fllm-005fdev.vscode-resource.vscode-cdn.net/)    Bfloat16 = TRUE. Xformers = 0.0.23.post1. FA = False.
 "-____-"     Free Apache license: http://github.com/unslothai/unsloth
Load failed.
Error(s) in loading state_dict for PeftModelForCausalLM:
	size mismatch for base_model.model.model.embed_tokens.modules_to_save.default.weight: copying a param with shape torch.Size([32001, 4096]) from checkpoint, the shape in current model is torch.Size([32000, 4096]).
	size mismatch for base_model.model.lm_head.modules_to_save.default.weight: copying a param with shape torch.Size([32001, 4096]) from checkpoint, the shape in current model is torch.Size([32000, 4096]).

# Loading the model as outlined above (with `resize_model_vocab` specified) works as intended and loads the model correctly

# Embed values comparison
Resized and uploaded are the same: True
Original and resized slice are the same: True
Original and uploaded slice are the same: True

The resized embeddings match the original ones (when taking a corresponding slice). The new embeddings match between between resized and uploaded models (which is expected, loading works correctly).

Fix is trivial, does not affect other functionalities, and works as expected, so it can be merged without any problem.
It is essentially the same as #272, but resize happens conditionally only when new vocab size is specified. Also there is no need to set model.config.vocab_size = model.config.vocab_size + new_token_num manually, it happens automatically.
The other difference is that #272 can only increase the number of tokens, whereas I opted for a general approach when new size can be arbitrary which may be relevant for research cases.

P.S. I also added gitignore file, because otherwise GIT track python cache files. It must be added anyways.

UPD. I also checked lm_head weights to be 100% sure, they do match.