구현체를 통해 PEFT lora 를 알아보자

기존에 썼던 짧은 lora config글과 통합하여, Lora를 근본적으로 이해하기 위해서 새로운 글을 시작한다. lora가 어떤 원리인지 구현체를 통해서 이해해 보고, huggingface peft lora를 사용할 때 어떤 config를 "왜" 조정해야 하는지 알아보자. 

 

 

0.  정말 간단한 요약

모델이 너무 무겁다. 이걸 어떻게 다 학습하나. 그래서 학습할 때 일부 layer들에 대해서(transformer에선 주로 attention block에 대해서만 적용하고 있다) weight를 따로 똑 떼어내서 학습하고자 한다. 

 

아래 그림을 보면, 기존 모델을 그대로 학습한다면 모든 weight를 다 학습에 사용한다. 반면에 lora를 학습에 이용한다면 기존 weight는 학습에서 freeze시키고 lora a, lora b를 학습 때 update 한다. 다시 모델을 Merge하거나 forward연산을 수행할 땐 freeze한 weight에 lora A, B matrics * scaling factor를 곱한 매트릭스를 통해 진행한다. 보라색의 r은 lora를 통해 얼마나 압축할 것인지 표현하는 계수임을 알 수 있다. 

lora의 원리

 

 

1. config 설명 

from peft import LoraConfig

Transformer 모델은 attention block에만 적용되는 세팅. 앞 단락에서 언급했듯이, 기존 matrix는 내버려두고(freeze) lora의 update matrices를 사용해 학습하는 구조다.

• r : update matrics의 rank. int형. 작을수록 trainable param이 적어짐. original weight matrix를 얼마나 줄여버릴 거냐에 대한 계수이므로 작을수록 많이 압축하는 느낌. 
• target_modules : lora로 바꿀 모듈. (ex. LLama)
• lora_alpha : LoRA scaling factor. scaling 값이 lora_alpha / r 로 들어감.
• bias : bias 학습할 건지 고르기. 'none', 'all', 'lora_only'. default는 none. 이를 켜면 weight뿐 아닌 bias도 lora로 처리한다. 
• layers_to_transform : LoRA로 바꿀 layer 정하기. default는 모든 레이어(taret_modules).
• rank_pattern, alpha_pattern : regexp 쓰거나 layer name 쓰는 부분. (lora_alpha랑 r 적용하고 학습 후에 다시 load 하는 과정에서 layer들을 불러올 때.) 

 

 

2. lora 구현체 파헤치기

get_peft_model 

함수 위치: hf peft github src/peft/mapping.py

AutoModelorCasualLM을 통해 불러오는 Model을 get_peft_model 함수에 인자로 넣어준다. 이를 통해 peft model을 준비함을 알 수 있다. 

model = get_peft_model(model, peft_config)

 

이 get_peft_model 함수는 어떻게 구현되어 있을까? https://github.com/huggingface/peft/blob/main/src/peft/mapping.py를 보면 PeftMixedModel 또는 PeftModel을 return하고 있음을 알 수 있다. 

def get_peft_model(
    model: PreTrainedModel, peft_config: PeftConfig, adapter_name: str = "default", mixed: bool = False
) -> PeftModel | PeftMixedModel:
    """
    Returns a Peft model object from a model and a config.

    Args:
        model ([`transformers.PreTrainedModel`]):
            Model to be wrapped.
        peft_config ([`PeftConfig`]):
            Configuration object containing the parameters of the Peft model.
        adapter_name (`str`, `optional`, defaults to `"default"`):
            The name of the adapter to be injected, if not provided, the default adapter name is used ("default").
        mixed (`bool`, `optional`, defaults to `False`):
            Whether to allow mixing different (compatible) adapter types.
    """
    model_config = getattr(model, "config", {"model_type": "custom"})
    if hasattr(model_config, "to_dict"):
        model_config = model_config.to_dict()

    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)

    if mixed:
    	#return값이 PeftMixedModel 또는 PeftModel이다. get_peft_model의 인자가 거의 그대로 들어간다.
        return PeftMixedModel(model, peft_config, adapter_name=adapter_name)

    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not peft_config.is_prompt_learning:
        return PeftModel(model, peft_config, adapter_name=adapter_name)

 

