PEFT로 LoRA Checkpoint 로드시 size mismatch 해결법

이슈


size mismatch for base_model.model.gpt_neox.layers.0.attention.query_key_value.lora_A.weight: copying a param with shape torch.Size([16, 5120]) from checkpoint, the shape in current model is torch.Size([8, 5120]).
        size mismatch for base_model.model.gpt_neox.layers.0.attention.query_key_value.lora_B.weight: copying a param with shape torch.Size([10240, 8, 1]) from checkpoint, the shape in current model is torch.Size([15360, 8]).

LoRA로 학습한 모델의 checkpoint를 PEFT Model 로드시 문제가 생기는 현상.


import torch
from peft import PeftModel
from transformers import LlamaTokenizer, LlamaForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM

def load_model(base, finetuned, multi_gpu):
    tokenizer = AutoTokenizer.from_pretrained(base)
    # tokenizer.pad_token_id = 0
    tokenizer.padding_side = "left"

    if not multi_gpu:
        model = AutoModelForCausalLM.from_pretrained(
            base,
            load_in_8bit=True,
            device_map="auto",
        )
        
        model = PeftModel.from_pretrained(
            model, finetuned, device_map={'': 0}
        )
        return model, tokenizer
    else:
        model = AutoModelForCausalLM.from_pretrained(
            base,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        
        model = PeftModel.from_pretrained(
            model, finetuned, torch_dtype=torch.float16
        )
        model.half()
        return model, tokenizer

위와 같은 코드에서 Peft Model 로드시에 모델 사이즈가 맞지 않다는 에러가 난다.

원인분석

현재 버전의 peft 0.2.0의 경우, 코드상에는 LoRA Model이 아래와 같이 peft_config.target_modules 이 1일 경우, 즉 GPT-NeoX와 같이 q 와 k 가 모두 한 layer에 들어가있는 경우 → query_key_value

분명히 아래와 같이 enable_lora 가 [True, False, True] 로, fan_in_fan_out 이 True 로 설정되어야 한다.

peft/mapping.py at v0.2.0 · huggingface/peft

🤗 PEFT: State-of-the-art Parameter-Efficient Fine-Tuning. - peft/mapping.py at v0.2.0 · huggingface/peft

https://github.com/huggingface/peft/blob/v0.2.0/src/peft/mapping.py#L114-L124

하지만 실제로 LLAMA LoRA에서 사용했던 것 처럼 아래와 같이 config가 null, false 로 각각 설정될 경우에는 해당 코드 설정이 제대로 반영되지 않는다.


{
  "base_model_name_or_path": "decapoda-research/llama-65b-hf",
  "bias": "none",
  "enable_lora": null,
  "fan_in_fan_out": false,
  "inference_mode": true,
  "lora_alpha": 16,
  "lora_dropout": 0.05,
  "merge_weights": false,
  "modules_to_save": null,
  "peft_type": "LORA",
  "r": 8,
  "target_modules": [
    "q_proj",
    "v_proj"
  ],
  "task_type": "CAUSAL_LM"
}

해결법: `adapter_config.json` 수정하기!

As-Is


{
  "base_model_name_or_path": "beomi/polyglot-ko-12.8b",
  "bias": "none",
  "enable_lora": null,
  "fan_in_fan_out": false,
  "inference_mode": true,
  "lora_alpha": 16,
  "lora_dropout": 0.05,
  "merge_weights": false,
  "modules_to_save": null,
  "peft_type": "LORA",
  "r": 8,
  "target_modules": [
    "query_key_value"
  ],
  "task_type": "CAUSAL_LM"
}

To-Be


{
  "base_model_name_or_path": "beomi/polyglot-ko-12.8b",
  "bias": "none",
  "enable_lora": [true, false, true],
  "fan_in_fan_out": true,
  "inference_mode": true,
  "lora_alpha": 16,
  "lora_dropout": 0.05,
  "merge_weights": false,
  "modules_to_save": null,
  "peft_type": "LORA",
  "r": 8,
  "target_modules": [
    "query_key_value"
  ],
  "task_type": "CAUSAL_LM"
}

PEFT로 LoRA Checkpoint 로드시 size mismatch 해결법

이슈

원인분석

해결법: adapter_config.json 수정하기!

As-Is

To-Be

해결법: `adapter_config.json` 수정하기!