Size mismatch after fine-tuning text generation model

[pahautelman]

Hi!
I recently attempted to fine-tune the microsoft/phi-2 model using SFT and DPO trainers. However, upon attempting to reload the model, I encountered the following error message:

Error(s) in loading state_dict for PhiForCausalLM: size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([50295, 2560]) from checkpoint, the shape in the current model is torch.Size([51200, 2560]). size mismatch for lm_head.weight: copying a param with shape torch.Size([50295, 2560]) from checkpoint, the shape in the current model is torch.Size([51200, 2560]). size mismatch for lm_head.bias: copying a param with shape torch.Size([50295]) from checkpoint, the shape in the current model is torch.Size([51200]).

I used the following hyperparameters:
{ "block_size": 1024, "model_max_length": 2048, "padding": "right", "use_flash_attention_2": false, "disable_gradient_checkpointing": false, "logging_steps": -1, "evaluation_strategy": "epoch", "save_total_limit": 1, "save_strategy": "epoch", "auto_find_batch_size": false, "mixed_precision": "bf16", "lr": 3e-05, "epochs": 1, "batch_size": 2, "warmup_ratio": 0.05, "gradient_accumulation": 1, "optimizer": "adamw_torch", "scheduler": "linear", "weight_decay": 0.0, "max_grad_norm": 1.0, "seed": 42, "apply_chat_template": false, "quantization": "int4", "target_modules": "", "merge_adapter": false, "peft": true, "lora_r": 16, "lora_alpha": 32, "lora_dropout": 0.05, "model_ref": null, "dpo_beta": 0.1 }

Has anyone encountered this issue before, and does anyone have insights on how to resolve it? Thanks for the help!

[abhishekkrthakur]

not sure if i understand when the error occurs, during training?

[pahautelman]

I encountered the error while attempting to load the trained model, similar to the example shown on the Hugging Face card. Please note that I fine-tuned the model using Hugging Face's Autotrain space.

tokenizer = AutoTokenizer.from_pretrained(model_path)

# use BitsAndBytesConfig to load model in 4bit
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    quantization_config=bnb_config
).eval()

[abhishekkrthakur]

there should be a handler.py in your model repo. that shows how to load model in case you used peft. could you try following the model loading part from there please?

[pahautelman]

I will try and update shortly. Thank you very much for your help and prompt response, @abhishekkrthakur!

[abhishekkrthakur]

please let me know. if it still doesnt work, we'll take a deeper look and convert this to an issue

[pahautelman]

It works like a charm like a charm now! Thank you!

Used code:

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    trust_remote_code=True,
)
model.resize_token_embeddings(len(tokenizer))
model = PeftModel.from_pretrained(model, model_path)
model = model.merge_and_unload()

I was able to successfully load the model in Google Colab. However, when I tried to run the same code on my personal machine, I encountered an error for model = model.merge_and_unload(): RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling 'cublasCreate(handle)'. Do you happen to know why this error occurs?"

[abhishekkrthakur]

glad to know. the error seems like it went out of memory