PeftModel

클래스 위치: src/peft/peft_model

PeftMixedModel도 어차피 PeftModel과 원리는 비슷할 것이다. 아래는 PeftModel 클래스의 __init__에선 active_adapter에 adapter_name을 할당하고 있다. 이 인자는 get_peft_model의 마지막 line에서 들어오는 값인데, 기존에 학습된 adapter가 없다면 'default'라는 string 형태로 들어간다. 

 

cls를 통해 base model 에 adapter name을 이식해 주는 것도 볼 수 있다. 

class PeftModel(PushToHubMixin, torch.nn.Module):
    """
    Base model encompassing various Peft methods.

    Args:
        model ([`~transformers.PreTrainedModel`]): The base transformer model used for Peft.
        peft_config ([`PeftConfig`]): The configuration of the Peft model.
        adapter_name (`str`,  *optional*): The name of the adapter, defaults to `"default"`.
    """

    def __init__(self, model: PreTrainedModel, peft_config: PeftConfig, adapter_name: str = "default") -> None:
        super().__init__()
        self.modules_to_save = None
        self.active_adapter = adapter_name
        #위 get_peft_model에서 넣어줬다. 안넣어줬으면 'default'로 들어왔다. 
        self.peft_type = peft_config.peft_type
        self._is_prompt_learning = peft_config.is_prompt_learning
        if self._is_prompt_learning:
            self._peft_config = {adapter_name: peft_config}
            self.base_model = model
            self.add_adapter(adapter_name, peft_config)
        else:
            self._peft_config = None
            cls = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]
            # base model 선언.
            self.base_model = cls(model, {adapter_name: peft_config}, adapter_name)  
            self.set_additional_trainable_modules(peft_config, adapter_name)

        if getattr(model, "is_gradient_checkpointing", True):
            model = self._prepare_model_for_gradient_checkpointing(model)

        # the `pretraining_tp` is set for some models to simulate Tensor Parallelism during inference to avoid
        # numerical differences, https://github.com/pytorch/pytorch/issues/76232 - to avoid any unexpected
        # behavior we disable that in this line.
        if hasattr(self.base_model, "config") and hasattr(self.base_model.config, "pretraining_tp"):
            self.base_model.config.pretraining_tp = 1

 

 

이 클래스의 forward가 어떻게 정의되었는지도 확인해 보면 단순하게 self.get_base_model()을 호출하고 있음을 보인다. 

 def forward(self, *args: Any, **kwargs: Any):
        """
        Forward pass of the model.
        """
        # 간단하다.
        return self.get_base_model()(*args, **kwargs)

 

 

get_base_model에서도 단순하게 self.base_model을 return 하는 게 전부다. 

 

아하, PeftModel에선 정말 껍데기로 lora model을 감싸고만 있구나. 

 

 

LoraModel

클래스 위치: src/peft/tuners/lora/model/LoraModel

이제 Lora 모델이 어떻게 만들어지고 선언되는지를 볼 차례다. 아래는 LoraModel 클래스의  _create_and_replace 함수다. target을 업데이트하는 함수를 호출하고 있다. 이 target은 layer일 것이기 때문에, LoraLayer 클래스를 찾아가 보자. 

    def _create_and_replace(self, lora_config, ....):
    # 단순하게 layer를 업데이트시키는 함수만 부르고 있다. 
    # if 문을 확인하면 target이 LoraLayer임을 알 수 있다.
    # LoraLayer에서 update_layer의 디테일을 확인해보자.
        if isinstance(target, LoraLayer) and not isinstance(target, AdaLoraLayer):
            target.update_layer(
                adapter_name,
                r,
                lora_alpha=alpha,
                lora_dropout=lora_config.lora_dropout,
                init_lora_weights=lora_config.init_lora_weights,
                use_rslora=lora_config.use_rslora,
                use_dora=lora_config.use_dora,
            )
        else:
            new_module = self._create_new_module(lora_config, adapter_name, target, **kwargs)
            if adapter_name != self.active_adapter:
                # adding an additional adapter: it is not automatically trainable
                new_module.requires_grad_(False)
            self._replace_module(parent, target_name, new_module, target)

 

 

 

LoraLayer

클래스 위치: src/peft/tuners/lora/model/LoraLayer

LoraLayer의 선언을 확인해보니 모든 layer를 A와 B로 관리하고 있음을 볼 수 있다. 

class LoraLayer(BaseTunerLayer):
    # All names of layers that may contain (trainable) adapter weights
    adapter_layer_names = ("lora_A", "lora_B", "lora_embedding_A", "lora_embedding_B")
    # All names of other parameters that may contain adapter-related parameters
    other_param_names = ("r", "lora_alpha", "scaling", "lora_dropout")

 

 

LoraModel에서  target.update_layer()가 호출되고 있던 것을 따라 여기까지 타고 왔다. update_layer를 보니 기존 weight가 [in_features, out_features] 형태였다면 lora A, B에게 각각 [in_features, r], [r, out_features] 형태로 쪼개졌음을 볼 수 있다. 

def update_layer(
        self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights, use_rslora, use_dora: bool = False
    ):
        # This code works for linear layers, override for other layer types
        if r <= 0:
            raise ValueError(f"`r` should be a positive integer value but the value passed is {r}")

        self.r[adapter_name] = r
        self.lora_alpha[adapter_name] = lora_alpha
        if lora_dropout > 0.0:
            lora_dropout_layer = nn.Dropout(p=lora_dropout)
        else:
            lora_dropout_layer = nn.Identity()

        self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer}))
        # Actual trainable parameters
        self.lora_A[adapter_name] = nn.Linear(self.in_features, r, bias=False)
        self.lora_B[adapter_name] = nn.Linear(r, self.out_features, bias=False)

 

 

Lora는 정말 문서에 적혀있는대로 [in_features, out_features]  -> [in_features, r], [r, out_features] 즉 2개의 small matrics로 쪼개지고 있구나. 

 

Linear

lora가 어떻게 구성되는지는 알았다. 남은 하나의 의문이 있다. 기존 weight와 Merge는 어떻게 하는 거지? 이는 아래 Linear class에서 나온다. LoraLayer를 상속받은 모습이다. 

모듈 key에 해당하는('proj_k', 'proj_q' 등등이 여기에 해당한다) 레이어 weight의 get_delta_weight output과 기존의 base weight를 단순히 더하는 형태임을 알 수 있다. 

 class Linear(nn.Module, LoraLayer):
 ... # LoraLayer를 상속받고 있다. 
     def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None:
		....
        for active_adapter in adapter_names:
            # 어차피 lora_A랑 lora_B key값들은 똑같다. 
            if active_adapter in self.lora_A.keys():
            	# key ex) 'proj_j', 'proj_q', ..... 
                base_layer = self.get_base_layer()
                if safe_merge:
                    # Note that safe_merge will be slower than the normal merge
                    # because of the copy operation.
                    orig_weights = base_layer.weight.data.clone()
                    # get_delta_weight로 lora 결과를 받아오고 있다. 
                    delta_weight = self.get_delta_weight(active_adapter)
                    if not self.use_dora[active_adapter]:
                    	# 정말 간단하다.
                        # base_model의 weight랑, lora 결과를 합치고 있다. 
                        orig_weights = orig_weights + delta_weight
        else:
                    delta_weight = self.get_delta_weight(active_adapter)
                    if not self.use_dora[active_adapter]:
                        base_layer.weight.data = base_layer.weight.data + delta_weight
                    else:
                        # handle dora
                        # since delta_weight already includes scaling, set it to 1 here
                        weight_norm = self._get_weight_norm(base_layer.weight, delta_weight, scaling=1).detach()
                        # We need to cache weight_norm because it has to be based on the original weights. We
                        # cannot calculate it on the fly based on the merged weights when unmerging because its a
                        # different value
                        self._cache_store(f"{active_adapter}-weight_norm", weight_norm)
                        dora_factor = self.lora_magnitude_vector[active_adapter] / weight_norm
                        new_weight = dora_factor.view(-1, 1) * (base_layer.weight.data + delta_weight)
                        base_layer.weight.data = new_weight

                self.merged_adapters.append(active_adapter)

 

 

 

중요하게 볼 부분은 여기다. 

		if not self.use_dora[active_adapter]:
                    	# 정말 간단하다.
                        # base_model의 weight랑, lora 결과를 합치고 있다. 
                        orig_weights = orig_weights + delta_weight

 

 

그럼 이제 delta_weight를 return해주는 self.get_delta_weight를 확인해보자. dtype에 따라 fp32인지 fp16인지 정해주고, A, B의 weight transpose연산을 진행해서 하나로 합쳐준다. 여기에 self.scaling 값이 곱해지는 걸 알 수 있다. scaling값은 맨 위 config 설명에서 말했던 것처럼 lora_alpha / r 값에 해당한다. lora scaling factor가 바로 여기서 등장하는 것이다. 

def get_delta_weight(self, adapter) -> torch.Tensor:
        """
        Compute the delta weight for the given adapter.

        Args:
            adapter (str):
                The name of the adapter for which the delta weight should be computed.
        """
        device = self.lora_B[adapter].weight.device
        dtype = self.lora_B[adapter].weight.dtype

        # In case users wants to merge the adapter weights that are in
        # float16 while being on CPU, we need to cast the weights to float32, perform the merge and then cast back to
        # float16 because the `@` and matmul operation in general is not supported in torch + cpu + fp16.
        cast_to_fp32 = device.type == "cpu" and dtype == torch.float16

        weight_A = self.lora_A[adapter].weight
        weight_B = self.lora_B[adapter].weight

        if cast_to_fp32:
            weight_A = weight_A.float()
            weight_B = weight_B.float()
            
        # lora weight들을 하나로 합쳐주고 있다. 
        # 맨 처음에 인자로 주었던 config가 여기서 쓰인다. 
        # 참고: @ 연산자는 matmul과 유사하다. 
        output_tensor = transpose(weight_B @ weight_A, self.fan_in_fan_out) * self.scaling[adapter]

        if cast_to_fp32:
            output_tensor = output_tensor.to(dtype=dtype)

            # cast back the weights
            self.lora_A[adapter].weight.data = weight_A.to(dtype)
            self.lora_B[adapter].weight.data = weight_B.to(dtype)

        return output_tensor

 

 

Linear.forward

하나 더, Lora가 어떻게 학습되고 있는지 확인해보면 위에서 저런 merge를 하는지와 일맥상통함을 알 수 있다. lora A, lora B의 결과를 scaling 해서 origin result와 lora 결과를 더해서 result로 return 한다. 즉,  freeze하고있는 기존 weight + ( loraA, lora B 합친 버전. dim을 생각해 보면 된다. ) * scaling factor의 값이 backpropagation의 대상이 되는 것이다. 

 def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
 	#self.base_layer는 LoraLayer정의시 가져온 모델 본연의 layer들에 해당한다. 
        if self.disable_adapters:
            if self.merged:
                self.unmerge()
            result = self.base_layer(x, *args, **kwargs)
        elif self.merged:
            result = self.base_layer(x, *args, **kwargs)
        else:
            result = self.base_layer(x, *args, **kwargs)
            torch_result_dtype = result.dtype
            for active_adapter in self.active_adapters:
                if active_adapter not in self.lora_A.keys():
                    continue
                lora_A = self.lora_A[active_adapter]
                lora_B = self.lora_B[active_adapter]
                dropout = self.lora_dropout[active_adapter]
                scaling = self.scaling[active_adapter]
                x = x.to(lora_A.weight.dtype)

                if not self.use_dora[active_adapter]:
                	# merge쪽과 유사하다. 
                	# lora_B(lora_A) : (in_feature, r) * (r, out_feature) -> (in_feature)
                    result = result + lora_B(lora_A(dropout(x))) * scaling
                else:
                    x = dropout(x)
                    result = result + self._apply_dora(x, lora_A, lora_B, scaling, active_adapter)

            result = result.to(torch_result_dtype)
        return result

 

 

---

오늘은 이렇게 Lora의 동작방식에 대해 정리해보았다. 대규모 모델들이 많이 나오면서 lora를 쓰는 분들에게 도움이 되기를.